refactor: make create_feature_view method more generic - drop BigQuery naming

Move BigQuerySource to util file to avoid circular dep + a few lint fixes.

PiperOrigin-RevId: 629484260

Author: vertex-sdk-bot

tests/unit/vertexai/test_feature_online_store.py

@@ -28,7 +28,7 @@
 from vertexai.resources.preview import (
     FeatureOnlineStore,
     FeatureOnlineStoreType,
-    FeatureView,
+    FeatureViewBigQuerySource,
     IndexConfig,
     DistanceMeasureType,
     TreeAhConfig,
@@ -361,15 +361,26 @@ def test_delete(force, delete_fos_mock, get_fos_mock, fos_logger_mock, sync=True
     )
 
 
-def test_create_bq_fv_none_source_raises_error(get_fos_mock):
+def test_create_fv_none_source_raises_error(get_fos_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
     fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
 
     with pytest.raises(
         ValueError,
-        match=re.escape("Please specify valid big_query_source."),
+        match=re.escape("Please specify a valid source."),
     ):
-        fos.create_feature_view_from_big_query("bq_fv", None)
+        fos.create_feature_view("bq_fv", None)
+
+
+def test_create_fv_wrong_object_type_raises_error(get_fos_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape("Only FeatureViewBigQuerySource is a supported source."),
+    ):
+        fos.create_feature_view("bq_fv", fos)
 
 
 def test_create_bq_fv_bad_uri_raises_error(get_fos_mock):
@@ -378,11 +389,11 @@ def test_create_bq_fv_bad_uri_raises_error(get_fos_mock):
 
     with pytest.raises(
         ValueError,
-        match=re.escape("Please specify URI in big_query_source."),
+        match=re.escape("Please specify URI in BigQuery source."),
     ):
-        fos.create_feature_view_from_big_query(
+        fos.create_feature_view(
             "bq_fv",
-            FeatureView.BigQuerySource(uri=None, entity_id_columns=["entity_id"]),
+            FeatureViewBigQuerySource(uri=None, entity_id_columns=["entity_id"]),
         )
 
 
@@ -395,11 +406,11 @@ def test_create_bq_fv_bad_entity_id_columns_raises_error(
 
     with pytest.raises(
         ValueError,
-        match=re.escape("Please specify entity ID columns in big_query_source."),
+        match=re.escape("Please specify entity ID columns in BigQuery source."),
     ):
-        fos.create_feature_view_from_big_query(
+        fos.create_feature_view(
             "bq_fv",
-            FeatureView.BigQuerySource(uri="hi", entity_id_columns=entity_id_columns),
+            FeatureViewBigQuerySource(uri="hi", entity_id_columns=entity_id_columns),
         )
 
 
@@ -416,9 +427,9 @@ def test_create_bq_fv(
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
     fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
 
-    fv = fos.create_feature_view_from_big_query(
+    fv = fos.create_feature_view(
         _TEST_FV1_ID,
-        FeatureView.BigQuerySource(
+        FeatureViewBigQuerySource(
             uri=_TEST_FV1_BQ_URI, entity_id_columns=_TEST_FV1_ENTITY_ID_COLUMNS
         ),
         labels=_TEST_FV1_LABELS,
@@ -478,9 +489,9 @@ def test_create_embedding_fv(
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
     fos = FeatureOnlineStore(_TEST_ESF_OPTIMIZED_FOS_ID)
 
-    embedding_fv = fos.create_feature_view_from_big_query(
+    embedding_fv = fos.create_feature_view(
         _TEST_OPTIMIZED_EMBEDDING_FV_ID,
-        FeatureView.BigQuerySource(uri="hi", entity_id_columns=["entity_id"]),
+        FeatureViewBigQuerySource(uri="hi", entity_id_columns=["entity_id"]),
         index_config=IndexConfig(
             embedding_column="embedding",
             dimensions=1536,

vertexai/resources/preview/__init__.py

@@ -38,6 +38,7 @@
     FeatureOnlineStore,
     FeatureOnlineStoreType,
     FeatureView,
+    FeatureViewBigQuerySource,
     FeatureViewReadResponse,
     IndexConfig,
     TreeAhConfig,
@@ -60,6 +61,7 @@
     "FeatureOnlineStoreType",
     "FeatureOnlineStore",
     "FeatureView",
+    "FeatureViewBigQuerySource",
     "FeatureViewReadResponse",
     "IndexConfig",
     "TreeAhConfig",

vertexai/resources/preview/feature_store/__init__.py

@@ -26,6 +26,7 @@
 )
 
 from vertexai.resources.preview.feature_store.utils import (
+    FeatureViewBigQuerySource,
     FeatureViewReadResponse,
     IndexConfig,
     TreeAhConfig,
@@ -38,6 +39,7 @@
     FeatureOnlineStoreType,
     FeatureOnlineStore,
     FeatureView,
+    FeatureViewBigQuerySource,
     FeatureViewReadResponse,
     IndexConfig,
     IndexConfig,

vertexai/resources/preview/feature_store/feature_online_store.py

@@ -38,6 +38,7 @@
 )
 from vertexai.resources.preview.feature_store.utils import (
     IndexConfig,
+    FeatureViewBigQuerySource,
 )
 
 
@@ -377,10 +378,10 @@ def labels(self) -> Dict[str, str]:
         return self._gca_resource.labels
 
     @base.optional_sync()
-    def create_feature_view_from_big_query(
+    def create_feature_view(
         self,
         name: str,
-        big_query_source: FeatureView.BigQuerySource,
+        source: FeatureViewBigQuerySource,
         labels: Optional[Dict[str, str]] = None,
         sync_config: Optional[str] = None,
         index_config: Optional[IndexConfig] = None,
@@ -390,7 +391,7 @@ def create_feature_view_from_big_query(
         request_metadata: Optional[Sequence[Tuple[str, str]]] = None,
         create_request_timeout: Optional[float] = None,
         sync: bool = True,
-    ) -> "FeatureView":
+    ) -> FeatureView:
         """Creates a FeatureView from a BigQuery source.
 
         Example Usage:
@@ -422,9 +423,9 @@ def create_feature_view_from_big_query(
         ```
         Args:
             name: The name of the feature view.
-            big_query_source:
-                The BigQuery source to load data from when a feature view sync
-                runs.
+            source:
+                The source to load data from when a feature view sync runs.
+                Currently supports a BigQuery source.
             labels:
                 The labels with user-defined metadata to organize your
                 FeatureViews.
@@ -445,6 +446,11 @@ def create_feature_view_from_big_query(
                 for each entity ID of this FeatureView are made ready for online
                 serving. Example format: "TZ=America/New_York 0 9 * * *" (sync
                 daily at 9 AM EST).
+            index_config:
+                Configuration for index preparation for vector search. It
+                contains the required configurations to create an index from
+                source data, so that approximate nearest neighbor (a.k.a ANN)
+                algorithms search can be performed during online serving.
             project:
                 Project to create feature view in. If unset, the project set in
                 aiplatform.init will be used.
@@ -467,17 +473,24 @@ def create_feature_view_from_big_query(
         Returns:
             FeatureView - the FeatureView resource object.
         """
-        if not big_query_source:
-            raise ValueError("Please specify valid big_query_source.")
-        elif not big_query_source.uri:
-            raise ValueError("Please specify URI in big_query_source.")
-        elif not big_query_source.entity_id_columns:
-            raise ValueError("Please specify entity ID columns in big_query_source.")
+        if not source:
+            raise ValueError("Please specify a valid source.")
+
+        # Only BigQuery source is supported right now.
+        if not isinstance(source, FeatureViewBigQuerySource):
+            raise ValueError("Only FeatureViewBigQuerySource is a supported source.")
+
+        # BigQuery source validation.
+        if not source.uri:
+            raise ValueError("Please specify URI in BigQuery source.")
+
+        if not source.entity_id_columns:
+            raise ValueError("Please specify entity ID columns in BigQuery source.")
 
         gapic_feature_view = gca_feature_view.FeatureView(
             big_query_source=gca_feature_view.FeatureView.BigQuerySource(
-                uri=big_query_source.uri,
-                entity_id_columns=big_query_source.entity_id_columns,
+                uri=source.uri,
+                entity_id_columns=source.entity_id_columns,
             ),
             sync_config=gca_feature_view.FeatureView.SyncConfig(cron=sync_config)
             if sync_config

vertexai/resources/preview/feature_store/feature_view.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 #
 
-from dataclasses import dataclass
 import re
 from typing import List, Optional
 from google.cloud.aiplatform import initializer
@@ -342,8 +341,7 @@ def search(
             )
         else:
             raise ValueError(
-                f"Either entity_id or embedding_value needs to be provided for"
-                f" search."
+                "Either entity_id or embedding_value needs to be provided for search."
             )
         response = self._get_online_store_client.search_nearest_entities(
             request=fos_service.SearchNearestEntitiesRequest(
@@ -365,11 +363,6 @@ def search(
         )
         return fv_utils.SearchNearestEntitiesResponse(response)
 
-    @dataclass
-    class BigQuerySource:
-        uri: str
-        entity_id_columns: List[str]
-
     class FeatureViewSync(base.VertexAiResourceNounWithFutureManager):
         """Class for managing Feature View Sync resources."""
 

vertexai/resources/preview/feature_store/utils.py

@@ -44,6 +44,12 @@ class PublicEndpointNotFoundError(RuntimeError):
     """Public endpoint has not been created yet."""
 
 
+@dataclass
+class FeatureViewBigQuerySource:
+    uri: str
+    entity_id_columns: List[str]
+
+
 @dataclass
 class FeatureViewReadResponse:
     _response: fos_service.FetchFeatureValuesResponse