1+ import json
12import logging
23from typing import Dict , List , Optional
34
3233logger = logging .getLogger (__name__ )
3334
3435
36+ def _build_feature_view_to_datasource_map (
37+ grpc_handler , project : str , allow_cache : bool = True
38+ ) -> Dict [str , str ]:
39+ """Build a map of feature_view_name → data_source_name for a project."""
40+ try :
41+ fv_req = RegistryServer_pb2 .ListAllFeatureViewsRequest (
42+ project = project , allow_cache = allow_cache
43+ )
44+ fv_response = grpc_call (grpc_handler .ListAllFeatureViews , fv_req )
45+ feature_views = fv_response .get ("featureViews" , [])
46+ except Exception :
47+ return {}
48+
49+ fv_ds_map : Dict [str , str ] = {}
50+ for any_fv in feature_views :
51+ # AnyFeatureView wraps under "featureView", "onDemandFeatureView", etc.
52+ fv_data = any_fv .get ("featureView" ) or any_fv .get ("streamFeatureView" ) or {}
53+ spec = fv_data .get ("spec" , {})
54+ fv_name = spec .get ("name" , "" )
55+ if not fv_name :
56+ continue
57+ batch_source = spec .get ("batchSource" , {})
58+ ds_name = batch_source .get ("name" , "" ) if batch_source else ""
59+ if ds_name :
60+ fv_ds_map [fv_name ] = ds_name
61+ return fv_ds_map
62+
63+
64+ def _get_dataset_data_sources (dataset : dict , fv_ds_map : Dict [str , str ]) -> List [str ]:
65+ """Get data source names for a dataset via its features' feature view lineage."""
66+ spec = dataset .get ("spec" , dataset )
67+ features = spec .get ("features" , [])
68+ data_sources : List [str ] = []
69+ seen : set = set ()
70+ for feature_ref in features :
71+ if isinstance (feature_ref , str ) and ":" in feature_ref :
72+ fv_name = feature_ref .split (":" )[0 ]
73+ ds_name = fv_ds_map .get (fv_name )
74+ if ds_name and ds_name not in seen :
75+ seen .add (ds_name )
76+ data_sources .append (ds_name )
77+ return data_sources
78+
79+
3580class RegisterDatasetRequest (BaseModel ):
3681 name : str
3782 project : str
3883 features : List [str ] = []
3984 join_keys : List [str ] = []
4085 storage_path : Optional [str ] = None
4186 storage_type : str = "file"
87+ storage_file_format : Optional [str ] = "parquet"
4288 tags : Dict [str , str ] = {}
4389 full_feature_names : bool = False
4490 feature_service_name : Optional [str ] = None
@@ -64,6 +110,7 @@ class CreateDatasetRequest(BaseModel):
64110 # Storage
65111 storage_type : str = "file"
66112 storage_path : str = ""
113+ storage_file_format : Optional [str ] = "parquet"
67114 tags : Dict [str , str ] = {}
68115 allow_overwrite : bool = False
69116
@@ -94,7 +141,7 @@ def list_saved_datasets_all(
94141 if collection is not None :
95142 extra_params ["collection" ] = collection
96143
97- return aggregate_across_projects (
144+ result = aggregate_across_projects (
98145 grpc_handler = grpc_handler ,
99146 list_method = grpc_handler .ListSavedDatasets ,
100147 request_cls = RegistryServer_pb2 .ListSavedDatasetsRequest ,
@@ -109,6 +156,21 @@ def list_saved_datasets_all(
109156 extra_request_params = extra_params or None ,
110157 )
111158
159+ saved_datasets = result .get ("savedDatasets" , [])
160+ projects_seen : Dict [str , Dict [str , str ]] = {}
161+ for sd in saved_datasets :
162+ proj = sd .get ("project" , "" )
163+ if proj not in projects_seen :
164+ projects_seen [proj ] = _build_feature_view_to_datasource_map (
165+ grpc_handler , proj , allow_cache
166+ )
167+ ds_names = _get_dataset_data_sources (sd , projects_seen [proj ])
168+ if ds_names :
169+ spec = sd .get ("spec" , sd )
170+ spec ["dataSources" ] = ds_names
171+
172+ return result
173+
112174 @router .get ("/saved_datasets/data/{name}" )
113175 def get_dataset_data (
114176 name : str ,
@@ -209,6 +271,14 @@ def get_saved_dataset(
209271
210272 result = saved_dataset
211273
274+ fv_ds_map = _build_feature_view_to_datasource_map (
275+ grpc_handler , project , allow_cache
276+ )
277+ ds_names = _get_dataset_data_sources (result , fv_ds_map )
278+ if ds_names :
279+ spec = result .get ("spec" , result )
280+ spec ["dataSources" ] = ds_names
281+
212282 if include_relationships :
213283 relationships = get_object_relationships (
214284 grpc_handler , "savedDataset" , name , project , allow_cache
@@ -270,6 +340,15 @@ def list_saved_datasets(
270340 response = grpc_call (grpc_handler .ListSavedDatasets , req )
271341 saved_datasets = response .get ("savedDatasets" , [])
272342
343+ fv_ds_map = _build_feature_view_to_datasource_map (
344+ grpc_handler , project , allow_cache
345+ )
346+ for sd in saved_datasets :
347+ ds_names = _get_dataset_data_sources (sd , fv_ds_map )
348+ if ds_names :
349+ spec = sd .get ("spec" , sd )
350+ spec ["dataSources" ] = ds_names
351+
273352 result = {
274353 "savedDatasets" : saved_datasets ,
275354 "pagination" : response .get ("pagination" , {}),
@@ -332,12 +411,48 @@ def register_saved_dataset(payload: RegisterDatasetRequest = Body(...)):
332411 )
333412 elif payload .storage_type == "spark" :
334413 storage_proto .spark_storage .CopyFrom (
335- DataSourceProto .SparkOptions (path = path )
414+ DataSourceProto .SparkOptions (
415+ path = path ,
416+ file_format = payload .storage_file_format or "parquet" ,
417+ )
336418 )
337419 elif payload .storage_type == "athena" :
338420 storage_proto .athena_storage .CopyFrom (
339421 DataSourceProto .AthenaOptions (table = path )
340422 )
423+ elif payload .storage_type == "postgres" :
424+ config = json .dumps ({"name" : None , "query" : None , "table" : path })
425+ storage_proto .custom_storage .CopyFrom (
426+ DataSourceProto .CustomSourceOptions (
427+ configuration = config .encode ("utf-8" )
428+ )
429+ )
430+ elif payload .storage_type == "clickhouse" :
431+ config = json .dumps ({"name" : None , "query" : None , "table" : path })
432+ storage_proto .custom_storage .CopyFrom (
433+ DataSourceProto .CustomSourceOptions (
434+ configuration = config .encode ("utf-8" )
435+ )
436+ )
437+ elif payload .storage_type == "couchbase" :
438+ parts = path .split ("." , 2 )
439+ database = parts [0 ] if len (parts ) > 0 else ""
440+ scope = parts [1 ] if len (parts ) > 1 else ""
441+ collection = parts [2 ] if len (parts ) > 2 else ""
442+ config = json .dumps (
443+ {
444+ "name" : None ,
445+ "query" : None ,
446+ "database" : database ,
447+ "scope" : scope ,
448+ "collection" : collection ,
449+ }
450+ )
451+ storage_proto .custom_storage .CopyFrom (
452+ DataSourceProto .CustomSourceOptions (
453+ configuration = config .encode ("utf-8" )
454+ )
455+ )
341456 elif payload .storage_type == "custom" :
342457 storage_proto .custom_storage .CopyFrom (
343458 DataSourceProto .CustomSourceOptions (configuration = path .encode ("utf-8" ))
0 commit comments