fix: FeatureGroup.create() - remove unused entity id column param

PiperOrigin-RevId: 646648227
googleapis · Jun 25, 2024 · c46f3e9 · c46f3e9
1 parent c8be56a
commit c46f3e9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/vertexai/resources/preview/feature_store/feature_group.py b/vertexai/resources/preview/feature_store/feature_group.py
@@ -92,7 +92,6 @@ def create(
         cls,
         name: str,
         source: FeatureGroupBigQuerySource = None,
-        entity_id_columns: Optional[List[str]] = None,
         labels: Optional[Dict[str, str]] = None,
         description: Optional[str] = None,
         project: Optional[str] = None,
@@ -107,9 +106,6 @@ def create(
         Args:
             name: The name of the feature group.
             source: The BigQuery source of the feature group.
-            entity_id_columns:
-                The entity ID columns. If not specified, defaults to
-                ['entity_id'].
             labels:
                 The labels with user-defined metadata to organize your
                 FeatureGroup.

diff --git a/vertexai/resources/preview/feature_store/utils.py b/vertexai/resources/preview/feature_store/utils.py
@@ -170,5 +170,7 @@ def as_dict(self) -> Dict[str, Any]:
 class FeatureGroupBigQuerySource:
     """BigQuery source for the Feature Group."""
 
+    # The URI for the BigQuery table/view.
     uri: str
-    entity_id_columns: List[str]
+    # The entity ID columns. If not specified, defaults to ['entity_id'].
+    entity_id_columns: Optional[List[str]] = None