chore: test version column name

PiperOrigin-RevId: 635841448

tests/unit/vertexai/conftest.py

@@ -65,6 +65,7 @@
     _TEST_PSC_OPTIMIZED_FOS,
     _TEST_OPTIMIZED_EMBEDDING_FV,
     _TEST_FG1_F1,
+    _TEST_FG1_F2,
 )
 
 _TEST_PROJECT = "test-project"
@@ -510,3 +511,13 @@ def get_feature_mock():
     ) as get_fg_mock:
         get_fg_mock.return_value = _TEST_FG1_F1
         yield get_fg_mock
+
+
+@pytest.fixture
+def get_feature_with_version_column_mock():
+    with patch.object(
+        feature_registry_service_client.FeatureRegistryServiceClient,
+        "get_feature",
+    ) as get_fg_mock:
+        get_fg_mock.return_value = _TEST_FG1_F2
+        yield get_fg_mock

tests/unit/vertexai/feature_store_constants.py

@@ -328,3 +328,20 @@
     labels=_TEST_FG1_F1_LABELS,
     point_of_contact=_TEST_FG1_F1_POINT_OF_CONTACT,
 )
+
+
+_TEST_FG1_F2_ID = "my_fg1_f2"
+_TEST_FG1_F2_PATH = (
+    f"{_TEST_PARENT}/featureGroups/{_TEST_FG1_ID}/features/{_TEST_FG1_F2_ID}"
+)
+_TEST_FG1_F2_DESCRIPTION = "My feature 2 in feature group 1"
+_TEST_FG1_F2_LABELS = {"my_fg1_feature": "f2"}
+_TEST_FG1_F2_POINT_OF_CONTACT = "fg1-f2-announce-list"
+_TEST_FG1_F2_VERSION_COLUMN_NAME = "specific_column_for_feature_2"
+_TEST_FG1_F2 = types.feature.Feature(
+    name=_TEST_FG1_F2_PATH,
+    version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
+    description=_TEST_FG1_F2_DESCRIPTION,
+    labels=_TEST_FG1_F2_LABELS,
+    point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
+)

tests/unit/vertexai/test_feature.py

@@ -16,7 +16,7 @@
 #
 
 import re
-from typing import Dict
+from typing import Dict, Optional
 
 from google.cloud import aiplatform
 from google.cloud.aiplatform import base
@@ -35,6 +35,12 @@
     _TEST_FG1_F1_DESCRIPTION,
     _TEST_FG1_F1_LABELS,
     _TEST_FG1_F1_POINT_OF_CONTACT,
+    _TEST_FG1_F2_ID,
+    _TEST_FG1_F2_PATH,
+    _TEST_FG1_F2_VERSION_COLUMN_NAME,
+    _TEST_FG1_F2_DESCRIPTION,
+    _TEST_FG1_F2_LABELS,
+    _TEST_FG1_F2_POINT_OF_CONTACT,
 )
 
 
@@ -50,6 +56,7 @@ def feature_eq(
     description: str,
     labels: Dict[str, str],
     point_of_contact: str,
+    version_column_name: Optional[str] = None,
 ):
     """Check if a Feature has the appropriate values set."""
     assert feature_to_check.name == name
@@ -60,6 +67,9 @@ def feature_eq(
     assert feature_to_check.labels == labels
     assert feature_to_check.point_of_contact == point_of_contact
 
+    if version_column_name:
+        assert feature_to_check.version_column_name == version_column_name
+
 
 def test_init_with_feature_id_and_no_fg_id_raises_error(get_feature_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
@@ -108,6 +118,31 @@ def test_init_with_feature_id(get_feature_mock):
     )
 
 
+def test_init_with_feature_id_for_explicit_version_column(
+    get_feature_with_version_column_mock,
+):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    feature = Feature(_TEST_FG1_F2_ID, feature_group_id=_TEST_FG1_ID)
+
+    get_feature_with_version_column_mock.assert_called_once_with(
+        name=_TEST_FG1_F2_PATH,
+        retry=base._DEFAULT_RETRY,
+    )
+
+    feature_eq(
+        feature,
+        name=_TEST_FG1_F2_ID,
+        resource_name=_TEST_FG1_F2_PATH,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        description=_TEST_FG1_F2_DESCRIPTION,
+        labels=_TEST_FG1_F2_LABELS,
+        point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
+        version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
+    )
+
+
 def test_init_with_feature_path(get_feature_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
@@ -128,3 +163,28 @@ def test_init_with_feature_path(get_feature_mock):
         labels=_TEST_FG1_F1_LABELS,
         point_of_contact=_TEST_FG1_F1_POINT_OF_CONTACT,
     )
+
+
+def test_init_with_feature_path_for_explicit_version_column(
+    get_feature_with_version_column_mock,
+):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    feature = Feature(_TEST_FG1_F2_PATH)
+
+    get_feature_with_version_column_mock.assert_called_once_with(
+        name=_TEST_FG1_F2_PATH,
+        retry=base._DEFAULT_RETRY,
+    )
+
+    feature_eq(
+        feature,
+        name=_TEST_FG1_F2_ID,
+        resource_name=_TEST_FG1_F2_PATH,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
+        description=_TEST_FG1_F2_DESCRIPTION,
+        labels=_TEST_FG1_F2_LABELS,
+        point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
+    )

tests/unit/vertexai/test_feature_group.py

@@ -66,6 +66,13 @@
     _TEST_FG1_F1_DESCRIPTION,
     _TEST_FG1_F1_LABELS,
     _TEST_FG1_F1_POINT_OF_CONTACT,
+    _TEST_FG1_F2,
+    _TEST_FG1_F2_ID,
+    _TEST_FG1_F2_PATH,
+    _TEST_FG1_F2_DESCRIPTION,
+    _TEST_FG1_F2_LABELS,
+    _TEST_FG1_F2_POINT_OF_CONTACT,
+    _TEST_FG1_F2_VERSION_COLUMN_NAME,
 )
 from test_feature import feature_eq
 
