diff --git a/sdk/python/feast/api/registry/rest/data_sources.py b/sdk/python/feast/api/registry/rest/data_sources.py index 4e522644ed3..ed0f6f983fa 100644 --- a/sdk/python/feast/api/registry/rest/data_sources.py +++ b/sdk/python/feast/api/registry/rest/data_sources.py @@ -31,6 +31,8 @@ class FileOptionsModel(BaseModel): uri: str = "" + file_format: Optional[str] = "parquet" + s3_endpoint_override: Optional[str] = "" class BigQueryOptionsModel(BaseModel): @@ -42,22 +44,59 @@ class SnowflakeOptionsModel(BaseModel): table: str = "" database: str = "" schema_: str = "" + query: Optional[str] = "" + warehouse: Optional[str] = "" class RedshiftOptionsModel(BaseModel): table: str = "" database: str = "" schema_: str = "" + query: Optional[str] = "" class KafkaOptionsModel(BaseModel): kafka_bootstrap_servers: str = "" topic: str = "" + message_format: Optional[str] = "json" + watermark_delay_threshold: Optional[str] = "" class SparkOptionsModel(BaseModel): table: str = "" path: str = "" + query: Optional[str] = "" + file_format: Optional[str] = "" + table_format: Optional[str] = "" + table_format_catalog: Optional[str] = "" + table_format_namespace: Optional[str] = "" + table_format_properties: Optional[str] = "" + date_partition_column: Optional[str] = "" + date_partition_column_format: Optional[str] = "" + + +class KinesisOptionsModel(BaseModel): + region: str = "" + stream_name: str = "" + record_format: Optional[str] = "json" + + +class TrinoOptionsModel(BaseModel): + table: str = "" + query: str = "" + + +class AthenaOptionsModel(BaseModel): + table: str = "" + query: str = "" + database: str = "" + data_source: Optional[str] = "" + + +class CustomOptionsModel(BaseModel): + configuration: Optional[str] = "" + class_name: Optional[str] = "" + config: Optional[str] = "" class ApplyDataSourceRequestBody(BaseModel): @@ -66,6 +105,7 @@ class ApplyDataSourceRequestBody(BaseModel): type: Optional[int] = None timestamp_field: Optional[str] = "" created_timestamp_column: Optional[str] = "" + date_partition_column: Optional[str] = "" description: Optional[str] = "" tags: Optional[Dict[str, str]] = {} owner: Optional[str] = "" @@ -75,6 +115,11 @@ class ApplyDataSourceRequestBody(BaseModel): redshift_options: Optional[RedshiftOptionsModel] = None kafka_options: Optional[KafkaOptionsModel] = None spark_options: Optional[SparkOptionsModel] = None + trino_options: Optional[TrinoOptionsModel] = None + athena_options: Optional[AthenaOptionsModel] = None + kinesis_options: Optional[KinesisOptionsModel] = None + custom_options: Optional[CustomOptionsModel] = None + data_source_class_type: Optional[str] = None def get_data_source_router(grpc_handler) -> APIRouter: @@ -221,8 +266,15 @@ def apply_data_source(body: ApplyDataSourceRequestBody): if body.type is not None: ds_proto.type = body.type # type: ignore[assignment] + if body.date_partition_column: + ds_proto.date_partition_column = body.date_partition_column + if body.file_options: ds_proto.file_options.uri = body.file_options.uri + if body.file_options.s3_endpoint_override: + ds_proto.file_options.s3_endpoint_override = ( + body.file_options.s3_endpoint_override + ) elif body.bigquery_options: ds_proto.bigquery_options.table = body.bigquery_options.table ds_proto.bigquery_options.query = body.bigquery_options.query @@ -230,10 +282,14 @@ def apply_data_source(body: ApplyDataSourceRequestBody): ds_proto.snowflake_options.table = body.snowflake_options.table ds_proto.snowflake_options.database = body.snowflake_options.database ds_proto.snowflake_options.schema = body.snowflake_options.schema_ + if body.snowflake_options.query: + ds_proto.snowflake_options.query = body.snowflake_options.query elif body.redshift_options: ds_proto.redshift_options.table = body.redshift_options.table ds_proto.redshift_options.database = body.redshift_options.database ds_proto.redshift_options.schema = body.redshift_options.schema_ + if body.redshift_options.query: + ds_proto.redshift_options.query = body.redshift_options.query elif body.kafka_options: ds_proto.kafka_options.kafka_bootstrap_servers = ( body.kafka_options.kafka_bootstrap_servers @@ -242,6 +298,34 @@ def apply_data_source(body: ApplyDataSourceRequestBody): elif body.spark_options: ds_proto.spark_options.table = body.spark_options.table ds_proto.spark_options.path = body.spark_options.path + if body.spark_options.query: + ds_proto.spark_options.query = body.spark_options.query + if body.spark_options.file_format: + ds_proto.spark_options.file_format = body.spark_options.file_format + if body.spark_options.date_partition_column_format: + ds_proto.spark_options.date_partition_column_format = ( + body.spark_options.date_partition_column_format + ) + elif body.trino_options: + ds_proto.trino_options.table = body.trino_options.table + ds_proto.trino_options.query = body.trino_options.query + elif body.athena_options: + ds_proto.athena_options.table = body.athena_options.table + ds_proto.athena_options.query = body.athena_options.query + ds_proto.athena_options.database = body.athena_options.database + if body.athena_options.data_source: + ds_proto.athena_options.data_source = body.athena_options.data_source + elif body.kinesis_options: + ds_proto.kinesis_options.region = body.kinesis_options.region + ds_proto.kinesis_options.stream_name = body.kinesis_options.stream_name + elif body.custom_options: + if body.custom_options.configuration: + ds_proto.custom_options.configuration = ( + body.custom_options.configuration.encode("utf-8") + ) + + if body.data_source_class_type: + ds_proto.data_source_class_type = body.data_source_class_type req = RegistryServer_pb2.ApplyDataSourceRequest( data_source=ds_proto, diff --git a/sdk/python/feast/api/registry/rest/saved_datasets.py b/sdk/python/feast/api/registry/rest/saved_datasets.py index 4509c65f53d..9a03b6755c8 100644 --- a/sdk/python/feast/api/registry/rest/saved_datasets.py +++ b/sdk/python/feast/api/registry/rest/saved_datasets.py @@ -1,3 +1,4 @@ +import json import logging from typing import Dict, List, Optional @@ -32,6 +33,50 @@ logger = logging.getLogger(__name__) +def _build_feature_view_to_datasource_map( + grpc_handler, project: str, allow_cache: bool = True +) -> Dict[str, str]: + """Build a map of feature_view_name → data_source_name for a project.""" + try: + fv_req = RegistryServer_pb2.ListAllFeatureViewsRequest( + project=project, allow_cache=allow_cache + ) + fv_response = grpc_call(grpc_handler.ListAllFeatureViews, fv_req) + feature_views = fv_response.get("featureViews", []) + except Exception: + return {} + + fv_ds_map: Dict[str, str] = {} + for any_fv in feature_views: + # AnyFeatureView wraps under "featureView", "onDemandFeatureView", etc. + fv_data = any_fv.get("featureView") or any_fv.get("streamFeatureView") or {} + spec = fv_data.get("spec", {}) + fv_name = spec.get("name", "") + if not fv_name: + continue + batch_source = spec.get("batchSource", {}) + ds_name = batch_source.get("name", "") if batch_source else "" + if ds_name: + fv_ds_map[fv_name] = ds_name + return fv_ds_map + + +def _get_dataset_data_sources(dataset: dict, fv_ds_map: Dict[str, str]) -> List[str]: + """Get data source names for a dataset via its features' feature view lineage.""" + spec = dataset.get("spec", dataset) + features = spec.get("features", []) + data_sources: List[str] = [] + seen: set = set() + for feature_ref in features: + if isinstance(feature_ref, str) and ":" in feature_ref: + fv_name = feature_ref.split(":")[0] + ds_name = fv_ds_map.get(fv_name) + if ds_name and ds_name not in seen: + seen.add(ds_name) + data_sources.append(ds_name) + return data_sources + + class RegisterDatasetRequest(BaseModel): name: str project: str @@ -39,6 +84,7 @@ class RegisterDatasetRequest(BaseModel): join_keys: List[str] = [] storage_path: Optional[str] = None storage_type: str = "file" + storage_file_format: Optional[str] = "parquet" tags: Dict[str, str] = {} full_feature_names: bool = False feature_service_name: Optional[str] = None @@ -61,6 +107,7 @@ class CreateDatasetRequest(BaseModel): # Storage storage_type: str = "file" storage_path: str = "" + storage_file_format: Optional[str] = "parquet" tags: Dict[str, str] = {} allow_overwrite: bool = False @@ -79,7 +126,7 @@ def list_saved_datasets_all( False, description="Include relationships for each saved dataset" ), ): - return aggregate_across_projects( + result = aggregate_across_projects( grpc_handler=grpc_handler, list_method=grpc_handler.ListSavedDatasets, request_cls=RegistryServer_pb2.ListSavedDatasetsRequest, @@ -93,6 +140,21 @@ def list_saved_datasets_all( include_relationships=include_relationships, ) + saved_datasets = result.get("savedDatasets", []) + projects_seen: Dict[str, Dict[str, str]] = {} + for sd in saved_datasets: + proj = sd.get("project", "") + if proj not in projects_seen: + projects_seen[proj] = _build_feature_view_to_datasource_map( + grpc_handler, proj, allow_cache + ) + ds_names = _get_dataset_data_sources(sd, projects_seen[proj]) + if ds_names: + spec = sd.get("spec", sd) + spec["dataSources"] = ds_names + + return result + @router.get("/saved_datasets/data/{name}") def get_dataset_data( name: str, @@ -193,6 +255,14 @@ def get_saved_dataset( result = saved_dataset + fv_ds_map = _build_feature_view_to_datasource_map( + grpc_handler, project, allow_cache + ) + ds_names = _get_dataset_data_sources(result, fv_ds_map) + if ds_names: + spec = result.get("spec", result) + spec["dataSources"] = ds_names + if include_relationships: relationships = get_object_relationships( grpc_handler, "savedDataset", name, project, allow_cache @@ -246,6 +316,15 @@ def list_saved_datasets( response = grpc_call(grpc_handler.ListSavedDatasets, req) saved_datasets = response.get("savedDatasets", []) + fv_ds_map = _build_feature_view_to_datasource_map( + grpc_handler, project, allow_cache + ) + for sd in saved_datasets: + ds_names = _get_dataset_data_sources(sd, fv_ds_map) + if ds_names: + spec = sd.get("spec", sd) + spec["dataSources"] = ds_names + result = { "savedDatasets": saved_datasets, "pagination": response.get("pagination", {}), @@ -308,12 +387,48 @@ def register_saved_dataset(payload: RegisterDatasetRequest = Body(...)): ) elif payload.storage_type == "spark": storage_proto.spark_storage.CopyFrom( - DataSourceProto.SparkOptions(path=path) + DataSourceProto.SparkOptions( + path=path, + file_format=payload.storage_file_format or "parquet", + ) ) elif payload.storage_type == "athena": storage_proto.athena_storage.CopyFrom( DataSourceProto.AthenaOptions(table=path) ) + elif payload.storage_type == "postgres": + config = json.dumps({"name": None, "query": None, "table": path}) + storage_proto.custom_storage.CopyFrom( + DataSourceProto.CustomSourceOptions( + configuration=config.encode("utf-8") + ) + ) + elif payload.storage_type == "clickhouse": + config = json.dumps({"name": None, "query": None, "table": path}) + storage_proto.custom_storage.CopyFrom( + DataSourceProto.CustomSourceOptions( + configuration=config.encode("utf-8") + ) + ) + elif payload.storage_type == "couchbase": + parts = path.split(".", 2) + database = parts[0] if len(parts) > 0 else "" + scope = parts[1] if len(parts) > 1 else "" + collection = parts[2] if len(parts) > 2 else "" + config = json.dumps( + { + "name": None, + "query": None, + "database": database, + "scope": scope, + "collection": collection, + } + ) + storage_proto.custom_storage.CopyFrom( + DataSourceProto.CustomSourceOptions( + configuration=config.encode("utf-8") + ) + ) elif payload.storage_type == "custom": storage_proto.custom_storage.CopyFrom( DataSourceProto.CustomSourceOptions(configuration=path.encode("utf-8")) diff --git a/sdk/python/feast/dataset_utils.py b/sdk/python/feast/dataset_utils.py index 3a5b617caf4..cea6c9330a2 100644 --- a/sdk/python/feast/dataset_utils.py +++ b/sdk/python/feast/dataset_utils.py @@ -22,10 +22,12 @@ def coerce_value(val: str) -> Union[int, float, str]: return val -def build_saved_dataset_storage(storage_type: str, path: str): +def build_saved_dataset_storage( + storage_type: str, path: str, file_format: Optional[str] = None +): """Build a SavedDatasetStorage object from type string and path/table reference. - Supports: file (default), bigquery, snowflake, redshift. + Supports: file (default), bigquery, snowflake, redshift, spark, trino, athena. Unknown types fall back to file storage. """ from feast.infra.offline_stores.file_source import SavedDatasetFileStorage @@ -47,6 +49,36 @@ def build_saved_dataset_storage(storage_type: str, path: str): "SavedDatasetRedshiftStorage", "redshift", ), + ( + "feast.infra.offline_stores.contrib.spark_offline_store.spark_source", + "SavedDatasetSparkStorage", + "spark", + ), + ( + "feast.infra.offline_stores.contrib.trino_offline_store.trino_source", + "SavedDatasetTrinoStorage", + "trino", + ), + ( + "feast.infra.offline_stores.contrib.athena_offline_store.athena_source", + "SavedDatasetAthenaStorage", + "athena", + ), + ( + "feast.infra.offline_stores.contrib.postgres_offline_store.postgres_source", + "SavedDatasetPostgreSQLStorage", + "postgres", + ), + ( + "feast.infra.offline_stores.contrib.clickhouse_offline_store.clickhouse_source", + "SavedDatasetClickhouseStorage", + "clickhouse", + ), + ( + "feast.infra.offline_stores.contrib.couchbase_offline_store.couchbase_source", + "SavedDatasetCouchbaseColumnarStorage", + "couchbase", + ), ]: try: m = importlib.import_module(mod) @@ -60,6 +92,26 @@ def build_saved_dataset_storage(storage_type: str, path: str): return storage_classes["snowflake"](table_ref=path) elif storage_type == "redshift" and "redshift" in storage_classes: return storage_classes["redshift"](table_ref=path) + elif storage_type == "spark" and "spark" in storage_classes: + return storage_classes["spark"](path=path, file_format=file_format or "parquet") + elif storage_type == "trino" and "trino" in storage_classes: + return storage_classes["trino"](table=path) + elif storage_type == "athena" and "athena" in storage_classes: + return storage_classes["athena"](table_ref=path) + elif storage_type == "postgres" and "postgres" in storage_classes: + return storage_classes["postgres"](table_ref=path) + elif storage_type == "clickhouse" and "clickhouse" in storage_classes: + return storage_classes["clickhouse"](table_ref=path) + elif storage_type == "couchbase" and "couchbase" in storage_classes: + parts = path.split(".", 2) + database_ref = parts[0] if len(parts) > 0 else "" + scope_ref = parts[1] if len(parts) > 1 else "" + collection_ref = parts[2] if len(parts) > 2 else "" + return storage_classes["couchbase"]( + database_ref=database_ref, + scope_ref=scope_ref, + collection_ref=collection_ref, + ) else: return SavedDatasetFileStorage(path=path) diff --git a/ui/src/components/DataSourceFormModal.tsx b/ui/src/components/DataSourceFormModal.tsx index ca95bb6aa0e..e3c838a349b 100644 --- a/ui/src/components/DataSourceFormModal.tsx +++ b/ui/src/components/DataSourceFormModal.tsx @@ -47,6 +47,10 @@ const SOURCE_TYPE_OPTIONS = [ value: String(feast.core.DataSource.SourceType.BATCH_ATHENA), text: "AWS Athena", }, + { + value: String(feast.core.DataSource.SourceType.BATCH_ICEBERG), + text: "Iceberg / Unity Catalog", + }, { value: String(feast.core.DataSource.SourceType.STREAM_KAFKA), text: "Kafka", @@ -84,42 +88,87 @@ interface DataSourceFormData { timestampField: string; createdTimestampColumn: string; tags: TagEntry[]; + // File source fileUri: string; + fileFormat: string; + fileS3EndpointOverride: string; + // BigQuery bigqueryTable: string; bigqueryQuery: string; + bigqueryDatePartitionColumn: string; + // Snowflake snowflakeTable: string; snowflakeDatabase: string; snowflakeSchema: string; + snowflakeQuery: string; + snowflakeWarehouse: string; + // Redshift redshiftTable: string; redshiftDatabase: string; redshiftSchema: string; + redshiftQuery: string; + // Kafka kafkaBootstrapServers: string; kafkaTopic: string; + kafkaMessageFormat: string; + kafkaWatermarkDelay: string; + // Spark sparkTable: string; sparkPath: string; + sparkQuery: string; + sparkFileFormat: string; + sparkTableFormat: string; + sparkTableFormatCatalog: string; + sparkTableFormatNamespace: string; + sparkTableFormatProperties: string; + sparkDatePartitionColumn: string; + sparkDatePartitionFormat: string; + // Kinesis kinesisRegion: string; kinesisStreamName: string; + kinesisRecordFormat: string; + // Trino trinoTable: string; trinoQuery: string; + // Athena athenaTable: string; athenaQuery: string; athenaDatabase: string; athenaDataSource: string; + athenaDatePartitionColumn: string; + // Custom customSourceClassName: string; customSourceConfig: string; - // Contrib source fields + // Iceberg / Unity Catalog + icebergCatalogType: string; + icebergEndpoint: string; + icebergWarehouse: string; + icebergNamespace: string; + icebergTable: string; + icebergTokenEnvVar: string; + icebergCredentialVending: string; + icebergCatalogProperties: string; + // Ray rayReaderType: string; rayPath: string; rayReaderOptions: string; + // Postgres postgresTable: string; postgresQuery: string; + // MongoDB mongodbCollection: string; + // ClickHouse clickhouseTable: string; clickhouseQuery: string; + // MSSQL mssqlTable: string; mssqlConnectionStr: string; + mssqlDatePartitionColumn: string; + // Oracle oracleTable: string; oracleConnectionStr: string; + oracleDatePartitionColumn: string; + // Couchbase couchbaseDatabase: string; couchbaseScope: string; couchbaseCollection: string; @@ -144,28 +193,54 @@ const EMPTY_FORM: DataSourceFormData = { createdTimestampColumn: "", tags: [], fileUri: "", + fileFormat: "parquet", + fileS3EndpointOverride: "", bigqueryTable: "", bigqueryQuery: "", + bigqueryDatePartitionColumn: "", snowflakeTable: "", snowflakeDatabase: "", snowflakeSchema: "", + snowflakeQuery: "", + snowflakeWarehouse: "", redshiftTable: "", redshiftDatabase: "", redshiftSchema: "", + redshiftQuery: "", kafkaBootstrapServers: "", kafkaTopic: "", + kafkaMessageFormat: "json", + kafkaWatermarkDelay: "", sparkTable: "", sparkPath: "", + sparkQuery: "", + sparkFileFormat: "parquet", + sparkTableFormat: "", + sparkTableFormatCatalog: "", + sparkTableFormatNamespace: "", + sparkTableFormatProperties: "", + sparkDatePartitionColumn: "", + sparkDatePartitionFormat: "%Y-%m-%d", kinesisRegion: "", kinesisStreamName: "", + kinesisRecordFormat: "json", trinoTable: "", trinoQuery: "", athenaTable: "", athenaQuery: "", athenaDatabase: "", athenaDataSource: "", + athenaDatePartitionColumn: "", customSourceClassName: "", customSourceConfig: "", + icebergCatalogType: "rest", + icebergEndpoint: "", + icebergWarehouse: "", + icebergNamespace: "", + icebergTable: "", + icebergTokenEnvVar: "", + icebergCredentialVending: "true", + icebergCatalogProperties: "", rayReaderType: "parquet", rayPath: "", rayReaderOptions: "", @@ -176,8 +251,10 @@ const EMPTY_FORM: DataSourceFormData = { clickhouseQuery: "", mssqlTable: "", mssqlConnectionStr: "", + mssqlDatePartitionColumn: "", oracleTable: "", oracleConnectionStr: "", + oracleDatePartitionColumn: "", couchbaseDatabase: "", couchbaseScope: "", couchbaseCollection: "", @@ -192,6 +269,7 @@ const BATCH_SOURCE_TYPES = new Set([ String(feast.core.DataSource.SourceType.BATCH_SPARK), String(feast.core.DataSource.SourceType.BATCH_TRINO), String(feast.core.DataSource.SourceType.BATCH_ATHENA), + String(feast.core.DataSource.SourceType.BATCH_ICEBERG), "RAY_SOURCE", "POSTGRES_SOURCE", "MONGODB_SOURCE", @@ -264,23 +342,36 @@ const DataSourceFormModal: React.FC = ({ } else if ( st === String(feast.core.DataSource.SourceType.BATCH_SNOWFLAKE) ) { - if (!formData.snowflakeTable.trim()) { - newErrors.snowflakeTable = "Table name is required for Snowflake."; + if (!formData.snowflakeTable.trim() && !formData.snowflakeQuery.trim()) { + newErrors.snowflakeTable = + "Either a table or a query is required for Snowflake."; } if (!formData.snowflakeDatabase.trim()) { newErrors.snowflakeDatabase = "Database is required for Snowflake."; } } else if (st === String(feast.core.DataSource.SourceType.BATCH_REDSHIFT)) { - if (!formData.redshiftTable.trim()) { - newErrors.redshiftTable = "Table name is required for Redshift."; + if (!formData.redshiftTable.trim() && !formData.redshiftQuery.trim()) { + newErrors.redshiftTable = + "Either a table or a query is required for Redshift."; } if (!formData.redshiftDatabase.trim()) { newErrors.redshiftDatabase = "Database is required for Redshift."; } } else if (st === String(feast.core.DataSource.SourceType.BATCH_SPARK)) { - if (!formData.sparkTable.trim() && !formData.sparkPath.trim()) { + if ( + !formData.sparkTable.trim() && + !formData.sparkPath.trim() && + !formData.sparkQuery.trim() + ) { newErrors.sparkTable = - "Either a table reference or a path is required for Spark."; + "Either a table, path, or query is required for Spark."; + } else if ( + formData.sparkPath.trim() && + !formData.sparkTableFormat && + !formData.sparkFileFormat + ) { + newErrors.sparkFileFormat = + "File format is required when using a path without table format."; } } else if (st === String(feast.core.DataSource.SourceType.BATCH_TRINO)) { if (!formData.trinoTable.trim() && !formData.trinoQuery.trim()) { @@ -295,6 +386,16 @@ const DataSourceFormModal: React.FC = ({ if (!formData.athenaDatabase.trim()) { newErrors.athenaDatabase = "Database is required for Athena."; } + } else if (st === String(feast.core.DataSource.SourceType.BATCH_ICEBERG)) { + if (!formData.icebergWarehouse.trim()) { + newErrors.icebergWarehouse = "Warehouse is required."; + } + if (!formData.icebergNamespace.trim()) { + newErrors.icebergNamespace = "Namespace is required."; + } + if (!formData.icebergTable.trim()) { + newErrors.icebergTable = "Table name is required."; + } } else if (st === String(feast.core.DataSource.SourceType.STREAM_KAFKA)) { if (!formData.kafkaBootstrapServers.trim()) { newErrors.kafkaBootstrapServers = "Bootstrap servers are required."; @@ -402,19 +503,52 @@ const DataSourceFormModal: React.FC = ({ }; const renderFileSourceFields = () => ( - - updateField("fileUri", e.target.value)} + <> + - + error={errors.fileUri} + helpText="Path to the data file accessible by the Feast server (e.g. s3://bucket/path/data.parquet, gs://bucket/data.csv, file:///mnt/data/features.parquet)." + > + updateField("fileUri", e.target.value)} + isInvalid={!!errors.fileUri} + placeholder="s3://bucket/path/to/data.parquet" + /> + + + + + updateField("fileFormat", e.target.value)} + /> + + + + + + updateField("fileS3EndpointOverride", e.target.value) + } + placeholder="http://minio:9000" + /> + + + + ); const renderSourceTypeHeader = () => { @@ -489,6 +623,18 @@ const DataSourceFormModal: React.FC = ({ rows={3} /> + + + updateField("bigqueryDatePartitionColumn", e.target.value) + } + placeholder="date_partition" + /> + ); } @@ -525,16 +671,48 @@ const DataSourceFormModal: React.FC = ({ + + + + + updateField("snowflakeTable", e.target.value) + } + isInvalid={!!errors.snowflakeTable} + placeholder="MY_TABLE" + /> + + + + + + updateField("snowflakeWarehouse", e.target.value) + } + placeholder="COMPUTE_WH" + /> + + + - updateField("snowflakeTable", e.target.value)} - isInvalid={!!errors.snowflakeTable} - placeholder="MY_TABLE" + updateField("snowflakeQuery", e.target.value)} + placeholder="SELECT * FROM MY_TABLE WHERE ..." + rows={3} /> @@ -577,6 +755,7 @@ const DataSourceFormModal: React.FC = ({ label="Table" isInvalid={!!errors.redshiftTable} error={errors.redshiftTable} + helpText="Provide either a table or a query." > = ({ placeholder="my_table" /> + + updateField("redshiftQuery", e.target.value)} + placeholder="SELECT * FROM my_table WHERE ..." + rows={3} + /> + ); } @@ -607,16 +797,50 @@ const DataSourceFormModal: React.FC = ({ placeholder="broker1:9092,broker2:9092" /> + + + + updateField("kafkaTopic", e.target.value)} + isInvalid={!!errors.kafkaTopic} + placeholder="my-feature-topic" + /> + + + + + + updateField("kafkaMessageFormat", e.target.value) + } + /> + + + updateField("kafkaTopic", e.target.value)} - isInvalid={!!errors.kafkaTopic} - placeholder="my-feature-topic" + value={formData.kafkaWatermarkDelay} + onChange={(e) => + updateField("kafkaWatermarkDelay", e.target.value) + } + placeholder="30 seconds" /> @@ -630,7 +854,7 @@ const DataSourceFormModal: React.FC = ({ label="Table" isInvalid={!!errors.sparkTable} error={errors.sparkTable} - helpText="Spark catalog table (catalog.database.table). Provide either table or path." + helpText="Spark catalog table (catalog.database.table). Provide table, path, or query." > = ({ updateField("sparkPath", e.target.value)} - placeholder="s3://bucket/path/" + placeholder="s3://bucket/path/ or abfss://container@account/path/" /> + + updateField("sparkQuery", e.target.value)} + placeholder="SELECT * FROM catalog.db.table WHERE ..." + rows={3} + /> + + + + + + updateField("sparkFileFormat", e.target.value) + } + /> + + + + + + updateField("sparkTableFormat", e.target.value) + } + /> + + + + {formData.sparkTableFormat && ( + <> + + + + + updateField("sparkTableFormatCatalog", e.target.value) + } + placeholder="my_catalog" + /> + + + + + + updateField("sparkTableFormatNamespace", e.target.value) + } + placeholder="my_db" + /> + + + + + + updateField("sparkTableFormatProperties", e.target.value) + } + placeholder='{"warehouse": "s3://bucket/warehouse"}' + rows={2} + /> + + + )} + + + + + updateField("sparkDatePartitionColumn", e.target.value) + } + placeholder="date_partition" + /> + + + + + + updateField("sparkDatePartitionFormat", e.target.value) + } + placeholder="%Y-%m-%d" + /> + + + ); } @@ -734,25 +1087,168 @@ const DataSourceFormModal: React.FC = ({ rows={3} /> + + + updateField("athenaDatePartitionColumn", e.target.value) + } + placeholder="date_partition" + /> + ); } - if (st === String(feast.core.DataSource.SourceType.STREAM_KINESIS)) { + if (st === String(feast.core.DataSource.SourceType.BATCH_ICEBERG)) { return ( <> + + updateField("icebergCatalogType", e.target.value) + } + /> + + + updateField("icebergEndpoint", e.target.value)} + placeholder="http://localhost:8080/api/2.1/unity-catalog/iceberg" + /> + + + + + + updateField("icebergWarehouse", e.target.value) + } + isInvalid={!!errors.icebergWarehouse} + placeholder="unity" + /> + + + + + + updateField("icebergNamespace", e.target.value) + } + isInvalid={!!errors.icebergNamespace} + placeholder="default" + /> + + + + updateField("kinesisRegion", e.target.value)} - isInvalid={!!errors.kinesisRegion} - placeholder="us-east-1" + value={formData.icebergTable} + onChange={(e) => updateField("icebergTable", e.target.value)} + isInvalid={!!errors.icebergTable} + placeholder="driver_stats" /> + + + updateField("icebergTokenEnvVar", e.target.value) + } + placeholder="DATABRICKS_TOKEN" + /> + + + + updateField("icebergCatalogProperties", e.target.value) + } + placeholder='{"s3.region": "us-east-1"}' + rows={3} + /> + + + ); + } + + if (st === String(feast.core.DataSource.SourceType.STREAM_KINESIS)) { + return ( + <> + + + + updateField("kinesisRegion", e.target.value)} + isInvalid={!!errors.kinesisRegion} + placeholder="us-east-1" + /> + + + + + + updateField("kinesisRecordFormat", e.target.value) + } + /> + + + = ({ placeholder="mssql+pyodbc://user:pass@host/db" // pragma: allowlist secret /> + + + updateField("mssqlDatePartitionColumn", e.target.value) + } + placeholder="date_partition" + /> + ); } @@ -974,6 +1482,18 @@ const DataSourceFormModal: React.FC = ({ placeholder="oracle+cx_oracle://user:pass@host:1521/service" // pragma: allowlist secret /> + + + updateField("oracleDatePartitionColumn", e.target.value) + } + placeholder="date_partition" + /> + ); } diff --git a/ui/src/components/FeatureViewFormModal.tsx b/ui/src/components/FeatureViewFormModal.tsx index 9e0655d96cf..b39f09e192e 100644 --- a/ui/src/components/FeatureViewFormModal.tsx +++ b/ui/src/components/FeatureViewFormModal.tsx @@ -245,12 +245,19 @@ const FeatureViewFormModal: React.FC = ({ const st = dsData.sourceType; if (st === String(feast.core.DataSource.SourceType.BATCH_FILE)) { - payload.file_options = { uri: dsData.fileUri }; + payload.file_options = { + uri: dsData.fileUri, + file_format: dsData.fileFormat || "parquet", + s3_endpoint_override: dsData.fileS3EndpointOverride || "", + }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_BIGQUERY)) { payload.bigquery_options = { table: dsData.bigqueryTable, query: dsData.bigqueryQuery, }; + if (dsData.bigqueryDatePartitionColumn) { + payload.date_partition_column = dsData.bigqueryDatePartitionColumn; + } } else if ( st === String(feast.core.DataSource.SourceType.BATCH_SNOWFLAKE) ) { @@ -258,22 +265,35 @@ const FeatureViewFormModal: React.FC = ({ table: dsData.snowflakeTable, database: dsData.snowflakeDatabase, schema_: dsData.snowflakeSchema, + query: dsData.snowflakeQuery || "", + warehouse: dsData.snowflakeWarehouse || "", }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_REDSHIFT)) { payload.redshift_options = { table: dsData.redshiftTable, database: dsData.redshiftDatabase, schema_: dsData.redshiftSchema, + query: dsData.redshiftQuery || "", }; } else if (st === String(feast.core.DataSource.SourceType.STREAM_KAFKA)) { payload.kafka_options = { kafka_bootstrap_servers: dsData.kafkaBootstrapServers, topic: dsData.kafkaTopic, + message_format: dsData.kafkaMessageFormat || "json", + watermark_delay_threshold: dsData.kafkaWatermarkDelay || "", }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_SPARK)) { payload.spark_options = { table: dsData.sparkTable, path: dsData.sparkPath, + query: dsData.sparkQuery || "", + file_format: dsData.sparkFileFormat || "", + table_format: dsData.sparkTableFormat || "", + table_format_catalog: dsData.sparkTableFormatCatalog || "", + table_format_namespace: dsData.sparkTableFormatNamespace || "", + table_format_properties: dsData.sparkTableFormatProperties || "", + date_partition_column: dsData.sparkDatePartitionColumn || "", + date_partition_column_format: dsData.sparkDatePartitionFormat || "", }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_TRINO)) { payload.trino_options = { @@ -287,10 +307,32 @@ const FeatureViewFormModal: React.FC = ({ database: dsData.athenaDatabase, data_source: dsData.athenaDataSource, }; + if (dsData.athenaDatePartitionColumn) { + payload.date_partition_column = dsData.athenaDatePartitionColumn; + } + } else if (st === String(feast.core.DataSource.SourceType.BATCH_ICEBERG)) { + const catalogProps = dsData.icebergCatalogProperties?.trim() + ? JSON.parse(dsData.icebergCatalogProperties) + : {}; + payload.custom_options = { + configuration: JSON.stringify({ + catalog_type: dsData.icebergCatalogType || "rest", + endpoint: dsData.icebergEndpoint, + warehouse: dsData.icebergWarehouse, + namespace: dsData.icebergNamespace, + table: dsData.icebergTable, + token_env_var: dsData.icebergTokenEnvVar || null, + credential_vending: dsData.icebergCredentialVending !== "false", + catalog_properties: catalogProps, + }), + }; + payload.data_source_class_type = + "feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source.IcebergSource"; } else if (st === String(feast.core.DataSource.SourceType.STREAM_KINESIS)) { payload.kinesis_options = { region: dsData.kinesisRegion, stream_name: dsData.kinesisStreamName, + record_format: dsData.kinesisRecordFormat || "json", }; } else if (st === String(feast.core.DataSource.SourceType.CUSTOM_SOURCE)) { payload.custom_options = { diff --git a/ui/src/pages/data-sources/DataSourceOverviewTab.tsx b/ui/src/pages/data-sources/DataSourceOverviewTab.tsx index 42a8d2536f0..2eb8d0644a6 100644 --- a/ui/src/pages/data-sources/DataSourceOverviewTab.tsx +++ b/ui/src/pages/data-sources/DataSourceOverviewTab.tsx @@ -45,58 +45,176 @@ const buildEditFormData = (ds: any): DataSourceFormData => { createdTimestampColumn: spec.createdTimestampColumn || ds.createdTimestampColumn || "", tags, + // File fileUri: spec.fileOptions?.uri || ds.fileOptions?.uri || "", + fileFormat: + spec.fileOptions?.fileFormat || ds.fileOptions?.fileFormat || "parquet", + fileS3EndpointOverride: + spec.fileOptions?.s3EndpointOverride || + ds.fileOptions?.s3EndpointOverride || + "", + // BigQuery bigqueryTable: spec.bigqueryOptions?.table || ds.bigqueryOptions?.table || "", bigqueryQuery: spec.bigqueryOptions?.query || ds.bigqueryOptions?.query || "", + bigqueryDatePartitionColumn: + spec.datePartitionColumn || ds.datePartitionColumn || "", + // Snowflake snowflakeTable: spec.snowflakeOptions?.table || ds.snowflakeOptions?.table || "", snowflakeDatabase: spec.snowflakeOptions?.database || ds.snowflakeOptions?.database || "", snowflakeSchema: spec.snowflakeOptions?.schema || ds.snowflakeOptions?.schema || "", + snowflakeQuery: + spec.snowflakeOptions?.query || ds.snowflakeOptions?.query || "", + snowflakeWarehouse: + spec.snowflakeOptions?.warehouse || ds.snowflakeOptions?.warehouse || "", + // Redshift redshiftTable: spec.redshiftOptions?.table || ds.redshiftOptions?.table || "", redshiftDatabase: spec.redshiftOptions?.database || ds.redshiftOptions?.database || "", redshiftSchema: spec.redshiftOptions?.schema || ds.redshiftOptions?.schema || "", + redshiftQuery: + spec.redshiftOptions?.query || ds.redshiftOptions?.query || "", + // Kafka kafkaBootstrapServers: spec.kafkaOptions?.kafkaBootstrapServers || ds.kafkaOptions?.kafkaBootstrapServers || "", kafkaTopic: spec.kafkaOptions?.topic || ds.kafkaOptions?.topic || "", + kafkaMessageFormat: + spec.kafkaOptions?.messageFormat || + ds.kafkaOptions?.messageFormat || + "json", + kafkaWatermarkDelay: + spec.kafkaOptions?.watermarkDelayThreshold || + ds.kafkaOptions?.watermarkDelayThreshold || + "", + // Spark sparkTable: spec.sparkOptions?.table || ds.sparkOptions?.table || "", sparkPath: spec.sparkOptions?.path || ds.sparkOptions?.path || "", + sparkQuery: spec.sparkOptions?.query || ds.sparkOptions?.query || "", + sparkFileFormat: + spec.sparkOptions?.fileFormat || ds.sparkOptions?.fileFormat || "", + sparkTableFormat: + spec.sparkOptions?.tableFormat?.formatType || + ds.sparkOptions?.tableFormat?.formatType || + "", + sparkTableFormatCatalog: + spec.sparkOptions?.tableFormat?.catalog || + ds.sparkOptions?.tableFormat?.catalog || + "", + sparkTableFormatNamespace: + spec.sparkOptions?.tableFormat?.namespace || + ds.sparkOptions?.tableFormat?.namespace || + "", + sparkTableFormatProperties: (() => { + const props = + spec.sparkOptions?.tableFormat?.properties || + ds.sparkOptions?.tableFormat?.properties; + return props ? JSON.stringify(props) : ""; + })(), + sparkDatePartitionColumn: + spec.sparkOptions?.datePartitionColumn || + ds.sparkOptions?.datePartitionColumn || + spec.datePartitionColumn || + ds.datePartitionColumn || + "", + sparkDatePartitionFormat: + spec.sparkOptions?.datePartitionColumnFormat || + ds.sparkOptions?.datePartitionColumnFormat || + "%Y-%m-%d", + // Kinesis kinesisRegion: spec.kinesisOptions?.region || ds.kinesisOptions?.region || "", kinesisStreamName: spec.kinesisOptions?.streamName || ds.kinesisOptions?.streamName || "", + kinesisRecordFormat: + spec.kinesisOptions?.recordFormat || + ds.kinesisOptions?.recordFormat || + "json", + // Trino trinoTable: spec.trinoOptions?.table || ds.trinoOptions?.table || "", trinoQuery: spec.trinoOptions?.query || ds.trinoOptions?.query || "", + // Athena athenaTable: spec.athenaOptions?.table || ds.athenaOptions?.table || "", athenaQuery: spec.athenaOptions?.query || ds.athenaOptions?.query || "", athenaDatabase: spec.athenaOptions?.database || ds.athenaOptions?.database || "", athenaDataSource: spec.athenaOptions?.dataSource || ds.athenaOptions?.dataSource || "", + athenaDatePartitionColumn: + spec.datePartitionColumn || ds.datePartitionColumn || "", + // Custom customSourceClassName: spec.customOptions?.className || ds.customOptions?.className || "", customSourceConfig: spec.customOptions?.config || ds.customOptions?.config || "", + // Iceberg + ...(() => { + const configStr = + spec.customOptions?.configuration || + ds.customOptions?.configuration || + ""; + if ( + String(spec.type ?? ds.type ?? 0) === + String(feast.core.DataSource.SourceType.BATCH_ICEBERG) && + configStr + ) { + try { + const cfg = JSON.parse(configStr); + return { + icebergCatalogType: cfg.catalog_type || "rest", + icebergEndpoint: cfg.endpoint || "", + icebergWarehouse: cfg.warehouse || "", + icebergNamespace: cfg.namespace || "", + icebergTable: cfg.table || "", + icebergTokenEnvVar: cfg.token_env_var || "", + icebergCredentialVending: String(cfg.credential_vending ?? true), + icebergCatalogProperties: cfg.catalog_properties + ? JSON.stringify(cfg.catalog_properties) + : "", + }; + } catch { + /* ignore parse errors */ + } + } + return { + icebergCatalogType: "rest", + icebergEndpoint: "", + icebergWarehouse: "", + icebergNamespace: "", + icebergTable: "", + icebergTokenEnvVar: "", + icebergCredentialVending: "true", + icebergCatalogProperties: "", + }; + })(), + // Ray rayReaderType: "", rayPath: "", rayReaderOptions: "", + // Postgres postgresTable: "", postgresQuery: "", + // MongoDB mongodbCollection: "", + // ClickHouse clickhouseTable: "", clickhouseQuery: "", + // MSSQL mssqlTable: "", mssqlConnectionStr: "", + mssqlDatePartitionColumn: "", + // Oracle oracleTable: "", oracleConnectionStr: "", + oracleDatePartitionColumn: "", + // Couchbase couchbaseDatabase: "", couchbaseScope: "", couchbaseCollection: "", @@ -120,33 +238,53 @@ const formDataToPayload = (formData: DataSourceFormData, project: string) => { const st = formData.sourceType; if (st === String(feast.core.DataSource.SourceType.BATCH_FILE)) { - payload.file_options = { uri: formData.fileUri }; + payload.file_options = { + uri: formData.fileUri, + file_format: formData.fileFormat || "parquet", + s3_endpoint_override: formData.fileS3EndpointOverride || "", + }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_BIGQUERY)) { payload.bigquery_options = { table: formData.bigqueryTable, query: formData.bigqueryQuery, }; + if (formData.bigqueryDatePartitionColumn) { + payload.date_partition_column = formData.bigqueryDatePartitionColumn; + } } else if (st === String(feast.core.DataSource.SourceType.BATCH_SNOWFLAKE)) { payload.snowflake_options = { table: formData.snowflakeTable, database: formData.snowflakeDatabase, schema_: formData.snowflakeSchema, + query: formData.snowflakeQuery || "", + warehouse: formData.snowflakeWarehouse || "", }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_REDSHIFT)) { payload.redshift_options = { table: formData.redshiftTable, database: formData.redshiftDatabase, schema_: formData.redshiftSchema, + query: formData.redshiftQuery || "", }; } else if (st === String(feast.core.DataSource.SourceType.STREAM_KAFKA)) { payload.kafka_options = { kafka_bootstrap_servers: formData.kafkaBootstrapServers, topic: formData.kafkaTopic, + message_format: formData.kafkaMessageFormat || "json", + watermark_delay_threshold: formData.kafkaWatermarkDelay || "", }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_SPARK)) { payload.spark_options = { table: formData.sparkTable, path: formData.sparkPath, + query: formData.sparkQuery || "", + file_format: formData.sparkFileFormat || "", + table_format: formData.sparkTableFormat || "", + table_format_catalog: formData.sparkTableFormatCatalog || "", + table_format_namespace: formData.sparkTableFormatNamespace || "", + table_format_properties: formData.sparkTableFormatProperties || "", + date_partition_column: formData.sparkDatePartitionColumn || "", + date_partition_column_format: formData.sparkDatePartitionFormat || "", }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_TRINO)) { payload.trino_options = { @@ -160,10 +298,32 @@ const formDataToPayload = (formData: DataSourceFormData, project: string) => { database: formData.athenaDatabase, data_source: formData.athenaDataSource, }; + if (formData.athenaDatePartitionColumn) { + payload.date_partition_column = formData.athenaDatePartitionColumn; + } + } else if (st === String(feast.core.DataSource.SourceType.BATCH_ICEBERG)) { + const catalogProps = formData.icebergCatalogProperties.trim() + ? JSON.parse(formData.icebergCatalogProperties) + : {}; + payload.custom_options = { + configuration: JSON.stringify({ + catalog_type: formData.icebergCatalogType || "rest", + endpoint: formData.icebergEndpoint, + warehouse: formData.icebergWarehouse, + namespace: formData.icebergNamespace, + table: formData.icebergTable, + token_env_var: formData.icebergTokenEnvVar || null, + credential_vending: formData.icebergCredentialVending !== "false", + catalog_properties: catalogProps, + }), + }; + payload.data_source_class_type = + "feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source.IcebergSource"; } else if (st === String(feast.core.DataSource.SourceType.STREAM_KINESIS)) { payload.kinesis_options = { region: formData.kinesisRegion, stream_name: formData.kinesisStreamName, + record_format: formData.kinesisRecordFormat || "json", }; } else if (st === String(feast.core.DataSource.SourceType.CUSTOM_SOURCE)) { payload.custom_options = { @@ -275,6 +435,72 @@ const DataSourceOverviewTab = () => { {spec?.fileOptions || spec?.bigqueryOptions ? ( + ) : String(sourceType) === + String(feast.core.DataSource.SourceType.BATCH_ICEBERG) ? ( + (() => { + let cfg: any = {}; + try { + cfg = JSON.parse( + spec?.customOptions?.configuration || "{}", + ); + } catch { + /* ignore */ + } + return ( + + + Source Type + + + Iceberg / Unity Catalog + + + Catalog Type + + + {cfg.catalog_type || "rest"} + + {cfg.endpoint && ( + <> + + Endpoint + + + {cfg.endpoint} + + + )} + + Warehouse + + + {cfg.warehouse || "—"} + + + Namespace + + + {cfg.namespace || "—"} + + + Table + + + {cfg.table || "—"} + + {cfg.token_env_var && ( + <> + + Token Env Variable + + + {cfg.token_env_var} + + + )} + + ); + })() ) : sourceType ? ( diff --git a/ui/src/pages/data-sources/Index.tsx b/ui/src/pages/data-sources/Index.tsx index e041058f995..c8422d3943a 100644 --- a/ui/src/pages/data-sources/Index.tsx +++ b/ui/src/pages/data-sources/Index.tsx @@ -67,33 +67,53 @@ const formDataToPayload = (formData: DataSourceFormData, project: string) => { const st = formData.sourceType; if (st === String(feast.core.DataSource.SourceType.BATCH_FILE)) { - payload.file_options = { uri: formData.fileUri }; + payload.file_options = { + uri: formData.fileUri, + file_format: formData.fileFormat || "parquet", + s3_endpoint_override: formData.fileS3EndpointOverride || "", + }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_BIGQUERY)) { payload.bigquery_options = { table: formData.bigqueryTable, query: formData.bigqueryQuery, }; + if (formData.bigqueryDatePartitionColumn) { + payload.date_partition_column = formData.bigqueryDatePartitionColumn; + } } else if (st === String(feast.core.DataSource.SourceType.BATCH_SNOWFLAKE)) { payload.snowflake_options = { table: formData.snowflakeTable, database: formData.snowflakeDatabase, schema_: formData.snowflakeSchema, + query: formData.snowflakeQuery || "", + warehouse: formData.snowflakeWarehouse || "", }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_REDSHIFT)) { payload.redshift_options = { table: formData.redshiftTable, database: formData.redshiftDatabase, schema_: formData.redshiftSchema, + query: formData.redshiftQuery || "", }; } else if (st === String(feast.core.DataSource.SourceType.STREAM_KAFKA)) { payload.kafka_options = { kafka_bootstrap_servers: formData.kafkaBootstrapServers, topic: formData.kafkaTopic, + message_format: formData.kafkaMessageFormat || "json", + watermark_delay_threshold: formData.kafkaWatermarkDelay || "", }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_SPARK)) { payload.spark_options = { table: formData.sparkTable, path: formData.sparkPath, + query: formData.sparkQuery || "", + file_format: formData.sparkFileFormat || "", + table_format: formData.sparkTableFormat || "", + table_format_catalog: formData.sparkTableFormatCatalog || "", + table_format_namespace: formData.sparkTableFormatNamespace || "", + table_format_properties: formData.sparkTableFormatProperties || "", + date_partition_column: formData.sparkDatePartitionColumn || "", + date_partition_column_format: formData.sparkDatePartitionFormat || "", }; } else if (st === String(feast.core.DataSource.SourceType.BATCH_TRINO)) { payload.trino_options = { @@ -107,10 +127,32 @@ const formDataToPayload = (formData: DataSourceFormData, project: string) => { database: formData.athenaDatabase, data_source: formData.athenaDataSource, }; + if (formData.athenaDatePartitionColumn) { + payload.date_partition_column = formData.athenaDatePartitionColumn; + } + } else if (st === String(feast.core.DataSource.SourceType.BATCH_ICEBERG)) { + const catalogProps = formData.icebergCatalogProperties.trim() + ? JSON.parse(formData.icebergCatalogProperties) + : {}; + payload.custom_options = { + configuration: JSON.stringify({ + catalog_type: formData.icebergCatalogType || "rest", + endpoint: formData.icebergEndpoint, + warehouse: formData.icebergWarehouse, + namespace: formData.icebergNamespace, + table: formData.icebergTable, + token_env_var: formData.icebergTokenEnvVar || null, + credential_vending: formData.icebergCredentialVending !== "false", + catalog_properties: catalogProps, + }), + }; + payload.data_source_class_type = + "feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source.IcebergSource"; } else if (st === String(feast.core.DataSource.SourceType.STREAM_KINESIS)) { payload.kinesis_options = { region: formData.kinesisRegion, stream_name: formData.kinesisStreamName, + record_format: formData.kinesisRecordFormat || "json", }; } else if (st === String(feast.core.DataSource.SourceType.CUSTOM_SOURCE)) { payload.custom_options = { @@ -157,28 +199,54 @@ const Index = () => { createdTimestampColumn: "", tags: [] as { key: string; value: string }[], fileUri: "", + fileFormat: "parquet", + fileS3EndpointOverride: "", bigqueryTable: "", bigqueryQuery: "", + bigqueryDatePartitionColumn: "", snowflakeTable: "", snowflakeDatabase: "", snowflakeSchema: "", + snowflakeQuery: "", + snowflakeWarehouse: "", redshiftTable: "", redshiftDatabase: "", redshiftSchema: "", + redshiftQuery: "", kafkaBootstrapServers: "", kafkaTopic: "", + kafkaMessageFormat: "json", + kafkaWatermarkDelay: "", sparkTable: "", sparkPath: "", + sparkQuery: "", + sparkFileFormat: "parquet", + sparkTableFormat: "", + sparkTableFormatCatalog: "", + sparkTableFormatNamespace: "", + sparkTableFormatProperties: "", + sparkDatePartitionColumn: "", + sparkDatePartitionFormat: "%Y-%m-%d", kinesisRegion: "", kinesisStreamName: "", + kinesisRecordFormat: "json", trinoTable: "", trinoQuery: "", athenaTable: "", athenaQuery: "", athenaDatabase: "", athenaDataSource: "", + athenaDatePartitionColumn: "", customSourceClassName: "", customSourceConfig: "", + icebergCatalogType: "rest", + icebergEndpoint: "", + icebergWarehouse: "", + icebergNamespace: "", + icebergTable: "", + icebergTokenEnvVar: "", + icebergCredentialVending: "true", + icebergCatalogProperties: "", rayReaderType: "parquet", rayPath: "", rayReaderOptions: "", @@ -189,8 +257,10 @@ const Index = () => { clickhouseQuery: "", mssqlTable: "", mssqlConnectionStr: "", + mssqlDatePartitionColumn: "", oracleTable: "", oracleConnectionStr: "", + oracleDatePartitionColumn: "", couchbaseDatabase: "", couchbaseScope: "", couchbaseCollection: "", diff --git a/ui/src/pages/saved-data-sets/CreateDatasetForm.tsx b/ui/src/pages/saved-data-sets/CreateDatasetForm.tsx index a3be0bae8df..265d6e663e9 100644 --- a/ui/src/pages/saved-data-sets/CreateDatasetForm.tsx +++ b/ui/src/pages/saved-data-sets/CreateDatasetForm.tsx @@ -69,6 +69,21 @@ const STORAGE_TYPES = [ { value: "spark", label: "Spark", placeholder: "s3://bucket/path/" }, { value: "trino", label: "Trino", placeholder: "catalog.schema.table" }, { value: "athena", label: "Athena", placeholder: "database.table" }, + { + value: "postgres", + label: "PostgreSQL", + placeholder: "schema.table_name", + }, + { + value: "clickhouse", + label: "ClickHouse", + placeholder: "database.table_name", + }, + { + value: "couchbase", + label: "Couchbase Columnar", + placeholder: "database.scope.collection", + }, ]; const CreateDatasetForm = ({ onClose }: CreateDatasetFormProps) => { @@ -105,6 +120,7 @@ const CreateDatasetForm = ({ onClose }: CreateDatasetFormProps) => { const [datasetName, setDatasetName] = useState(""); const [storageType, setStorageType] = useState("file"); const [storagePath, setStoragePath] = useState(""); + const [storageFileFormat, setStorageFileFormat] = useState("parquet"); const [tags, setTags] = useState([]); const [allowOverwrite] = useState(false); @@ -199,6 +215,11 @@ const CreateDatasetForm = ({ onClose }: CreateDatasetFormProps) => { if (dsType === 7) detectedTypes.add("spark"); if (dsType === 8) detectedTypes.add("trino"); if (dsType === 9) detectedTypes.add("athena"); + const classType = + spec.dataSourceClassType || ds.dataSourceClassType || ""; + if (classType.includes("postgres")) detectedTypes.add("postgres"); + if (classType.includes("clickhouse")) detectedTypes.add("clickhouse"); + if (classType.includes("couchbase")) detectedTypes.add("couchbase"); } // Always include file as a fallback @@ -256,6 +277,8 @@ const CreateDatasetForm = ({ onClose }: CreateDatasetFormProps) => { project: projectName || "", storage_type: storageType, storage_path: storagePath.trim(), + storage_file_format: + storageType === "spark" ? storageFileFormat : undefined, entity_source_type: entitySourceType, allow_overwrite: allowOverwrite, tags: tags.reduce( @@ -588,6 +611,44 @@ const CreateDatasetForm = ({ onClose }: CreateDatasetFormProps) => { + {storageType === "spark" && ( + <> + + + Parquet, + }, + { + value: "avro", + inputDisplay: "Avro", + dropdownDisplay: Avro, + }, + { + value: "csv", + inputDisplay: "CSV", + dropdownDisplay: CSV, + }, + { + value: "json", + inputDisplay: "JSON", + dropdownDisplay: JSON, + }, + ]} + valueOfSelected={storageFileFormat} + onChange={setStorageFileFormat} + fullWidth + /> + + + )} + diff --git a/ui/src/pages/saved-data-sets/DatasetOverviewTab.tsx b/ui/src/pages/saved-data-sets/DatasetOverviewTab.tsx index b765da70241..9c3d013aaab 100644 --- a/ui/src/pages/saved-data-sets/DatasetOverviewTab.tsx +++ b/ui/src/pages/saved-data-sets/DatasetOverviewTab.tsx @@ -15,6 +15,7 @@ import { EuiCallOut, } from "@elastic/eui"; import { useParams } from "react-router-dom"; +import EuiCustomLink from "../../components/EuiCustomLink"; import DatasetFeaturesTable from "./DatasetFeaturesTable"; import DatasetJoinKeysTable from "./DatasetJoinKeysTable"; import useLoadDataset from "./useLoadDataset"; @@ -56,7 +57,7 @@ function formatTimestamp(ts: any): string { } const DatasetOverviewTab = () => { - let { datasetName } = useParams(); + const { datasetName, projectName } = useParams(); if (!datasetName) { throw new Error( @@ -96,6 +97,7 @@ const DatasetOverviewTab = () => { const tags = data.spec?.tags || {}; const featureServiceName = data.spec?.featureServiceName || data.spec?.feature_service_name; + const dataSources: string[] = data.spec?.dataSources || []; const createdTs = data.meta?.createdTimestamp || data.meta?.created_timestamp; const minEventTs = data.meta?.minEventTimestamp || data.meta?.min_event_timestamp; @@ -180,6 +182,26 @@ const DatasetOverviewTab = () => { + {dataSources.length > 0 && ( + <> + + Data Source{dataSources.length > 1 ? "s" : ""} + + + {dataSources.map((dsName, idx) => ( + + {idx > 0 && ", "} + + {dsName} + + + ))} + + + )} + {featureServiceName && ( <> diff --git a/ui/src/pages/saved-data-sets/DatasetsCardGrid.tsx b/ui/src/pages/saved-data-sets/DatasetsCardGrid.tsx index 2c96a0ee109..955f50681a3 100644 --- a/ui/src/pages/saved-data-sets/DatasetsCardGrid.tsx +++ b/ui/src/pages/saved-data-sets/DatasetsCardGrid.tsx @@ -11,6 +11,7 @@ import { EuiToolTip, EuiIcon, EuiCopy, + EuiLink, } from "@elastic/eui"; import { useNavigate, useParams } from "react-router-dom"; @@ -213,6 +214,28 @@ const DatasetCard: React.FC = ({ )} + {spec.dataSources && spec.dataSources.length > 0 && ( + + + Source + + + {spec.dataSources.map((dsName: string, idx: number) => ( + + {idx > 0 && ", "} + { + e.stopPropagation(); + navigate(`/p/${datasetProject}/data-source/${dsName}`); + }} + > + {dsName} + + + ))} + + + )} {/* Spacer to push footer */} diff --git a/ui/src/pages/saved-data-sets/DatasetsListingTable.tsx b/ui/src/pages/saved-data-sets/DatasetsListingTable.tsx index 823f42176f3..762ca3e7a7e 100644 --- a/ui/src/pages/saved-data-sets/DatasetsListingTable.tsx +++ b/ui/src/pages/saved-data-sets/DatasetsListingTable.tsx @@ -15,6 +15,9 @@ function detectStorageType(dataset: any): string { if (storage.snowflakeStorage) return "Snowflake"; if (storage.redshiftStorage) return "Redshift"; if (storage.sparkStorage) return "Spark"; + if (storage.trinoStorage) return "Trino"; + if (storage.athenaStorage) return "Athena"; + if (storage.customStorage) return "Custom"; return "—"; } @@ -53,6 +56,26 @@ const DatasetsListingTable = ({ datasets }: DatasetsListingTableProps) => { ), width: "120px", }, + { + name: "Data Source", + render: (item: any) => { + const dsList: string[] = item.spec?.dataSources || []; + if (dsList.length === 0) return "—"; + const itemProject = item.project || item.spec?.project || projectName; + return ( + <> + {dsList.map((dsName, idx) => ( + + {idx > 0 && ", "} + + {dsName} + + + ))} + + ); + }, + }, { name: "Feature Service", render: (item: any) => diff --git a/ui/src/pages/saved-data-sets/EditDatasetModal.tsx b/ui/src/pages/saved-data-sets/EditDatasetModal.tsx index db567a15703..c99e43d3824 100644 --- a/ui/src/pages/saved-data-sets/EditDatasetModal.tsx +++ b/ui/src/pages/saved-data-sets/EditDatasetModal.tsx @@ -99,6 +99,33 @@ const ALL_STORAGE_TYPES: StorageTypeDef[] = [ helpText: "Athena table reference.", sourceTypeMatch: ["BATCH_ATHENA"], }, + { + value: "postgres", + label: "PostgreSQL", + description: "PostgreSQL table reference", + placeholder: "schema.table_name", + helpText: + "PostgreSQL table reference. Data is read via the PostgreSQL offline store.", + sourceTypeMatch: ["CUSTOM_SOURCE"], + }, + { + value: "clickhouse", + label: "ClickHouse", + description: "ClickHouse table reference", + placeholder: "database.table_name", + helpText: + "ClickHouse table reference. Data is read via the ClickHouse offline store.", + sourceTypeMatch: ["CUSTOM_SOURCE"], + }, + { + value: "couchbase", + label: "Couchbase Columnar", + description: "Couchbase Columnar collection reference", + placeholder: "database.scope.collection", + helpText: + "Couchbase Columnar reference in format: database.scope.collection", + sourceTypeMatch: ["CUSTOM_SOURCE"], + }, { value: "custom", label: "Custom", @@ -128,6 +155,11 @@ function detectDataSourceTypes(dataSources: any[]): Set { if (ds.spec?.trinoOptions || ds.trinoOptions) types.add("BATCH_TRINO"); if (ds.spec?.athenaOptions || ds.athenaOptions) types.add("BATCH_ATHENA"); if (ds.spec?.customOptions || ds.customOptions) types.add("CUSTOM_SOURCE"); + const classType = + ds.spec?.dataSourceClassType || ds.dataSourceClassType || ""; + if (classType.includes("postgres")) types.add("CUSTOM_SOURCE"); + if (classType.includes("clickhouse")) types.add("CUSTOM_SOURCE"); + if (classType.includes("couchbase")) types.add("CUSTOM_SOURCE"); } return types; } @@ -160,7 +192,25 @@ function detectStorageType(dataset: any): string { if (storage.sparkStorage) return "spark"; if (storage.trinoStorage) return "trino"; if (storage.athenaStorage) return "athena"; - if (storage.customStorage) return "custom"; + if (storage.customStorage) { + try { + const config = storage.customStorage.configuration || ""; + const parsed = typeof config === "string" ? JSON.parse(config) : config; + if (parsed.database && parsed.scope && parsed.collection) + return "couchbase"; + if (parsed.table) { + const classType = + dataset?.spec?.dataSourceClassType || + dataset?.dataSourceClassType || + ""; + if (classType.includes("postgres")) return "postgres"; + if (classType.includes("clickhouse")) return "clickhouse"; + } + } catch { + // fall through + } + return "custom"; + } return "file"; } @@ -180,6 +230,14 @@ function extractStoragePath(dataset: any): string { return ""; } +function extractStorageFileFormat(dataset: any): string { + const storage = dataset?.spec?.storage; + if (storage?.sparkStorage?.fileFormat) return storage.sparkStorage.fileFormat; + if (storage?.sparkStorage?.file_format) + return storage.sparkStorage.file_format; + return "parquet"; +} + const EditDatasetModal = ({ dataset, onClose, @@ -307,6 +365,9 @@ const EditDatasetModal = ({ // Form state const [storagePath, setStoragePath] = useState(extractStoragePath(dataset)); const [storageType, setStorageType] = useState(detectStorageType(dataset)); + const [storageFileFormat, setStorageFileFormat] = useState( + extractStorageFileFormat(dataset), + ); const [featuresInput, setFeaturesInput] = useState( (spec.features || []).map((f: string) => ({ label: f })), ); @@ -365,6 +426,8 @@ const EditDatasetModal = ({ join_keys: joinKeysInput.map((o) => o.label), storage_path: storagePath.trim(), storage_type: storageType, + storage_file_format: + storageType === "spark" ? storageFileFormat : undefined, tags: tagsObj, full_feature_names: fullFeatureNames, feature_service_name: featureServiceName || undefined, @@ -490,6 +553,41 @@ const EditDatasetModal = ({ /> + {storageType === "spark" && ( + + Parquet, + }, + { + value: "avro", + inputDisplay: "Avro", + dropdownDisplay: Avro, + }, + { + value: "csv", + inputDisplay: "CSV", + dropdownDisplay: CSV, + }, + { + value: "json", + inputDisplay: "JSON", + dropdownDisplay: JSON, + }, + ]} + valueOfSelected={storageFileFormat} + onChange={setStorageFileFormat} + fullWidth + /> + + )} + diff --git a/ui/src/pages/saved-data-sets/RegisterDatasetModal.tsx b/ui/src/pages/saved-data-sets/RegisterDatasetModal.tsx index 90f93b8ca31..9c8ba9efc91 100644 --- a/ui/src/pages/saved-data-sets/RegisterDatasetModal.tsx +++ b/ui/src/pages/saved-data-sets/RegisterDatasetModal.tsx @@ -34,6 +34,7 @@ export interface RegisterDatasetPayload { join_keys: string[]; storage_path: string; storage_type: string; + storage_file_format?: string; tags: Record; full_feature_names: boolean; feature_service_name?: string; @@ -115,6 +116,33 @@ const ALL_STORAGE_TYPES: StorageTypeDefinition[] = [ helpText: "Athena table reference. Data is queried via Athena.", sourceTypeMatch: ["BATCH_ATHENA"], }, + { + value: "postgres", + label: "PostgreSQL", + description: "PostgreSQL table reference", + placeholder: "schema.table_name", + helpText: + "PostgreSQL table reference. Data is read via the PostgreSQL offline store.", + sourceTypeMatch: ["CUSTOM_SOURCE"], + }, + { + value: "clickhouse", + label: "ClickHouse", + description: "ClickHouse table reference", + placeholder: "database.table_name", + helpText: + "ClickHouse table reference. Data is read via the ClickHouse offline store.", + sourceTypeMatch: ["CUSTOM_SOURCE"], + }, + { + value: "couchbase", + label: "Couchbase Columnar", + description: "Couchbase Columnar collection reference", + placeholder: "database.scope.collection", + helpText: + "Couchbase Columnar reference in format: database.scope.collection", + sourceTypeMatch: ["CUSTOM_SOURCE"], + }, { value: "custom", label: "Custom", @@ -144,6 +172,11 @@ function detectDataSourceTypes(dataSources: any[]): Set { if (ds.spec?.trinoOptions || ds.trinoOptions) types.add("BATCH_TRINO"); if (ds.spec?.athenaOptions || ds.athenaOptions) types.add("BATCH_ATHENA"); if (ds.spec?.customOptions || ds.customOptions) types.add("CUSTOM_SOURCE"); + const classType = + ds.spec?.dataSourceClassType || ds.dataSourceClassType || ""; + if (classType.includes("postgres")) types.add("CUSTOM_SOURCE"); + if (classType.includes("clickhouse")) types.add("CUSTOM_SOURCE"); + if (classType.includes("couchbase")) types.add("CUSTOM_SOURCE"); } return types; } @@ -243,6 +276,7 @@ const RegisterDatasetModal = ({ const [name, setName] = useState(""); const [storagePath, setStoragePath] = useState(""); const [storageType, setStorageType] = useState("file"); + const [storageFileFormat, setStorageFileFormat] = useState("parquet"); const [featuresInput, setFeaturesInput] = useState( [], ); @@ -336,6 +370,8 @@ const RegisterDatasetModal = ({ join_keys: joinKeysInput.map((o) => o.label), storage_path: storagePath.trim(), storage_type: storageType, + storage_file_format: + storageType === "spark" ? storageFileFormat : undefined, tags: tagsObj, full_feature_names: fullFeatureNames, feature_service_name: featureServiceName || undefined, @@ -464,6 +500,41 @@ const RegisterDatasetModal = ({ /> + {storageType === "spark" && ( + + Parquet, + }, + { + value: "avro", + inputDisplay: "Avro", + dropdownDisplay: Avro, + }, + { + value: "csv", + inputDisplay: "CSV", + dropdownDisplay: CSV, + }, + { + value: "json", + inputDisplay: "JSON", + dropdownDisplay: JSON, + }, + ]} + valueOfSelected={storageFileFormat} + onChange={setStorageFileFormat} + fullWidth + /> + + )} + diff --git a/ui/src/queries/mutations/useDataSourceMutations.ts b/ui/src/queries/mutations/useDataSourceMutations.ts index 737ced0dbfd..a59313bf7a7 100644 --- a/ui/src/queries/mutations/useDataSourceMutations.ts +++ b/ui/src/queries/mutations/useDataSourceMutations.ts @@ -15,6 +15,12 @@ interface ApplyDataSourcePayload { redshift_options?: { table: string; database: string; schema_: string }; kafka_options?: { kafka_bootstrap_servers: string; topic: string }; spark_options?: { table: string; path: string }; + custom_options?: { + configuration?: string; + class_name?: string; + config?: string; + }; + data_source_class_type?: string; } interface DeleteDataSourcePayload {