Skip to content

Commit 165fcf2

Browse files
ntkatholentkathole
committed
feat: Added optional namespace/colleciton to datasets
Signed-off-by: ntkathole <nikhilkathole2683@gmail.com>
1 parent 955403e commit 165fcf2

17 files changed

Lines changed: 266 additions & 104 deletions

‎protos/feast/core/SavedDataset.proto‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,19 @@ message SavedDatasetSpec {
4848

4949
// User defined metadata
5050
map<string, string> tags = 7;
51+
52+
// Optional logical namespace for hierarchical grouping.
53+
// Maps to the top-level prefix in Iceberg REST Catalog API.
54+
// Empty string means not set (no namespace scoping).
55+
string namespace = 9;
56+
57+
// Optional sub-grouping within a namespace.
58+
// Maps to the namespace level in Iceberg REST Catalog API.
59+
// Empty string means not set (dataset sits directly under namespace).
60+
string collection = 10;
61+
62+
// Description of the saved dataset.
63+
string description = 11;
5164
}
5265

5366
message SavedDatasetStorage {

‎protos/feast/registry/RegistryServer.proto‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -417,6 +417,10 @@ message ListSavedDatasetsRequest {
417417
map<string,string> tags = 3;
418418
PaginationParams pagination = 4;
419419
SortingParams sorting = 5;
420+
// Optional logical namespace filter. Empty string means no filter.
421+
string namespace = 6;
422+
// Optional collection filter. Empty string means no filter.
423+
string collection = 7;
420424
}
421425

422426
message ListSavedDatasetsResponse {

‎sdk/python/feast/api/registry/rest/saved_datasets.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,9 @@ class RegisterDatasetRequest(BaseModel):
4343
full_feature_names: bool = False
4444
feature_service_name: Optional[str] = None
4545
allow_override: bool = False
46+
namespace: str = ""
47+
collection: str = ""
48+
description: str = ""
4649

4750

4851
class CreateDatasetRequest(BaseModel):
@@ -75,10 +78,22 @@ def list_saved_datasets_all(
7578
limit: int = Query(50, ge=1, le=100),
7679
sort_by: str = Query(None),
7780
sort_order: str = Query("asc"),
81+
namespace: Optional[str] = Query(
82+
None, description="Filter by namespace (logical grouping)"
83+
),
84+
collection: Optional[str] = Query(
85+
None, description="Filter by collection (sub-grouping within namespace)"
86+
),
7887
include_relationships: bool = Query(
7988
False, description="Include relationships for each saved dataset"
8089
),
8190
):
91+
extra_params = {}
92+
if namespace is not None:
93+
extra_params["namespace"] = namespace
94+
if collection is not None:
95+
extra_params["collection"] = collection
96+
8297
return aggregate_across_projects(
8398
grpc_handler=grpc_handler,
8499
list_method=grpc_handler.ListSavedDatasets,
@@ -91,6 +106,7 @@ def list_saved_datasets_all(
91106
sort_by=sort_by,
92107
sort_order=sort_order,
93108
include_relationships=include_relationships,
109+
extra_request_params=extra_params or None,
94110
)
95111

96112
@router.get("/saved_datasets/data/{name}")
@@ -230,6 +246,12 @@ def list_saved_datasets(
230246
project: str = Query(...),
231247
allow_cache: bool = Query(default=True),
232248
tags: Dict[str, str] = Depends(parse_tags),
249+
namespace: Optional[str] = Query(
250+
None, description="Filter by namespace (logical grouping)"
251+
),
252+
collection: Optional[str] = Query(
253+
None, description="Filter by collection (sub-grouping within namespace)"
254+
),
233255
include_relationships: bool = Query(
234256
False, description="Include relationships for each saved dataset"
235257
),
@@ -240,6 +262,8 @@ def list_saved_datasets(
240262
project=project,
241263
allow_cache=allow_cache,
242264
tags=tags,
265+
namespace=namespace or "",
266+
collection=collection or "",
243267
pagination=create_grpc_pagination_params(pagination_params),
244268
sorting=create_grpc_sorting_params(sorting_params),
245269
)
@@ -347,6 +371,9 @@ def register_saved_dataset(payload: RegisterDatasetRequest = Body(...)):
347371
full_feature_names=payload.full_feature_names,
348372
storage=storage_proto,
349373
tags=payload.tags,
374+
namespace=payload.namespace,
375+
collection=payload.collection,
376+
description=payload.description,
350377
)
351378
if payload.feature_service_name:
352379
spec.feature_service_name = payload.feature_service_name

‎sdk/python/feast/feature_store.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4770,20 +4770,30 @@ def delete_project(self, name: str, commit: bool = True) -> None:
47704770
return self.registry.delete_project(name, commit=commit)
47714771

47724772
def list_saved_datasets(
4773-
self, allow_cache: bool = False, tags: Optional[dict[str, str]] = None
4773+
self,
4774+
allow_cache: bool = False,
4775+
tags: Optional[dict[str, str]] = None,
4776+
namespace: Optional[str] = None,
4777+
collection: Optional[str] = None,
47744778
) -> List[SavedDataset]:
47754779
"""
47764780
Retrieves the list of saved datasets from the registry.
47774781
47784782
Args:
47794783
allow_cache: Whether to allow returning saved datasets from a cached registry.
47804784
tags: Filter by tags.
4785+
namespace: Filter by logical namespace grouping.
4786+
collection: Filter by collection sub-grouping within namespace.
47814787
47824788
Returns:
47834789
A list of saved datasets.
47844790
"""
47854791
return self.registry.list_saved_datasets(
4786-
self.project, allow_cache=allow_cache, tags=tags
4792+
self.project,
4793+
allow_cache=allow_cache,
4794+
tags=tags,
4795+
namespace=namespace,
4796+
collection=collection,
47874797
)
47884798

47894799
async def initialize(self) -> None:

‎sdk/python/feast/infra/registry/base_registry.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -698,6 +698,8 @@ def list_saved_datasets(
698698
project: str,
699699
allow_cache: bool = False,
700700
tags: Optional[dict[str, str]] = None,
701+
namespace: Optional[str] = None,
702+
collection: Optional[str] = None,
701703
) -> List[SavedDataset]:
702704
"""
703705
Retrieves a list of all saved datasets in specified project
@@ -706,6 +708,8 @@ def list_saved_datasets(
706708
project: Feast project
707709
allow_cache: Whether to allow returning this dataset from a cached registry
708710
tags: Filter by tags
711+
namespace: Filter by logical namespace grouping
712+
collection: Filter by collection sub-grouping within namespace
709713
710714
Returns:
711715
Returns the list of SavedDatasets

‎sdk/python/feast/infra/registry/caching_registry.py‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -344,7 +344,11 @@ def get_saved_dataset(
344344

345345
@abstractmethod
346346
def _list_saved_datasets(
347-
self, project: str, tags: Optional[dict[str, str]] = None
347+
self,
348+
project: str,
349+
tags: Optional[dict[str, str]] = None,
350+
namespace: Optional[str] = None,
351+
collection: Optional[str] = None,
348352
) -> List[SavedDataset]:
349353
pass
350354

@@ -353,13 +357,21 @@ def list_saved_datasets(
353357
project: str,
354358
allow_cache: bool = False,
355359
tags: Optional[dict[str, str]] = None,
360+
namespace: Optional[str] = None,
361+
collection: Optional[str] = None,
356362
) -> List[SavedDataset]:
357363
if allow_cache:
358364
self._refresh_cached_registry_if_necessary()
359365
return proto_registry_utils.list_saved_datasets(
360-
self.cached_registry_proto, project, tags
366+
self.cached_registry_proto,
367+
project,
368+
tags,
369+
namespace=namespace,
370+
collection=collection,
361371
)
362-
return self._list_saved_datasets(project, tags)
372+
return self._list_saved_datasets(
373+
project, tags, namespace=namespace, collection=collection
374+
)
363375

364376
@abstractmethod
365377
def _get_validation_reference(self, name: str, project: str) -> ValidationReference:

‎sdk/python/feast/infra/registry/proto_registry_utils.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -418,14 +418,23 @@ def list_data_sources(
418418

419419
@registry_proto_cache_with_tags
420420
def list_saved_datasets(
421-
registry_proto: RegistryProto, project: str, tags: Optional[dict[str, str]]
421+
registry_proto: RegistryProto,
422+
project: str,
423+
tags: Optional[dict[str, str]],
424+
namespace: Optional[str] = None,
425+
collection: Optional[str] = None,
422426
) -> List[SavedDataset]:
423427
saved_datasets = []
424428
for saved_dataset in registry_proto.saved_datasets:
425-
if saved_dataset.spec.project == project and utils.has_all_tags(
426-
saved_dataset.spec.tags, tags
427-
):
428-
saved_datasets.append(SavedDataset.from_proto(saved_dataset))
429+
if saved_dataset.spec.project != project:
430+
continue
431+
if not utils.has_all_tags(saved_dataset.spec.tags, tags):
432+
continue
433+
if namespace is not None and saved_dataset.spec.namespace != namespace:
434+
continue
435+
if collection is not None and saved_dataset.spec.collection != collection:
436+
continue
437+
saved_datasets.append(SavedDataset.from_proto(saved_dataset))
429438
return saved_datasets
430439

431440

‎sdk/python/feast/infra/registry/registry.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1305,11 +1305,19 @@ def list_saved_datasets(
13051305
project: str,
13061306
allow_cache: bool = False,
13071307
tags: Optional[dict[str, str]] = None,
1308+
namespace: Optional[str] = None,
1309+
collection: Optional[str] = None,
13081310
) -> List[SavedDataset]:
13091311
registry_proto = self._get_registry_proto(
13101312
project=project, allow_cache=allow_cache
13111313
)
1312-
return proto_registry_utils.list_saved_datasets(registry_proto, project, tags)
1314+
return proto_registry_utils.list_saved_datasets(
1315+
registry_proto,
1316+
project,
1317+
tags,
1318+
namespace=namespace,
1319+
collection=collection,
1320+
)
13131321

13141322
def delete_saved_dataset(self, name: str, project: str, commit: bool = True):
13151323
self._prepare_registry_for_changes(project)

‎sdk/python/feast/infra/registry/remote.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -511,9 +511,15 @@ def list_saved_datasets(
511511
project: str,
512512
allow_cache: bool = False,
513513
tags: Optional[dict[str, str]] = None,
514+
namespace: Optional[str] = None,
515+
collection: Optional[str] = None,
514516
) -> List[SavedDataset]:
515517
request = RegistryServer_pb2.ListSavedDatasetsRequest(
516-
project=project, allow_cache=allow_cache, tags=tags
518+
project=project,
519+
allow_cache=allow_cache,
520+
tags=tags,
521+
namespace=namespace or "",
522+
collection=collection or "",
517523
)
518524
response = self.stub.ListSavedDatasets(request)
519525
return [

‎sdk/python/feast/infra/registry/snowflake.py‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -945,20 +945,31 @@ def list_saved_datasets(
945945
project: str,
946946
allow_cache: bool = False,
947947
tags: Optional[dict[str, str]] = None,
948+
namespace: Optional[str] = None,
949+
collection: Optional[str] = None,
948950
) -> List[SavedDataset]:
949951
if allow_cache:
950952
registry_proto = self._refresh_cached_registry_if_necessary()
951953
return proto_registry_utils.list_saved_datasets(
952-
registry_proto, project, tags
954+
registry_proto,
955+
project,
956+
tags,
957+
namespace=namespace,
958+
collection=collection,
953959
)
954-
return self._list_objects(
960+
results = self._list_objects(
955961
"SAVED_DATASETS",
956962
project,
957963
SavedDatasetProto,
958964
SavedDataset,
959965
"SAVED_DATASET_PROTO",
960966
tags=tags,
961967
)
968+
if namespace is not None:
969+
results = [sd for sd in results if sd.namespace == namespace]
970+
if collection is not None:
971+
results = [sd for sd in results if sd.collection == collection]
972+
return results
962973

963974
def list_stream_feature_views(
964975
self,

0 commit comments

Comments
 (0)