Skip to content

Commit f9011e0

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: add FeatureGroup delete
PiperOrigin-RevId: 633836995
1 parent 710f33d commit f9011e0

File tree

2 files changed

+71
-0
lines changed

2 files changed

+71
-0
lines changed

tests/unit/vertexai/test_feature_group.py

+41
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,17 @@ def list_fg_mock():
108108
yield list_fg_mock
109109

110110

111+
@pytest.fixture
112+
def delete_fg_mock():
113+
with patch.object(
114+
feature_registry_service_client.FeatureRegistryServiceClient,
115+
"delete_feature_group",
116+
) as delete_fg_mock:
117+
delete_fg_lro_mock = mock.Mock(ga_operation.Operation)
118+
delete_fg_mock.return_value = delete_fg_lro_mock
119+
yield delete_fg_mock
120+
121+
111122
def fg_eq(
112123
fg_to_check: FeatureGroup,
113124
name: str,
@@ -293,3 +304,33 @@ def test_list(list_fg_mock):
293304
location=_TEST_LOCATION,
294305
labels=_TEST_FG3_LABELS,
295306
)
307+
308+
309+
@pytest.mark.parametrize("force", [True, False])
310+
def test_delete(force, delete_fg_mock, get_fg_mock, fg_logger_mock, sync=True):
311+
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
312+
313+
fg = FeatureGroup(_TEST_FG1_ID)
314+
fg.delete(force=force, sync=sync)
315+
316+
if not sync:
317+
fg.wait()
318+
319+
delete_fg_mock.assert_called_once_with(
320+
name=_TEST_FG1_PATH,
321+
force=force,
322+
)
323+
324+
fg_logger_mock.assert_has_calls(
325+
[
326+
call(
327+
"Deleting FeatureGroup resource: projects/test-project/locations/us-central1/featureGroups/my_fg1"
328+
),
329+
call(
330+
f"Delete FeatureGroup backing LRO: {delete_fg_mock.return_value.operation.name}"
331+
),
332+
call(
333+
"FeatureGroup resource projects/test-project/locations/us-central1/featureGroups/my_fg1 deleted."
334+
),
335+
]
336+
)

vertexai/resources/preview/feature_store/feature_group.py

+30
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,36 @@ def create(
208208

209209
return feature_group_obj
210210

211+
@base.optional_sync()
212+
def delete(self, force: bool = False, sync: bool = True) -> None:
213+
"""Deletes this feature group.
214+
215+
WARNING: This deletion is permanent.
216+
217+
Args:
218+
force:
219+
If set to True, all features under this online store will be
220+
deleted prior to online store deletion. Otherwise, deletion
221+
will only succeed if the online store has no FeatureViews.
222+
223+
If set to true, any Features under this FeatureGroup will also
224+
be deleted. (Otherwise, the request will only work if the
225+
FeatureGroup has no Features.)
226+
sync:
227+
Whether to execute this deletion synchronously. If False, this
228+
method will be executed in concurrent Future and any downstream
229+
object will be immediately returned and synced when the Future
230+
has completed.
231+
"""
232+
233+
lro = getattr(self.api_client, self._delete_method)(
234+
name=self.resource_name,
235+
force=force,
236+
)
237+
_LOGGER.log_delete_with_lro(self, lro)
238+
lro.result()
239+
_LOGGER.log_delete_complete(self)
240+
211241
@property
212242
def source(self) -> FeatureGroupBigQuerySource:
213243
return FeatureGroupBigQuerySource(

0 commit comments

Comments
 (0)