@@ -138,6 +145,18 @@ def create_feature_mock():
         yield create_feature_mock
 
 
+@pytest.fixture
+def create_feature_with_version_column_mock():
+    with patch.object(
+        feature_registry_service_client.FeatureRegistryServiceClient,
+        "create_feature",
+    ) as create_feature_mock:
+        create_feature_lro_mock = mock.Mock(ga_operation.Operation)
+        create_feature_lro_mock.result.return_value = _TEST_FG1_F2
+        create_feature_mock.return_value = create_feature_lro_mock
+        yield create_feature_mock
+
+
 def fg_eq(
     fg_to_check: FeatureGroup,
     name: str,
@@ -442,3 +461,72 @@ def test_create_feature(
             ),
         ]
     )
+
+
+@pytest.mark.parametrize("create_request_timeout", [None, 1.0])
+@pytest.mark.parametrize("sync", [True, False])
+def test_create_feature_with_version_feature_column(
+    get_fg_mock,
+    create_feature_with_version_column_mock,
+    get_feature_with_version_column_mock,
+    fg_logger_mock,
+    create_request_timeout,
+    sync,
+):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    fg = FeatureGroup(_TEST_FG1_ID)
+    feature = fg.create_feature(
+        _TEST_FG1_F2_ID,
+        version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
+        description=_TEST_FG1_F2_DESCRIPTION,
+        labels=_TEST_FG1_F2_LABELS,
+        point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
+        create_request_timeout=create_request_timeout,
+        sync=sync,
+    )
+
+    if not sync:
+        feature.wait()
+
+    expected_feature = types.feature.Feature(
+        version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
+        description=_TEST_FG1_F2_DESCRIPTION,
+        labels=_TEST_FG1_F2_LABELS,
+        point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
+    )
+    create_feature_with_version_column_mock.assert_called_once_with(
+        parent=_TEST_FG1_PATH,
+        feature=expected_feature,
+        feature_id=_TEST_FG1_F2_ID,
+        metadata=(),
+        timeout=create_request_timeout,
+    )
+
+    feature_eq(
+        feature,
+        name=_TEST_FG1_F2_ID,
+        resource_name=_TEST_FG1_F2_PATH,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        description=_TEST_FG1_F2_DESCRIPTION,
+        labels=_TEST_FG1_F2_LABELS,
+        point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
+        version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
+    )
+
+    fg_logger_mock.assert_has_calls(
+        [
+            call("Creating Feature"),
+            call(
+                f"Create Feature backing LRO: {create_feature_with_version_column_mock.return_value.operation.name}"
+            ),
+            call(
+                "Feature created. Resource name: projects/test-project/locations/us-central1/featureGroups/my_fg1/features/my_fg1_f2"
+            ),
+            call("To use this Feature in another session:"),
+            call(
+                "feature = aiplatform.Feature('projects/test-project/locations/us-central1/featureGroups/my_fg1/features/my_fg1_f2')"
+            ),
+        ]
+    )

vertexai/resources/preview/feature_store/feature.py

@@ -105,6 +105,11 @@ def __init__(
 
         self._gca_resource = self._get_gca_resource(resource_name=feature)
 
+    @property
+    def version_column_name(self) -> str:
+        """The name of the BigQuery Table/View column hosting data for this version."""
+        return self._gca_resource.version_column_name
+
     @property
     def description(self) -> str:
         """The description of the feature."""