Skip to content

Commit fe78170

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
chore: test version column name
PiperOrigin-RevId: 635841448
1 parent 555ead7 commit fe78170

5 files changed

Lines changed: 182 additions & 1 deletion

File tree

‎tests/unit/vertexai/conftest.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@
6565
_TEST_PSC_OPTIMIZED_FOS,
6666
_TEST_OPTIMIZED_EMBEDDING_FV,
6767
_TEST_FG1_F1,
68+
_TEST_FG1_F2,
6869
)
6970

7071
_TEST_PROJECT = "test-project"
@@ -510,3 +511,13 @@ def get_feature_mock():
510511
) as get_fg_mock:
511512
get_fg_mock.return_value = _TEST_FG1_F1
512513
yield get_fg_mock
514+
515+
516+
@pytest.fixture
517+
def get_feature_with_version_column_mock():
518+
with patch.object(
519+
feature_registry_service_client.FeatureRegistryServiceClient,
520+
"get_feature",
521+
) as get_fg_mock:
522+
get_fg_mock.return_value = _TEST_FG1_F2
523+
yield get_fg_mock

‎tests/unit/vertexai/feature_store_constants.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -328,3 +328,20 @@
328328
labels=_TEST_FG1_F1_LABELS,
329329
point_of_contact=_TEST_FG1_F1_POINT_OF_CONTACT,
330330
)
331+
332+
333+
_TEST_FG1_F2_ID = "my_fg1_f2"
334+
_TEST_FG1_F2_PATH = (
335+
f"{_TEST_PARENT}/featureGroups/{_TEST_FG1_ID}/features/{_TEST_FG1_F2_ID}"
336+
)
337+
_TEST_FG1_F2_DESCRIPTION = "My feature 2 in feature group 1"
338+
_TEST_FG1_F2_LABELS = {"my_fg1_feature": "f2"}
339+
_TEST_FG1_F2_POINT_OF_CONTACT = "fg1-f2-announce-list"
340+
_TEST_FG1_F2_VERSION_COLUMN_NAME = "specific_column_for_feature_2"
341+
_TEST_FG1_F2 = types.feature.Feature(
342+
name=_TEST_FG1_F2_PATH,
343+
version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
344+
description=_TEST_FG1_F2_DESCRIPTION,
345+
labels=_TEST_FG1_F2_LABELS,
346+
point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
347+
)

‎tests/unit/vertexai/test_feature.py‎

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
#
1717

1818
import re
19-
from typing import Dict
19+
from typing import Dict, Optional
2020

2121
from google.cloud import aiplatform
2222
from google.cloud.aiplatform import base
@@ -35,6 +35,12 @@
3535
_TEST_FG1_F1_DESCRIPTION,
3636
_TEST_FG1_F1_LABELS,
3737
_TEST_FG1_F1_POINT_OF_CONTACT,
38+
_TEST_FG1_F2_ID,
39+
_TEST_FG1_F2_PATH,
40+
_TEST_FG1_F2_VERSION_COLUMN_NAME,
41+
_TEST_FG1_F2_DESCRIPTION,
42+
_TEST_FG1_F2_LABELS,
43+
_TEST_FG1_F2_POINT_OF_CONTACT,
3844
)
3945

4046

@@ -50,6 +56,7 @@ def feature_eq(
5056
description: str,
5157
labels: Dict[str, str],
5258
point_of_contact: str,
59+
version_column_name: Optional[str] = None,
5360
):
5461
"""Check if a Feature has the appropriate values set."""
5562
assert feature_to_check.name == name
@@ -60,6 +67,9 @@ def feature_eq(
6067
assert feature_to_check.labels == labels
6168
assert feature_to_check.point_of_contact == point_of_contact
6269

70+
if version_column_name:
71+
assert feature_to_check.version_column_name == version_column_name
72+
6373

6474
def test_init_with_feature_id_and_no_fg_id_raises_error(get_feature_mock):
6575
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
@@ -108,6 +118,31 @@ def test_init_with_feature_id(get_feature_mock):
108118
)
109119

110120

121+
def test_init_with_feature_id_for_explicit_version_column(
122+
get_feature_with_version_column_mock,
123+
):
124+
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
125+
126+
feature = Feature(_TEST_FG1_F2_ID, feature_group_id=_TEST_FG1_ID)
127+
128+
get_feature_with_version_column_mock.assert_called_once_with(
129+
name=_TEST_FG1_F2_PATH,
130+
retry=base._DEFAULT_RETRY,
131+
)
132+
133+
feature_eq(
134+
feature,
135+
name=_TEST_FG1_F2_ID,
136+
resource_name=_TEST_FG1_F2_PATH,
137+
project=_TEST_PROJECT,
138+
location=_TEST_LOCATION,
139+
description=_TEST_FG1_F2_DESCRIPTION,
140+
labels=_TEST_FG1_F2_LABELS,
141+
point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
142+
version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
143+
)
144+
145+
111146
def test_init_with_feature_path(get_feature_mock):
112147
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
113148

