feat: add document store

HaoXuAI · HaoXuAI · commit 2cd73d199f76 · 2024-03-31T16:45:50.000-07:00
diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres.py b/sdk/python/feast/infra/online_stores/contrib/postgres.py
@@ -341,13 +341,16 @@ def online_search(
         return result
 
     def create_index(
-        self, config: RepoConfig, index: str, index_config: DocumentStoreIndexConfig
+            self,
+            config: RepoConfig,
+            table: str
     ):
+        document_store_config = config.document_store_config
         with self._get_conn(config) as conn, conn.cursor() as cur:
             cur.execute(
                 CREATE_INDEX_QUERY_TEMPLATE.format(
-                    table_name=config.project,
-                    index_type=index,
-                    embeding_type=index_config.embedding_type,
+                    table=table,
+                    index_type=document_store_config.index_type,
+                    embeding_type=document_store_config.embedding_type,
                 )
             )
diff --git a/sdk/python/feast/infra/online_stores/document_store.py b/sdk/python/feast/infra/online_stores/document_store.py
@@ -7,29 +7,23 @@
 from feast.feature_view import FeatureView
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 from feast.repo_config import FeastConfigBaseModel, RepoConfig
-from infra.online_stores.online_store import OnlineStore
+from feast.infra.online_stores.online_store import OnlineStore
 
 
 class DocumentStoreIndexConfig(FeastConfigBaseModel):
+    index_type: Optional[str]
     embedding_type: Optional[str]
 
 
 class DocumentStore(OnlineStore):
-    index: Optional[str]
 
     @abstractmethod
     def online_search(
-        self,
-        config: RepoConfig,
-        table: FeatureView,
-        requested_feature: str,
-        embeddings: np.ndarray,
-        top_k: int,
+            self,
+            config: RepoConfig,
+            table: FeatureView,
+            requested_feature: str,
+            embeddings: np.ndarray,
+            top_k: int,
     ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
-        raise NotImplementedError("You have to implement this!")
-
-    @abstractmethod
-    def create_index(
-        self, config: RepoConfig, index: str, index_config: DocumentStoreIndexConfig
-    ):
-        raise NotImplementedError("You have to implement this!")
+        raise NotImplementedError("You have to implement this!")
