Fixed data source import bug

MatsMoll · Jun 11, 2024 · 17ca065 · 17ca065
1 parent d14b470
commit 17ca065
Show file tree

Hide file tree

Showing 4 changed files with 758 additions and 741 deletions.
diff --git a/aligned/data_source/batch_data_source.py b/aligned/data_source/batch_data_source.py
@@ -39,6 +39,9 @@ def __init__(self) -> None:
             DeltaFileSource,
             PartitionedParquetFileSource,
         )
+        from aligned.schemas.feature_view import FeatureViewReferenceSource
+        from aligned.schemas.model import ModelSource
+
         from aligned.sources.psql import PostgreSQLDataSource
         from aligned.sources.redshift import RedshiftSQLDataSource
         from aligned.sources.s3 import AwsS3CsvDataSource, AwsS3ParquetDataSource
@@ -48,8 +51,7 @@ def __init__(self) -> None:
             AzureBlobParquetDataSource,
             AzureBlobPartitionedParquetDataSource,
         )
-        from aligned.schemas.feature_view import FeatureViewReferenceSource
-        from aligned.schemas.model import ModelSource
+        from aligned.sources.lancedb import LanceDbTable
 
         source_types = [
             PostgreSQLDataSource,
@@ -67,6 +69,8 @@ def __init__(self) -> None:
             AzureBlobDeltaDataSource,
             AzureBlobParquetDataSource,
             AzureBlobPartitionedParquetDataSource,
+            # LanceDB
+            LanceDbTable,
             # Aligned Related Sources
             JoinDataSource,
             JoinAsofDataSource,

diff --git a/aligned/exposed_model/interface.py b/aligned/exposed_model/interface.py
@@ -35,6 +35,7 @@ def __init__(self):
             InMemMLFlowAlias,
             ShadowModel,
             ABTestModel,
+            DillPredictor,
         ]
         for predictor in types:
             self.supported_predictors[predictor.model_type] = predictor