@@ -128,3 +163,28 @@ def test_init_with_feature_path(get_feature_mock):
128163
labels=_TEST_FG1_F1_LABELS,
129164
point_of_contact=_TEST_FG1_F1_POINT_OF_CONTACT,
130165
)
166+
167+
168+
def test_init_with_feature_path_for_explicit_version_column(
169+
get_feature_with_version_column_mock,
170+
):
171+
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
172+
173+
feature = Feature(_TEST_FG1_F2_PATH)
174+
175+
get_feature_with_version_column_mock.assert_called_once_with(
176+
name=_TEST_FG1_F2_PATH,
177+
retry=base._DEFAULT_RETRY,
178+
)
179+
180+
feature_eq(
181+
feature,
182+
name=_TEST_FG1_F2_ID,
183+
resource_name=_TEST_FG1_F2_PATH,
184+
project=_TEST_PROJECT,
185+
location=_TEST_LOCATION,
186+
version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
187+
description=_TEST_FG1_F2_DESCRIPTION,
188+
labels=_TEST_FG1_F2_LABELS,
189+
point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
190+
)

‎tests/unit/vertexai/test_feature_group.py‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,13 @@
6666
_TEST_FG1_F1_DESCRIPTION,
6767
_TEST_FG1_F1_LABELS,
6868
_TEST_FG1_F1_POINT_OF_CONTACT,
69+
_TEST_FG1_F2,
70+
_TEST_FG1_F2_ID,
71+
_TEST_FG1_F2_PATH,
72+
_TEST_FG1_F2_DESCRIPTION,
73+
_TEST_FG1_F2_LABELS,
74+
_TEST_FG1_F2_POINT_OF_CONTACT,
75+
_TEST_FG1_F2_VERSION_COLUMN_NAME,
6976
)
7077
from test_feature import feature_eq
7178

@@ -138,6 +145,18 @@ def create_feature_mock():
138145
yield create_feature_mock
139146

140147

148+
@pytest.fixture
149+
def create_feature_with_version_column_mock():
150+
with patch.object(
151+
feature_registry_service_client.FeatureRegistryServiceClient,
152+
"create_feature",
153+
) as create_feature_mock:
154+
create_feature_lro_mock = mock.Mock(ga_operation.Operation)
155+
create_feature_lro_mock.result.return_value = _TEST_FG1_F2
156+
create_feature_mock.return_value = create_feature_lro_mock
157+
yield create_feature_mock
158+
159+
141160
def fg_eq(
142161
fg_to_check: FeatureGroup,
143162
name: str,
@@ -442,3 +461,72 @@ def test_create_feature(
442461
),
443462
]
444463
)
464+
465+
466+
@pytest.mark.parametrize("create_request_timeout", [None, 1.0])
467+
@pytest.mark.parametrize("sync", [True, False])
468+
def test_create_feature_with_version_feature_column(
469+
get_fg_mock,
470+
create_feature_with_version_column_mock,
471+
get_feature_with_version_column_mock,
472+
fg_logger_mock,
473+
create_request_timeout,
474+
sync,
475+
):
476+
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
477+
478+
fg = FeatureGroup(_TEST_FG1_ID)
479+
feature = fg.create_feature(
480+
_TEST_FG1_F2_ID,
481+
version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
482+
description=_TEST_FG1_F2_DESCRIPTION,
483+
labels=_TEST_FG1_F2_LABELS,
484+
point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
485+
create_request_timeout=create_request_timeout,
486+
sync=sync,
487+
)
488+
489+
if not sync:
490+
feature.wait()
491+
492+
expected_feature = types.feature.Feature(
493+
version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
494+
description=_TEST_FG1_F2_DESCRIPTION,
495+
labels=_TEST_FG1_F2_LABELS,
496+
point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
497+
)
498+
create_feature_with_version_column_mock.assert_called_once_with(
499+
parent=_TEST_FG1_PATH,
500+
feature=expected_feature,
501+
feature_id=_TEST_FG1_F2_ID,
502+
metadata=(),
503+
timeout=create_request_timeout,
504+
)
505+
506+
feature_eq(
507+
feature,
508+
name=_TEST_FG1_F2_ID,
509+
resource_name=_TEST_FG1_F2_PATH,
510+
project=_TEST_PROJECT,
511+
location=_TEST_LOCATION,
512+
description=_TEST_FG1_F2_DESCRIPTION,
513+
labels=_TEST_FG1_F2_LABELS,
514+
point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
515+
version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
516+
)
517+
518+
fg_logger_mock.assert_has_calls(
519+
[
520+
call("Creating Feature"),
521+
call(
522+
f"Create Feature backing LRO: {create_feature_with_version_column_mock.return_value.operation.name}"
523+
),
524+
call(
525+
"Feature created. Resource name: projects/test-project/locations/us-central1/featureGroups/my_fg1/features/my_fg1_f2"
526+
),
527+
call("To use this Feature in another session:"),
528+
call(
529+
"feature = aiplatform.Feature('projects/test-project/locations/us-central1/featureGroups/my_fg1/features/my_fg1_f2')"
530+
),
531+
]
532+
)

‎vertexai/resources/preview/feature_store/feature.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,11 @@ def __init__(
105105

106106
self._gca_resource = self._get_gca_resource(resource_name=feature)
107107

108+
@property
109+
def version_column_name(self) -> str:
110+
"""The name of the BigQuery Table/View column hosting data for this version."""
111+
return self._gca_resource.version_column_name
112+
108113
@property
109114
def description(self) -> str:
110115
"""The description of the feature."""

0 commit comments

Comments
 (0)