diff --git a/docs/getting-started/components/compute-engine.md b/docs/getting-started/components/compute-engine.md index 8b69b9d1a8b..4ba7ad8de5c 100644 --- a/docs/getting-started/components/compute-engine.md +++ b/docs/getting-started/components/compute-engine.md @@ -22,6 +22,7 @@ engines. |-------------------------|-------------------------------------------------------------------------------------------------|------------|------| | LocalComputeEngine | Runs on Arrow + Pandas/Polars/Dask etc., designed for light weight transformation. | ✅ | | | SparkComputeEngine | Runs on Apache Spark, designed for large-scale distributed feature generation. | ✅ | | +| TrinoComputeEngine | Runs on Trino, designed for scalable feature generation using Trino SQL. | ✅ | [docs](../../reference/compute-engine/trino.md) | | SnowflakeComputeEngine | Runs on Snowflake, designed for scalable feature generation using Snowflake SQL. | ✅ | | | LambdaComputeEngine | Runs on AWS Lambda, designed for serverless feature generation. | ✅ | | | FlinkComputeEngine | Runs on Apache Flink, designed for distributed feature generation through PyFlink Table API. | ✅ | | diff --git a/docs/reference/compute-engine/trino.md b/docs/reference/compute-engine/trino.md new file mode 100644 index 00000000000..6377de7bee7 --- /dev/null +++ b/docs/reference/compute-engine/trino.md @@ -0,0 +1,149 @@ +# Trino + +## Description + +Trino Compute Engine provides a distributed execution engine for batch materialization operations (`materialize` and `materialize-incremental`) and historical retrieval operations (`get_historical_features`). + +It is designed to handle large-scale data processing directly on Trino clusters without moving raw data to the client machine. + +### Design + +The Trino Compute engine is implemented as a subclass of `feast.infra.compute_engines.base.ComputeEngine`. + +The engine supports the following features: +- **Pushdown SQL execution**: Compiles feature pipeline operations (filtering, deduplication, time-windowed aggregations, point-in-time joins, and transformations) into Trino SQL Common Table Expressions (CTEs) executed directly on the Trino cluster. +- **Streaming online materialization**: Streams query results in chunks as PyArrow record batches directly to the online store via a thread pool, avoiding loading entire datasets into client memory. +- **Atomic offline table writes**: Materializes offline tables using a temporary staging table and rename swap (`CREATE TABLE ...__staging AS ...; DROP TABLE ...; ALTER TABLE ...__staging RENAME TO ...`). +- **Configuration inheritance**: Inherits connection details from the Trino offline store when both are configured. + +--- + +## Example + +```yaml +project: feast_trino_project +registry: data/registry.db +provider: local + +offline_store: + type: trino.offline + host: localhost + port: 8080 + catalog: iceberg + dataset: feast_offline + user: feast_user + +batch_engine: + type: trino.engine + batch_size: 10000 + write_concurrency: 4 + +online_store: + type: redis + connection_string: localhost:6379 +``` + +--- + +## Example in Python + +```python +from datetime import timedelta +from feast import ( + BatchFeatureView, + Entity, + Field, +) +from feast.aggregation import Aggregation +from feast.infra.offline_stores.contrib.trino_offline_store.trino_source import TrinoSource +from feast.transformation.mode import TransformationMode +from feast.transformation.trino_transformation import TrinoTransformation +from feast.types import Float32, Int32 +from feast.value_type import ValueType + +# 1. Define Entity +driver = Entity( + name="driver_id", + value_type=ValueType.INT32, + join_keys=["driver_id"], + description="Driver identifier", +) + +# 2. Define Trino Batch Source +driver_stats_source = TrinoSource( + name="driver_stats_source", + table="iceberg.feast_offline.driver_hourly_stats", + timestamp_field="event_timestamp", + created_timestamp_column="created_timestamp", +) + +# 3. Define Pure Trino SQL Transformation +# The '{}' placeholder will be substituted with the upstream CTE/table name +transform = TrinoTransformation( + mode=TransformationMode.TRINO_SQL, + udf=""" + SELECT + driver_id, + event_timestamp, + conv_rate * 2.0 AS conv_rate, + acc_rate * 2.0 AS acc_rate + FROM {} + """, +) + +# 4. Define Feature View with SQL Transformation & Aggregations +driver_hourly_stats_fv = BatchFeatureView( + name="driver_hourly_stats", + entities=[driver], + ttl=timedelta(days=3), + feature_transformation=transform, + aggregations=[ + Aggregation(column="conv_rate", function="sum"), + Aggregation(column="acc_rate", function="avg"), + ], + schema=[ + Field(name="sum_conv_rate", dtype=Float32), + Field(name="avg_acc_rate", dtype=Float32), + Field(name="driver_id", dtype=Int32), + ], + online=True, + offline=True, + source=driver_stats_source, +) +``` + +--- + +## Retrieving and Materializing Features + +```python +import pandas as pd +from datetime import datetime, timezone +from feast import FeatureStore + +fs = FeatureStore(repo_path=".") + +entity_df: pd.DataFrame + +# Lazy Historical Retrieval: +# Compiles into a single Trino ANSI SQL query with CTEs. +job = fs.get_historical_features( + entity_df=entity_df, + features=[ + "driver_hourly_stats:sum_conv_rate", + "driver_hourly_stats:avg_acc_rate", + ], +) + +# Inspect the compiled pure SQL query: +print(job.to_sql()) + +# Execute on cluster and return pandas dataframe: +training_df = job.to_df() + +# Materialize to Online and Offline Stores: +fs.materialize( + start_date=datetime(2025, 1, 1, tzinfo=timezone.utc), + end_date=datetime(2025, 1, 15, tzinfo=timezone.utc), +) +``` diff --git a/sdk/python/feast/infra/compute_engines/dag/model.py b/sdk/python/feast/infra/compute_engines/dag/model.py index 263c5029f4f..203a40d172d 100644 --- a/sdk/python/feast/infra/compute_engines/dag/model.py +++ b/sdk/python/feast/infra/compute_engines/dag/model.py @@ -7,3 +7,4 @@ class DAGFormat(str, Enum): ARROW = "arrow" RAY = "ray" FLINK = "flink" + TRINO = "trino" diff --git a/sdk/python/feast/infra/compute_engines/trino/__init__.py b/sdk/python/feast/infra/compute_engines/trino/__init__.py new file mode 100644 index 00000000000..82706f0bd7a --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/trino/__init__.py @@ -0,0 +1,6 @@ +from feast.infra.compute_engines.trino.compute import ( + TrinoComputeEngine, + TrinoComputeEngineConfig, +) + +__all__ = ["TrinoComputeEngine", "TrinoComputeEngineConfig"] diff --git a/sdk/python/feast/infra/compute_engines/trino/compute.py b/sdk/python/feast/infra/compute_engines/trino/compute.py new file mode 100644 index 00000000000..c740b44100b --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/trino/compute.py @@ -0,0 +1,323 @@ +from __future__ import annotations + +import logging +from datetime import datetime +from typing import Any, Dict, Literal, Optional, Sequence, Union + +from pydantic import ConfigDict, Field, StrictStr +from trino.exceptions import TrinoConnectionError + +from feast import ( + BatchFeatureView, + Entity, + FeatureView, + OnDemandFeatureView, + StreamFeatureView, +) +from feast.infra.common.materialization_job import ( + MaterializationJob, + MaterializationJobStatus, + MaterializationTask, +) +from feast.infra.common.retrieval_task import HistoricalRetrievalTask +from feast.infra.compute_engines.base import ComputeEngine +from feast.infra.compute_engines.trino.feature_builder import TrinoFeatureBuilder +from feast.infra.compute_engines.trino.job import ( + TrinoDAGRetrievalJob, + TrinoMaterializationJob, +) +from feast.infra.compute_engines.trino.utils import ( + get_or_create_trino_client, + stream_trino_arrow_batches, + write_arrow_batches_to_online_store, +) +from feast.infra.offline_stores.contrib.trino_offline_store.trino import ( + AuthConfig, + TrinoOfflineStoreConfig, +) +from feast.infra.offline_stores.contrib.trino_offline_store.trino_queries import Trino +from feast.infra.offline_stores.offline_store import OfflineStore, RetrievalJob +from feast.infra.online_stores.online_store import OnlineStore +from feast.infra.registry.base_registry import BaseRegistry +from feast.repo_config import FeastConfigBaseModel, RepoConfig +from feast.utils import _get_column_names + +logger = logging.getLogger(__name__) + + +class TrinoComputeEngineConfig(FeastConfigBaseModel): + """Configuration for Trino Compute Engine.""" + + type: Literal["trino.engine", "trino"] = "trino.engine" + """ Compute engine type selector """ + + host: Optional[StrictStr] = None + """ Host of the Trino coordinator """ + + port: Optional[int] = 8080 + """ Port of the Trino coordinator """ + + catalog: Optional[StrictStr] = None + """ Target Trino catalog (e.g. iceberg, hive, delta, memory) """ + + dataset: Optional[StrictStr] = "feast" + """ Schema/dataset name for temporary staging tables """ + + user: Optional[StrictStr] = None + """ Username for connecting to Trino """ + + source: Optional[StrictStr] = "trino-python-client" + """ Application source identifier for Trino query tracking """ + + http_scheme: Literal["http", "https"] = Field(default="http", alias="http-scheme") + """ Protocol scheme (http or https) """ + + verify: bool = True + """ Whether to verify SSL certificates """ + + extra_credential: Optional[StrictStr] = None + """ Optional extra credential header passed to Trino """ + + connector: Optional[Dict[str, str]] = None + """ Connector arguments for Trino offline store table operations """ + + auth: Optional[AuthConfig] = None + """ Authentication mechanism (kerberos, basic, oauth2) """ + + batch_size: int = 10000 + """ Streaming batch chunk size for PyArrow record batches """ + + write_concurrency: int = 4 + """ Number of concurrent threads used when writing batches to online store """ + + offline_write_mode: Literal["append", "overwrite", "merge"] = "append" + """ Write mode for offline feature view materialization ('append', 'overwrite', or 'merge') """ + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + +class TrinoComputeEngine(ComputeEngine): + """Distributed Trino SQL compute engine for Feast materialization and retrieval.""" + + def __init__( + self, + *, + repo_config: RepoConfig, + offline_store: OfflineStore, + online_store: OnlineStore, + **kwargs: Any, + ): + super().__init__( + repo_config=repo_config, + offline_store=offline_store, + online_store=online_store, + **kwargs, + ) + + batch_engine_config = repo_config.batch_engine + if isinstance(batch_engine_config, TrinoComputeEngineConfig): + self.config = batch_engine_config + elif isinstance(batch_engine_config, dict): + self.config = TrinoComputeEngineConfig(**batch_engine_config) + elif isinstance(batch_engine_config, str): + self.config = TrinoComputeEngineConfig() + else: + self.config = TrinoComputeEngineConfig() + + # Automatic fallback: Inherit connection details from Trino offline store if unspecified + if ( + isinstance(repo_config.offline_store, TrinoOfflineStoreConfig) + and self.config.host is None + ): + offline_conf = repo_config.offline_store + self.config.host = offline_conf.host + self.config.port = offline_conf.port + self.config.catalog = offline_conf.catalog + self.config.user = offline_conf.user + self.config.source = offline_conf.source + self.config.http_scheme = offline_conf.http_scheme + self.config.verify = offline_conf.verify + self.config.extra_credential = offline_conf.extra_credential + self.config.connector = offline_conf.connector + self.config.dataset = offline_conf.dataset + self.config.auth = offline_conf.auth + + self.client: Optional[Trino] = None + if self.config.host and self.config.catalog and self.config.user: + self.client = get_or_create_trino_client(self.config) + + def update( + self, + project: str, + views_to_delete: Sequence[ + Union[BatchFeatureView, StreamFeatureView, FeatureView] + ], + views_to_keep: Sequence[ + Union[BatchFeatureView, StreamFeatureView, FeatureView, OnDemandFeatureView] + ], + entities_to_delete: Sequence[Entity], + entities_to_keep: Sequence[Entity], + ) -> None: + """Trino compute engine does not require persistent infrastructure updates.""" + logger.debug("Trino compute engine does not require updating") + pass + + def teardown_infra( + self, + project: str, + fvs: Sequence[Union[BatchFeatureView, StreamFeatureView, FeatureView]], + entities: Sequence[Entity], + ) -> None: + """Trino compute engine does not require infrastructure teardown.""" + logger.debug("Trino compute engine does not require teardown") + pass + + def _materialize_one( + self, + registry: BaseRegistry, + task: MaterializationTask, + from_offline_store: bool = False, + **kwargs: Any, + ) -> MaterializationJob: + if from_offline_store: + return self._materialize_from_offline_store( + registry=registry, + feature_view=task.feature_view, + start_date=task.start_time, + end_date=task.end_time, + project=task.project, + ) + + job_id = f"{task.feature_view.name}-{task.start_time}-{task.end_time}" + context = self.get_execution_context(registry, task) + + try: + if self.client is None: + raise RuntimeError( + "Trino client is not configured. Set host, catalog, and user " + "in batch_engine config or use a TrinoOfflineStoreConfig." + ) + builder = TrinoFeatureBuilder( + registry=registry, + client=self.client, + task=task, + ) + plan = builder.build() + plan.execute(context) + + return TrinoMaterializationJob( + job_id=job_id, status=MaterializationJobStatus.SUCCEEDED + ) + except TrinoConnectionError as ce: + logger.error("Connection to Trino failed during materialization: %s", ce) + return TrinoMaterializationJob( + job_id=job_id, status=MaterializationJobStatus.ERROR, error=ce + ) + except Exception as e: + logger.debug("Materialization failed for %s: %s", task.feature_view.name, e) + return TrinoMaterializationJob( + job_id=job_id, status=MaterializationJobStatus.ERROR, error=e + ) + + def _materialize_from_offline_store( + self, + registry: BaseRegistry, + feature_view: Union[BatchFeatureView, StreamFeatureView, FeatureView], + start_date: datetime, + end_date: datetime, + project: str, + ) -> MaterializationJob: + """Legacy materialization fallback method.""" + logger.warning( + "Materializing from offline store will be deprecated in the future. " + "Please use the unified DAG materialization engine." + ) + entities = [] + for entity_name in feature_view.entities: + entities.append(registry.get_entity(entity_name, project)) + + ( + join_key_columns, + feature_name_columns, + timestamp_field, + created_timestamp_column, + ) = _get_column_names(feature_view, entities) + + job_id = f"{feature_view.name}-{start_date}-{end_date}" + + try: + offline_job = self.offline_store.pull_latest_from_table_or_query( + config=self.repo_config, + data_source=feature_view.batch_source, # type: ignore[arg-type] + join_key_columns=join_key_columns, + feature_name_columns=feature_name_columns, + timestamp_field=timestamp_field, + created_timestamp_column=created_timestamp_column, + start_date=start_date, + end_date=end_date, + ) + + if not self.client: + raise RuntimeError("Trino client required for materialization.") + + batches = stream_trino_arrow_batches( + client=self.client, + query_text=offline_job.to_sql(), + batch_size=self.config.batch_size, + ) + write_arrow_batches_to_online_store( + batches=batches, + feature_view=feature_view, # type: ignore[arg-type] + online_store=self.online_store, + repo_config=self.repo_config, + concurrency=self.config.write_concurrency, + ) + + return TrinoMaterializationJob( + job_id=job_id, status=MaterializationJobStatus.SUCCEEDED + ) + except Exception as e: + return TrinoMaterializationJob( + job_id=job_id, status=MaterializationJobStatus.ERROR, error=e + ) + + def get_historical_features( + self, registry: BaseRegistry, task: HistoricalRetrievalTask + ) -> RetrievalJob: + """Compile and return lazy Trino historical retrieval job.""" + context = self.get_execution_context(registry, task) + + try: + if self.client is None: + raise RuntimeError( + "Trino client is not configured. Set host, catalog, and user " + "in batch_engine config or use a TrinoOfflineStoreConfig." + ) + builder = TrinoFeatureBuilder( + registry=registry, + client=self.client, + task=task, + ) + plan = builder.build() + + return TrinoDAGRetrievalJob( + client=self.client, + plan=plan, + context=context, + config=self.repo_config, + full_feature_names=task.full_feature_name, + on_demand_feature_views=getattr(task, "on_demand_feature_views", None), + metadata=getattr(task, "metadata", None), + ) + except Exception as e: + logger.error("Failed to build historical feature retrieval DAG: %s", e) + return TrinoDAGRetrievalJob( + client=self.client, + plan=None, + context=context, + config=self.repo_config, + full_feature_names=task.full_feature_name, + on_demand_feature_views=getattr(task, "on_demand_feature_views", None), + metadata=getattr(task, "metadata", None), + error=e, + ) diff --git a/sdk/python/feast/infra/compute_engines/trino/feature_builder.py b/sdk/python/feast/infra/compute_engines/trino/feature_builder.py new file mode 100644 index 00000000000..911f7c803dd --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/trino/feature_builder.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, List, Optional, Union + +import pandas as pd + +from feast.infra.common.materialization_job import MaterializationTask +from feast.infra.common.retrieval_task import HistoricalRetrievalTask +from feast.infra.compute_engines.dag.node import DAGNode +from feast.infra.compute_engines.feature_builder import FeatureBuilder +from feast.infra.compute_engines.trino.nodes import ( + TrinoAggregationNode, + TrinoDedupNode, + TrinoFilterNode, + TrinoJoinNode, + TrinoReadNode, + TrinoTransformationNode, + TrinoValidationNode, + TrinoWriteNode, +) +from feast.infra.compute_engines.trino.utils import from_feast_to_trino_type +from feast.infra.offline_stores.contrib.trino_offline_store.trino_queries import Trino +from feast.infra.registry.base_registry import BaseRegistry + +if TYPE_CHECKING: + from feast import BatchFeatureView, FeatureView, StreamFeatureView + +logger = logging.getLogger(__name__) + + +class TrinoFeatureBuilder(FeatureBuilder): + """Translates FeatureView definitions and tasks into a pure Trino SQL DAG ExecutionPlan.""" + + def __init__( + self, + registry: BaseRegistry, + client: Trino, + task: Union[MaterializationTask, HistoricalRetrievalTask], + ) -> None: + super().__init__(registry, task.feature_view, task) + self.client = client + + def _should_join_entity_df(self) -> bool: + return isinstance(self.task, HistoricalRetrievalTask) and ( + isinstance(self.task.entity_df, pd.DataFrame) + or ( + isinstance(self.task.entity_df, str) + and bool(self.task.entity_df.strip()) + ) + ) + + def _build( + self, + view: Union[BatchFeatureView, StreamFeatureView, FeatureView, Any], + input_nodes: Optional[List[DAGNode]], + ) -> DAGNode: + if getattr(view, "data_source", None): + last_node: DAGNode = self.build_source_node(view) + + if self._should_transform(view): + last_node = self.build_transformation_node(view, [last_node]) + + if self._should_join_entity_df(): + last_node = self.build_join_node(view, [last_node]) + + elif input_nodes: + if self._should_transform(view): + last_node = self.build_transformation_node(view, input_nodes) + else: + last_node = self.build_join_node(view, input_nodes) + else: + raise ValueError(f"FeatureView {view.name} has no valid source or inputs") + + last_node = self.build_filter_node(view, last_node) + + if self._should_aggregate(view): + last_node = self.build_aggregation_node(view, last_node) + elif self._should_dedupe(view): + last_node = self.build_dedup_node(view, last_node) + + if self._should_validate(view): + last_node = self.build_validation_node(view, last_node) + + return last_node + + def build_source_node( + self, view: Union[BatchFeatureView, StreamFeatureView, FeatureView] + ) -> TrinoReadNode: + start_time = getattr(self.task, "start_time", None) + end_time = getattr(self.task, "end_time", None) + source = view.batch_source + if source is None: + raise ValueError(f"FeatureView {view.name} has no batch_source defined") + column_info = self.get_column_info(view) + return TrinoReadNode( + name=f"{view.name}:source", + source=source, + column_info=column_info, + client=self.client, + start_time=start_time, + end_time=end_time, + ) + + def build_aggregation_node( + self, + view: Union[BatchFeatureView, StreamFeatureView, FeatureView], + input_node: DAGNode, + ) -> TrinoAggregationNode: + agg_specs = getattr(view, "aggregations", None) or [] + group_by_keys = getattr(view, "entities", []) or [] + batch_source = view.batch_source + timestamp_col = ( + batch_source.timestamp_field + if batch_source and hasattr(batch_source, "timestamp_field") + else None + ) + if not timestamp_col: + raise ValueError( + f"FeatureView {view.name} missing timestamp_field for aggregation" + ) + + enable_tiling = getattr(view, "enable_tiling", False) + hop_size = getattr(view, "tiling_hop_size", None) + + return TrinoAggregationNode( + name=f"{view.name}:agg", + aggregations=agg_specs, + group_by_keys=group_by_keys, + timestamp_col=timestamp_col, + client=self.client, + inputs=[input_node], + enable_tiling=enable_tiling, + hop_size=hop_size, + ) + + def build_join_node( + self, + view: Union[BatchFeatureView, StreamFeatureView, FeatureView], + input_nodes: List[DAGNode], + ) -> TrinoJoinNode: + column_info = self.get_column_info(view) + return TrinoJoinNode( + name=f"{view.name}:join", + column_info=column_info, + client=self.client, + inputs=input_nodes, + how="left", + ) + + def build_filter_node( + self, + view: Union[BatchFeatureView, StreamFeatureView, FeatureView], + input_node: DAGNode, + ) -> TrinoFilterNode: + """Build filter node applying TTL and feature view filter expression. + + Note: + ``view.filter`` is rendered directly into the Trino SQL WHERE clause and + must be trusted developer-authored input. + """ + filter_expr = getattr(view, "filter", None) + ttl = getattr(view, "ttl", None) + column_info = self.get_column_info(view) + return TrinoFilterNode( + name=f"{view.name}:filter", + column_info=column_info, + client=self.client, + ttl=ttl, + filter_condition=filter_expr, + inputs=[input_node], + ) + + def build_dedup_node( + self, + view: Union[BatchFeatureView, StreamFeatureView, FeatureView], + input_node: DAGNode, + ) -> TrinoDedupNode: + column_info = self.get_column_info(view) + return TrinoDedupNode( + name=f"{view.name}:dedup", + column_info=column_info, + client=self.client, + inputs=[input_node], + ) + + def build_transformation_node( + self, + view: Union[BatchFeatureView, StreamFeatureView, FeatureView], + input_nodes: List[DAGNode], + ) -> TrinoTransformationNode: + feature_transformation = getattr(view, "feature_transformation", None) + if feature_transformation is None: + raise ValueError( + f"FeatureView {view.name} has no feature_transformation defined" + ) + udf_name = feature_transformation.name + udf = feature_transformation.udf + udf_string = getattr(feature_transformation, "udf_string", "") or "" + return TrinoTransformationNode( + name=udf_name, + udf=udf, + inputs=input_nodes, + udf_string=udf_string, + ) + + def build_output_nodes( + self, + view: Union[BatchFeatureView, StreamFeatureView, FeatureView], + final_node: DAGNode, + ) -> TrinoWriteNode: + return TrinoWriteNode( + name=f"{view.name}:output", + feature_view=self.dag_root.view, + client=self.client, + inputs=[final_node], + ) + + def build_validation_node( + self, + view: Union[BatchFeatureView, StreamFeatureView, FeatureView], + input_node: DAGNode, + ) -> TrinoValidationNode: + features = getattr(view, "features", []) or [] + expected_columns = { + feature.name: from_feast_to_trino_type(feature.dtype) + for feature in features + } + json_columns = { + feature.name + for feature in features + if getattr(feature.dtype, "name", None) == "JSON" + } + + return TrinoValidationNode( + name=f"{view.name}:validate", + expected_columns=expected_columns, + json_columns=json_columns, + inputs=[input_node], + ) diff --git a/sdk/python/feast/infra/compute_engines/trino/job.py b/sdk/python/feast/infra/compute_engines/trino/job.py new file mode 100644 index 00000000000..b4849f02f49 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/trino/job.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import List, Optional + +import pandas as pd +import pyarrow as pa +from trino.exceptions import TrinoConnectionError, TrinoQueryError + +from feast import OnDemandFeatureView, RepoConfig +from feast.infra.common.materialization_job import ( + MaterializationJob, + MaterializationJobStatus, +) +from feast.infra.compute_engines.dag.context import ExecutionContext +from feast.infra.compute_engines.dag.plan import ExecutionPlan +from feast.infra.compute_engines.trino.sql_builder import TrinoQueryPlan +from feast.infra.offline_stores.contrib.trino_offline_store.trino_queries import Trino +from feast.infra.offline_stores.offline_store import RetrievalJob, RetrievalMetadata +from feast.saved_dataset import SavedDatasetStorage + +logger = logging.getLogger(__name__) + + +class TrinoDAGRetrievalJob(RetrievalJob): + """Lazy retrieval job returned by Trino DAG execution. + + Maintains the compiled SQL query without executing it eagerly. Data is streamed + from Trino only when .to_arrow() or .to_df() is explicitly called. + """ + + def __init__( + self, + client: Optional[Trino], + plan: Optional[ExecutionPlan], + context: Optional[ExecutionContext], + full_feature_names: bool, + config: RepoConfig, + on_demand_feature_views: Optional[List[OnDemandFeatureView]] = None, + metadata: Optional[RetrievalMetadata] = None, + error: Optional[BaseException] = None, + query: Optional[str] = None, + ): + self._client = client + self._plan = plan + self._context = context + self._full_feature_names = full_feature_names + self._config = config + self._on_demand_feature_views = on_demand_feature_views or [] + self._metadata = metadata + self._error = error + self._query = query or "" + self._arrow_table: Optional[pa.Table] = None + + @property + def full_feature_names(self) -> bool: + return self._full_feature_names + + @property + def on_demand_feature_views(self) -> List[OnDemandFeatureView]: + return self._on_demand_feature_views + + @property + def metadata(self) -> Optional[RetrievalMetadata]: + return self._metadata + + def error(self) -> Optional[BaseException]: + return self._error + + def to_sql(self) -> str: + """Return the compiled Trino ANSI SQL query for the feature retrieval DAG.""" + if self._query: + return self._query + + if self._plan and self._context: + result = self._plan.execute(self._context) + if isinstance(result.data, TrinoQueryPlan): + self._query = result.data.to_sql() + return self._query + elif isinstance(result.data, str): + self._query = result.data + return self._query + + return self._query + + def _ensure_executed(self) -> pa.Table: + """Execute the compiled SQL query against Trino cluster and cache as Arrow table.""" + if self._error: + raise self._error + + if self._arrow_table is not None: + return self._arrow_table + + query = self.to_sql() + if not self._client or not query: + raise RuntimeError( + "No Trino client or SQL query available to retrieve data." + ) + + try: + results = self._client.execute_query(query) + df = results.to_dataframe() + self._arrow_table = pa.Table.from_pandas(df) + return self._arrow_table + except TrinoConnectionError as ce: + logger.error( + "Connection failed while executing Trino retrieval query: %s", ce + ) + raise + except TrinoQueryError as e: + logger.debug("Trino query error during retrieval execution: %s", e) + raise + + def _to_arrow_internal(self, timeout: Optional[int] = None) -> pa.Table: + return self._ensure_executed() + + def _to_df_internal(self, timeout: Optional[int] = None) -> pd.DataFrame: + return self._ensure_executed().to_pandas() + + def persist( + self, + storage: SavedDatasetStorage, + allow_overwrite: Optional[bool] = False, + timeout: Optional[int] = None, + ) -> None: + """Persist training dataset directly to a Trino table via CREATE TABLE AS SELECT.""" + from feast.infra.offline_stores.contrib.trino_offline_store.trino_source import ( + SavedDatasetTrinoStorage, + ) + + if not isinstance(storage, SavedDatasetTrinoStorage): + raise ValueError( + f"Expected `SavedDatasetTrinoStorage` but received {type(storage)}" + ) + + destination_table = storage.trino_options.table + query = self.to_sql() + if not self._client or not query: + raise RuntimeError( + "Cannot persist dataset without client and compiled query." + ) + + if allow_overwrite: + try: + self._client.execute_query(f"DROP TABLE IF EXISTS {destination_table}") + except TrinoConnectionError: + raise + except TrinoQueryError as e: + logger.debug( + "Failed to drop existing table %s: %s", destination_table, e + ) + + create_query = f"CREATE TABLE {destination_table} AS ({query})" + self._client.execute_query(create_query) + + +@dataclass +class TrinoMaterializationJob(MaterializationJob): + """Tracks materialization jobs executed via TrinoComputeEngine.""" + + def __init__( + self, + job_id: str, + status: MaterializationJobStatus, + error: Optional[BaseException] = None, + ) -> None: + super().__init__() + self._job_id: str = job_id + self._status: MaterializationJobStatus = status + self._error: Optional[BaseException] = error + + def status(self) -> MaterializationJobStatus: + return self._status + + def error(self) -> Optional[BaseException]: + return self._error + + def should_be_retried(self) -> bool: + # TODO(@mrosti): we could probably retry some trino failures but no one uses this + return False + + def job_id(self) -> str: + return self._job_id + + def url(self) -> Optional[str]: + return None diff --git a/sdk/python/feast/infra/compute_engines/trino/nodes.py b/sdk/python/feast/infra/compute_engines/trino/nodes.py new file mode 100644 index 00000000000..54b35c35112 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/trino/nodes.py @@ -0,0 +1,864 @@ +from __future__ import annotations + +import logging +from datetime import datetime, timedelta +from typing import ( + Any, + Mapping, + Optional, + Sequence, + Set, + Tuple, + Union, +) + +import pandas as pd +from trino.exceptions import TrinoQueryError + +from feast import BatchFeatureView, FeatureView, StreamFeatureView +from feast.aggregation import Aggregation +from feast.data_source import DataSource +from feast.infra.compute_engines.dag.context import ColumnInfo, ExecutionContext +from feast.infra.compute_engines.dag.model import DAGFormat +from feast.infra.compute_engines.dag.node import DAGNode +from feast.infra.compute_engines.dag.value import DAGValue +from feast.infra.compute_engines.trino.sql_builder import TrinoQueryPlan +from feast.infra.compute_engines.trino.utils import ( + _has_unquoted_semicolon, + _trino_sql_literal, + quote_identifier, + stream_trino_arrow_batches, + unique_ordered, + write_arrow_batches_to_online_store, +) +from feast.infra.compute_engines.utils import ( + ENTITY_TS_ALIAS, + infer_entity_timestamp_column, +) +from feast.infra.offline_stores.contrib.trino_offline_store.connectors.upload import ( + upload_pandas_dataframe_to_trino, +) +from feast.infra.offline_stores.contrib.trino_offline_store.trino_queries import Trino +from feast.infra.offline_stores.offline_utils import get_temp_entity_table_name + +logger = logging.getLogger(__name__) + + +def _format_agg_expression(agg: Aggregation) -> str: + """Pure mapping of an Aggregation spec to its ANSI SQL Trino expression.""" + col_quoted = quote_identifier(agg.column) + resolved_alias = quote_identifier(agg.resolved_name(agg.time_window)) + + match agg.function.lower(): + case "count_distinct": + expr = f"approx_distinct({col_quoted})" + case "count" | "sum" | "avg" | "min" | "max" as func: + expr = f"{func.upper()}({col_quoted})" + case custom: + expr = f"{custom.upper()}({col_quoted})" + + return f"{expr} AS {resolved_alias}" + + +class TrinoReadNode(DAGNode): + """Generates the root Trino SQL CTE from the batch data source. + + Performs column projection, field_mapping aliasing, and timestamp filtering + 100% inside the Trino query without downloading raw source data. + """ + + def __init__( + self, + name: str, + source: DataSource, + column_info: ColumnInfo, + client: Trino, + start_time: Optional[datetime] = None, + end_time: Optional[datetime] = None, + ): + super().__init__(name) + self.source = source + self.column_info = column_info + self.client = client + self.start_time = start_time + self.end_time = end_time + + def execute(self, context: ExecutionContext) -> DAGValue: + table_ref = self.source.get_table_query_string() + try: + timestamp_col = self.column_info.timestamp_column + except (ValueError, AttributeError): + timestamp_col = None + try: + created_ts_col = self.column_info.created_timestamp_column + except (ValueError, AttributeError): + created_ts_col = None + join_keys = self.column_info.join_keys + field_mapping: dict[str, str] = self.column_info.field_mapping or {} + + if self.column_info.feature_cols: + cols_to_select = set(join_keys) + if timestamp_col: + cols_to_select.add(timestamp_col) + if created_ts_col: + cols_to_select.add(created_ts_col) + cols_to_select.update(self.column_info.feature_cols) + + projections = [ + f"{quote_identifier(str(col))} AS {quote_identifier(field_mapping[str(col)])}" + if str(col) in field_mapping and field_mapping[str(col)] != str(col) + else quote_identifier(str(col)) + for col in sorted(cols_to_select, key=str) + ] + select_clause = ", ".join(projections) + output_cols = tuple( + field_mapping.get(str(col), str(col)) + for col in sorted(cols_to_select, key=str) + ) + else: + select_clause = "*" + output_cols = None + + filters = [] + if timestamp_col and self.start_time: + filters.append( + f"{quote_identifier(timestamp_col)} >= {_trino_sql_literal(self.start_time)}" + ) + if timestamp_col and self.end_time: + filters.append( + f"{quote_identifier(timestamp_col)} <= {_trino_sql_literal(self.end_time)}" + ) + + where_clause = f"\nWHERE {' AND '.join(filters)}" if filters else "" + cte_name = f"_source_{self.name.replace(':', '_')}" + query = f"SELECT {select_clause}\nFROM {table_ref}{where_clause}" + + plan = TrinoQueryPlan( + ctes=((cte_name, query),), + current_from=cte_name, + columns=output_cols, + join_keys=tuple(join_keys), + timestamp_col=timestamp_col, + created_timestamp_col=created_ts_col, + metadata={"source": self.name}, + ) + + return DAGValue(data=plan, format=DAGFormat.TRINO) + + +class TrinoFilterNode(DAGNode): + """Appends a Trino SQL CTE applying point-in-time, TTL, and custom filters. + + .. warning:: + ``filter_condition`` is rendered directly into the generated SQL ``WHERE`` clause + without parameterization, matching Feast's standard design for developer-authored + feature view definitions. It must be trusted developer input; do not pass unsanitized + user input into feature view filter expressions. + """ + + def __init__( + self, + name: str, + column_info: ColumnInfo, + client: Trino, + ttl: Optional[timedelta] = None, + filter_condition: Optional[str] = None, + inputs: Optional[Sequence[DAGNode]] = None, + ): + """Initialize TrinoFilterNode. + + Args: + name: Identifier for the DAG node. + column_info: Column schema information for timestamp and entity keys. + client: Trino client instance. + ttl: Optional time-to-live timedelta for point-in-time expiration. + filter_condition: Optional raw SQL boolean expression string (e.g. from + ``FeatureView.filter``). Must be trusted developer-authored input. + inputs: Upstream DAG nodes providing input query plans. + """ + super().__init__(name, inputs=list(inputs) if inputs else None) + self.column_info = column_info + self.client = client + self.ttl = ttl + self.filter_condition = filter_condition + + def execute(self, context: ExecutionContext) -> DAGValue: + input_value = self.get_single_input_value(context) + input_value.assert_format(DAGFormat.TRINO) + upstream_plan: TrinoQueryPlan = input_value.data + + upstream_cte = upstream_plan.get_latest_cte_name() + ts_col = upstream_plan.timestamp_col + if not ts_col: + try: + ts_col = self.column_info.timestamp_column + except (ValueError, AttributeError): + ts_col = None + + conditions = [] + has_entity_ts = bool( + upstream_plan.metadata.get("entity_joined") + or (upstream_plan.columns and ENTITY_TS_ALIAS in upstream_plan.columns) + ) + if ts_col and has_entity_ts: + conditions.append( + f"{quote_identifier(ts_col)} <= {quote_identifier(ENTITY_TS_ALIAS)}" + ) + if self.ttl: + ttl_seconds = int(self.ttl.total_seconds()) + conditions.append( + f"{quote_identifier(ts_col)} >= {quote_identifier(ENTITY_TS_ALIAS)} - INTERVAL '{ttl_seconds}' SECOND" + ) + if self.filter_condition: + if _has_unquoted_semicolon(self.filter_condition): + raise ValueError( + f"Feature view filter condition contains prohibited semicolon (';'): " + f"{self.filter_condition!r}. Semicolons are statement terminators and " + f"cannot be used inside filter expressions." + ) + clean_filter = self.filter_condition.strip() + if clean_filter: + conditions.append(f"({clean_filter})") + + where_clause = f"\nWHERE {' AND '.join(conditions)}" if conditions else "" + cte_name = f"_filter_{self.name.replace(':', '_')}" + query = f"SELECT *\nFROM {upstream_cte}{where_clause}" + + plan = upstream_plan.add_cte( + name=cte_name, + query=query, + metadata={"filter_applied": True}, + ) + return DAGValue(data=plan, format=DAGFormat.TRINO) + + +class TrinoDedupNode(DAGNode): + """Appends a Trino SQL window CTE deduplicating rows to keep the latest per entity. + + Replaces Spark's Window.partitionBy().orderBy().row_number() with native Trino SQL: + ROW_NUMBER() OVER (PARTITION BY ... ORDER BY ... DESC) = 1. + """ + + def __init__( + self, + name: str, + column_info: ColumnInfo, + client: Trino, + inputs: Optional[Sequence[DAGNode]] = None, + ): + super().__init__(name, inputs=list(inputs) if inputs else None) + self.column_info = column_info + self.client = client + + def execute(self, context: ExecutionContext) -> DAGValue: + input_value = self.get_single_input_value(context) + input_value.assert_format(DAGFormat.TRINO) + upstream_plan: TrinoQueryPlan = input_value.data + + upstream_cte = upstream_plan.get_latest_cte_name() + join_keys = upstream_plan.join_keys or self.column_info.join_keys + ts_col = upstream_plan.timestamp_col + if not ts_col: + try: + ts_col = self.column_info.timestamp_column + except (ValueError, AttributeError): + ts_col = None + + created_ts_col = upstream_plan.created_timestamp_col + if not created_ts_col: + try: + created_ts_col = self.column_info.created_timestamp_column + except (ValueError, AttributeError): + created_ts_col = None + + if not join_keys or not ts_col: + return input_value + + candidate_cols = ( + *join_keys, + *(c for c in (ts_col, created_ts_col) if c), + *self.column_info.feature_cols, + ) + columns = upstream_plan.columns or unique_ordered(candidate_cols) + clean_cols = tuple(c for c in columns if c != "_feast_rn") + select_clause = ( + ", ".join(quote_identifier(c) for c in clean_cols) if clean_cols else "*" + ) + + partition_keys = [quote_identifier(k) for k in join_keys] + if ENTITY_TS_ALIAS in (upstream_plan.columns or ()): + partition_keys.append(quote_identifier(ENTITY_TS_ALIAS)) + partition_clause = ", ".join(partition_keys) + order_clauses = [f"{quote_identifier(ts_col)} DESC"] + if created_ts_col: + order_clauses.append(f"{quote_identifier(created_ts_col)} DESC") + + cte_name = f"_dedup_{self.name.replace(':', '_')}" + query = ( + f"SELECT {select_clause} FROM (\n" + f" SELECT *,\n" + f" ROW_NUMBER() OVER (\n" + f" PARTITION BY {partition_clause}\n" + f" ORDER BY {', '.join(order_clauses)}\n" + f" ) AS _feast_rn\n" + f" FROM {upstream_cte}\n" + f") AS _dedup_sub\n" + f"WHERE _feast_rn = 1" + ) + + plan = upstream_plan.add_cte( + name=cte_name, + query=query, + columns=clean_cols, + metadata={"deduped": True}, + ) + return DAGValue(data=plan, format=DAGFormat.TRINO) + + +class TrinoAggregationNode(DAGNode): + """Appends a Trino SQL aggregation CTE computing group by and windowed aggregates.""" + + def __init__( + self, + name: str, + aggregations: Sequence[Aggregation], + group_by_keys: Sequence[str], + timestamp_col: str, + client: Trino, + inputs: Optional[Sequence[DAGNode]] = None, + enable_tiling: bool = False, + hop_size: Optional[timedelta] = None, + ): + super().__init__(name, inputs=list(inputs) if inputs else None) + self.aggregations: Tuple[Aggregation, ...] = tuple(aggregations) + self.group_by_keys: Tuple[str, ...] = tuple(group_by_keys) + self.timestamp_col: str = timestamp_col + self.client: Trino = client + self.enable_tiling: bool = enable_tiling + self.hop_size: Optional[timedelta] = hop_size + + def execute(self, context: ExecutionContext) -> DAGValue: + input_value = self.get_single_input_value(context) + input_value.assert_format(DAGFormat.TRINO) + upstream_plan: TrinoQueryPlan = input_value.data + + upstream_cte = upstream_plan.get_latest_cte_name() + + # Pure list comprehension for aggregation expressions + agg_exprs = [_format_agg_expression(agg) for agg in self.aggregations] + group_by_clauses = [quote_identifier(k) for k in self.group_by_keys] + + has_time_windows = any(agg.time_window for agg in self.aggregations) + + if has_time_windows: + time_window = self.aggregations[0].time_window + assert time_window is not None + window_sec = int(time_window.total_seconds()) + + # Trino epoch bucketing expression + window_bucket_expr = ( + f"from_unixtime(floor(to_unixtime({quote_identifier(self.timestamp_col)}) / " + f"{window_sec}) * {window_sec})" + ) + projections = ( + *group_by_clauses, + f"{window_bucket_expr} AS {quote_identifier(self.timestamp_col)}", + *agg_exprs, + ) + group_by_all = (*group_by_clauses, window_bucket_expr) + else: + projections = ( + *group_by_clauses, + f"MAX({quote_identifier(self.timestamp_col)}) AS {quote_identifier(self.timestamp_col)}", + *agg_exprs, + ) + group_by_all = tuple(group_by_clauses) + + agg_cols = ( + *self.group_by_keys, + self.timestamp_col, + *( + agg.resolved_name(agg.time_window) + if has_time_windows + else agg.resolved_name() + for agg in self.aggregations + ), + ) + + cte_name = f"_agg_{self.name.replace(':', '_')}" + query = ( + f"SELECT {', '.join(projections)}\n" + f"FROM {upstream_cte}\n" + f"GROUP BY {', '.join(group_by_all)}" + ) + + plan = upstream_plan.add_cte( + name=cte_name, + query=query, + columns=agg_cols, + metadata={"aggregated": True}, + ) + return DAGValue(data=plan, format=DAGFormat.TRINO) + + +class TrinoJoinNode(DAGNode): + """Joins upstream Trino SQL query plans or joins with entity_df.""" + + def __init__( + self, + name: str, + column_info: ColumnInfo, + client: Trino, + inputs: Optional[Sequence[DAGNode]] = None, + how: str = "left", + ): + super().__init__(name, inputs=list(inputs) if inputs else []) + self.column_info = column_info + self.client = client + self.how = how.upper() + + def execute(self, context: ExecutionContext) -> DAGValue: + input_values = self.get_input_values(context) + for val in input_values: + val.assert_format(DAGFormat.TRINO) + + upstream_plans: Tuple[TrinoQueryPlan, ...] = tuple( + val.data for val in input_values + ) + join_keys = self.column_info.join_keys + + # Pure tuple comprehension across input plans + combined_ctes = tuple(cte for plan in upstream_plans for cte in plan.ctes) + + base_plan = upstream_plans[0] + curr_from = base_plan.get_latest_cte_name() + + # Build join steps + join_steps = [] + for plan in upstream_plans[1:]: + next_cte = plan.get_latest_cte_name() + if join_keys: + on_conditions = " AND ".join( + f"{curr_from}.{quote_identifier(k)} = {next_cte}.{quote_identifier(k)}" + for k in join_keys + ) + join_steps.append(f"{self.how} JOIN {next_cte} ON {on_conditions}") + else: + join_steps.append(f"CROSS JOIN {next_cte}") + + cte_name = f"_join_{self.name.replace(':', '_')}" + join_suffix = f" {' '.join(join_steps)}" if join_steps else "" + query = f"SELECT * FROM {curr_from}{join_suffix}" + + upstream_cols = unique_ordered( + c for p in upstream_plans if p.columns for c in p.columns + ) + + final_plan = TrinoQueryPlan( + ctes=(*combined_ctes, (cte_name, query)), + current_from=cte_name, + columns=upstream_cols or None, + join_keys=tuple(join_keys), + timestamp_col=base_plan.timestamp_col, + created_timestamp_col=base_plan.created_timestamp_col, + metadata={"joined": True}, + ) + + # Handle entity_df if provided in context + if context.entity_df is not None: + final_plan = self._join_with_entity_df(final_plan, context) + + return DAGValue(data=final_plan, format=DAGFormat.TRINO) + + def _join_with_entity_df( + self, plan: TrinoQueryPlan, context: ExecutionContext + ) -> TrinoQueryPlan: + """Upload entity_df to staging table and join with point-in-time correctness.""" + entity_df = context.entity_df + assert entity_df is not None + + if isinstance(entity_df, pd.DataFrame): + catalog = getattr(context.repo_config.batch_engine, "catalog", "memory") + dataset = getattr(context.repo_config.batch_engine, "dataset", "feast") + tbl_name = get_temp_entity_table_name() + full_tbl = f"{catalog}.{dataset}.{tbl_name}" + + entity_ts_col = infer_entity_timestamp_column( + dict(zip(entity_df.columns, entity_df.dtypes)) + ) + # Rename entity timestamp column to ENTITY_TS_ALIAS to prevent collisions + # and provide explicit binding for point-in-time filtering. + if entity_ts_col and entity_ts_col != ENTITY_TS_ALIAS: + upload_df = entity_df.rename(columns={entity_ts_col: ENTITY_TS_ALIAS}) + else: + upload_df = entity_df + + # Upload entity dataframe to Trino + connector = getattr( + context.repo_config.offline_store, + "connector", + {"type": "memory"}, + ) + if not isinstance(connector, dict) or "type" not in connector: + connector = {"type": "memory"} + upload_pandas_dataframe_to_trino( + client=self.client, + df=upload_df, + table=full_tbl, + connector_args=connector, + ) + + entity_source_ref = full_tbl + else: + entity_source_ref = f"({entity_df})" + entity_ts_col = ENTITY_TS_ALIAS + upload_df = None + + latest_cte = plan.get_latest_cte_name() + join_keys = plan.join_keys or self.column_info.join_keys + + on_cond = [ + f"{latest_cte}.{quote_identifier(k)} = _entity.{quote_identifier(k)}" + for k in join_keys + ] + if plan.timestamp_col: + on_cond.append( + f"{latest_cte}.{quote_identifier(plan.timestamp_col)} <= _entity.{quote_identifier(ENTITY_TS_ALIAS)}" + ) + + if on_cond: + join_clause = f"LEFT JOIN {latest_cte} ON {' AND '.join(on_cond)}" + else: + join_clause = f"CROSS JOIN {latest_cte}" + + entity_cols = ( + tuple(upload_df.columns) if isinstance(upload_df, pd.DataFrame) else () + ) + candidate_cols = plan.columns or tuple(self.column_info.feature_cols or ()) + rhs_cols = tuple( + c for c in candidate_cols if c not in entity_cols and c not in join_keys + ) + + if rhs_cols: + rhs_select = ", ".join( + f"{latest_cte}.{quote_identifier(c)}" for c in rhs_cols + ) + select_clause = f"_entity.*, {rhs_select}" + else: + select_clause = f"_entity.*, {latest_cte}.*" + + cte_name = f"_pit_join_{self.name.replace(':', '_')}" + query = ( + f"SELECT {select_clause}\n" + f"FROM {entity_source_ref} AS _entity\n" + f"{join_clause}" + ) + + joined_cols = (*entity_cols, *rhs_cols) if entity_cols else plan.columns + + return plan.add_cte( + name=cte_name, + query=query, + columns=joined_cols, + metadata={"entity_joined": True, "entity_ts_col": entity_ts_col}, + ) + + +class TrinoTransformationNode(DAGNode): + """Appends a Trino SQL CTE applying user-defined SQL feature transformations.""" + + def __init__( + self, + name: str, + udf: Any, + inputs: Sequence[DAGNode], + udf_string: str = "", + ): + super().__init__(name, list(inputs)) + self.udf = udf + self.udf_string = udf_string + + def execute(self, context: ExecutionContext) -> DAGValue: + input_values = self.get_input_values(context) + for val in input_values: + val.assert_format(DAGFormat.TRINO) + + upstream_plans: Tuple[TrinoQueryPlan, ...] = tuple( + val.data for val in input_values + ) + primary_plan = upstream_plans[0] + input_table_names = tuple(p.get_latest_cte_name() for p in upstream_plans) + + # Resolve SQL transformation string + if callable(self.udf): + sql_expr = self.udf(*input_table_names) + elif isinstance(self.udf, str): + sql_expr = ( + self.udf.format(*input_table_names) + if "{}" in self.udf + else f"SELECT * FROM {input_table_names[0]}" + ) + elif self.udf_string: + sql_expr = ( + self.udf_string.format(*input_table_names) + if "{}" in self.udf_string + else f"SELECT * FROM {input_table_names[0]}" + ) + else: + raise ValueError( + f"TrinoTransformationNode '{self.name}' requires SQL UDF string or callable." + ) + + cte_name = f"_transform_{self.name.replace(':', '_')}" + plan = primary_plan.add_cte( + name=cte_name, query=sql_expr, metadata={"transformed": True} + ) + return DAGValue(data=plan, format=DAGFormat.TRINO) + + +class TrinoValidationNode(DAGNode): + """Validates schema projections and data invariants in Trino SQL.""" + + def __init__( + self, + name: str, + expected_columns: Mapping[str, Optional[str]], + json_columns: Optional[Set[str]] = None, + inputs: Optional[Sequence[DAGNode]] = None, + ): + super().__init__(name, inputs=list(inputs) if inputs else None) + self.expected_columns: Mapping[str, Optional[str]] = dict(expected_columns) + self.json_columns: frozenset[str] = frozenset(json_columns or ()) + + def execute(self, context: ExecutionContext) -> DAGValue: + input_value = self.get_single_input_value(context) + input_value.assert_format(DAGFormat.TRINO) + context.node_outputs[self.name] = input_value + return input_value + + +class TrinoWriteNode(DAGNode): + """Terminal node executing Trino SQL plan and persisting to online or offline store. + + - Online Store: Streams PyArrow RecordBatches directly to online_write_batch. + - Offline Store: Generates in-cluster CTAS / INSERT INTO via atomic staging + swap. + """ + + def __init__( + self, + name: str, + feature_view: Union[BatchFeatureView, StreamFeatureView, FeatureView], + client: Trino, + inputs: Optional[Sequence[DAGNode]] = None, + ): + super().__init__(name, inputs=list(inputs) if inputs else None) + self.feature_view = feature_view + self.client = client + + def execute(self, context: ExecutionContext) -> DAGValue: + input_value = self.get_single_input_value(context) + input_value.assert_format(DAGFormat.TRINO) + plan: TrinoQueryPlan = input_value.data + full_sql = plan.to_sql() + + # 1. Write to online store via streaming PyArrow batches + if self.feature_view.online: + logger.info("Executing online store materialization stream via Trino...") + batches = stream_trino_arrow_batches( + client=self.client, + query_text=full_sql, + batch_size=getattr( + context.repo_config.batch_engine, "batch_size", 10000 + ), + ) + write_arrow_batches_to_online_store( + batches=batches, + feature_view=self.feature_view, + online_store=context.online_store, + repo_config=context.repo_config, + concurrency=getattr( + context.repo_config.batch_engine, "write_concurrency", 4 + ), + ) + + # 2. Write to offline store using atomic staging + rename-swap pattern + if getattr(self.feature_view, "offline", False): + self._execute_offline_write(plan, context) + + return DAGValue( + data=plan, + format=DAGFormat.TRINO, + metadata={ + "feature_view": self.feature_view.name, + "online": self.feature_view.online, + "offline": getattr(self.feature_view, "offline", False), + }, + ) + + def _execute_offline_write( + self, plan: TrinoQueryPlan, context: ExecutionContext + ) -> None: + """Execute materialization to an offline table. + + - 'append' (default): Appends new rows into the existing target table via + INSERT INTO, or creates it via CTAS if it does not yet exist. + - 'merge': Idempotently upserts rows using MERGE INTO (Iceberg/Delta Lake), + matching on (join_keys, timestamp_col). + - 'overwrite': Atomically replaces the target table using a safe 3-step + staging + backup-swap pattern (staging -> backup -> target) with rollback. + """ + target_source = getattr(self.feature_view, "batch_source", None) + if target_source is None: + return + + target_table = target_source.get_table_query_string() + repo_config = getattr(context, "repo_config", None) + batch_engine = getattr(repo_config, "batch_engine", None) + write_mode = getattr(batch_engine, "offline_write_mode", "append") + + match write_mode: + case "merge": + self._execute_offline_merge(plan, target_table) + case "overwrite": + self._execute_offline_overwrite(plan, target_table) + case _: + self._execute_offline_append(plan, target_table) + + def _execute_offline_merge(self, plan: TrinoQueryPlan, target_table: str) -> None: + """Upsert features using Trino MERGE INTO (Iceberg/Delta Lake). + + Updates matching (join_keys, timestamp_col) records in place and inserts + new records. If the target table does not yet exist, creates it via CTAS. + """ + if not self._table_exists(target_table): + create_sql = f"CREATE TABLE {target_table} AS (\n{plan.to_sql()}\n)" + self.client.execute_query(create_sql) + return + + join_keys = plan.join_keys or () + ts_col = plan.timestamp_col + all_cols = plan.columns or () + + if not join_keys or not ts_col: + self._execute_offline_append(plan, target_table) + return + + match_keys = (*join_keys, ts_col) + update_cols = tuple(c for c in all_cols if c not in match_keys) + + on_clause = " AND ".join( + f"target.{quote_identifier(k)} = stage.{quote_identifier(k)}" + for k in match_keys + ) + matched_clause = ( + f"WHEN MATCHED THEN\n UPDATE SET {', '.join(f'{quote_identifier(c)} = stage.{quote_identifier(c)}' for c in update_cols)}\n" + if update_cols + else "" + ) + insert_cols = ", ".join(quote_identifier(c) for c in all_cols) + insert_vals = ", ".join(f"stage.{quote_identifier(c)}" for c in all_cols) + + merge_sql = ( + f"MERGE INTO {target_table} AS target\n" + f"USING (\n{plan.to_sql()}\n) AS stage\n" + f"ON {on_clause}\n" + f"{matched_clause}" + f"WHEN NOT MATCHED THEN\n" + f" INSERT ({insert_cols})\n" + f" VALUES ({insert_vals})" + ) + self.client.execute_query(merge_sql) + + def _execute_offline_append(self, plan: TrinoQueryPlan, target_table: str) -> None: + """Append features to offline target table using INSERT INTO, or CTAS if new.""" + if self._table_exists(target_table): + col_clause = ( + f" ({', '.join(quote_identifier(c) for c in plan.columns)})" + if plan.columns + else "" + ) + insert_sql = f"INSERT INTO {target_table}{col_clause}\n{plan.to_sql()}" + self.client.execute_query(insert_sql) + else: + create_sql = f"CREATE TABLE {target_table} AS (\n{plan.to_sql()}\n)" + self.client.execute_query(create_sql) + + def _execute_offline_overwrite( + self, plan: TrinoQueryPlan, target_table: str + ) -> None: + """Safely overwrite target table: staging -> backup -> target -> drop backup.""" + target_prefix = ( + target_table.rsplit(".", 1)[0] + "." if "." in target_table else "" + ) + target_simple = target_table.rsplit(".", 1)[-1] + unquoted_simple = target_simple.strip('"') + + staging_simple = quote_identifier(f"{unquoted_simple}__staging") + backup_simple = quote_identifier(f"{unquoted_simple}__backup") + + staging_table = f"{target_prefix}{staging_simple}" + backup_table = f"{target_prefix}{backup_simple}" + + # Clean up any leftover staging table + try: + self.client.execute_query(f"DROP TABLE IF EXISTS {staging_table}") + except TrinoQueryError as e: + logger.debug("Failed to drop staging table %s: %s", staging_table, e) + + # Step 1: Create staging table from compiled plan + create_sql = f"CREATE TABLE {staging_table} AS (\n{plan.to_sql()}\n)" + self.client.execute_query(create_sql) + + target_backed_up = False + try: + # Step 2: If target exists, rename target -> backup + if self._table_exists(target_table): + try: + self.client.execute_query(f"DROP TABLE IF EXISTS {backup_table}") + except TrinoQueryError: + pass + self.client.execute_query( + f"ALTER TABLE {target_table} RENAME TO {backup_simple}" + ) + target_backed_up = True + + # Step 3: Rename staging -> target + self.client.execute_query( + f"ALTER TABLE {staging_table} RENAME TO {target_simple}" + ) + + # Step 4: Drop backup now that target is safely replaced + if target_backed_up: + try: + self.client.execute_query(f"DROP TABLE IF EXISTS {backup_table}") + except TrinoQueryError as e: + logger.debug("Failed to drop backup table %s: %s", backup_table, e) + + except Exception: + # Rollback: restore backup -> target if backup was created + if target_backed_up: + try: + self.client.execute_query( + f"ALTER TABLE {backup_table} RENAME TO {target_simple}" + ) + except TrinoQueryError as rollback_err: + logger.error( + "Failed to restore backup table %s to %s: %s", + backup_table, + target_table, + rollback_err, + ) + # Clean up staging table + try: + self.client.execute_query(f"DROP TABLE IF EXISTS {staging_table}") + except TrinoQueryError: + pass + raise + + def _table_exists(self, table_name: str) -> bool: + """Check whether a Trino table exists.""" + try: + self.client.execute_query(f"SELECT 1 FROM {table_name} WHERE 1 = 0") + return True + except Exception: + return False diff --git a/sdk/python/feast/infra/compute_engines/trino/sql_builder.py b/sdk/python/feast/infra/compute_engines/trino/sql_builder.py new file mode 100644 index 00000000000..0bed08dbcfa --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/trino/sql_builder.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Mapping, Optional, Sequence, Tuple + +from feast.infra.compute_engines.trino.utils import quote_identifier +from feast.infra.compute_engines.utils import ENTITY_TS_ALIAS + + +@dataclass(frozen=True) +class TrinoQueryPlan: + """Represents an immutable, compiled Trino SQL query plan built from DAG nodes. + + Maintains Common Table Expressions (CTEs), current projection columns, + and metadata (join keys, timestamp column) to support in-cluster execution. + Strictly immutable by design via copy-on-write transformations. + """ + + ctes: Tuple[Tuple[str, str], ...] = () + current_from: str = "" + columns: Optional[Tuple[str, ...]] = None + join_keys: Tuple[str, ...] = () + timestamp_col: Optional[str] = None + created_timestamp_col: Optional[str] = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "ctes", tuple(self.ctes)) + if self.columns is not None and not isinstance(self.columns, tuple): + object.__setattr__(self, "columns", tuple(self.columns)) + if not isinstance(self.join_keys, tuple): + object.__setattr__(self, "join_keys", tuple(self.join_keys)) + if not isinstance(self.metadata, dict): + object.__setattr__(self, "metadata", dict(self.metadata)) + + def add_cte( + self, + name: str, + query: str, + columns: Optional[Sequence[str]] = None, + timestamp_col: Optional[str] = None, + created_timestamp_col: Optional[str] = None, + join_keys: Optional[Sequence[str]] = None, + metadata: Optional[Mapping[str, Any]] = None, + ) -> TrinoQueryPlan: + """Append a CTE step to the query plan immutably and update the current reference.""" + return TrinoQueryPlan( + ctes=(*self.ctes, (name, query.strip())), + current_from=name, + columns=tuple(columns) if columns is not None else self.columns, + join_keys=tuple(join_keys) if join_keys is not None else self.join_keys, + timestamp_col=timestamp_col or self.timestamp_col, + created_timestamp_col=created_timestamp_col or self.created_timestamp_col, + metadata={**self.metadata, **(metadata or {})}, + ) + + def to_sql(self) -> str: + """Compile the CTE chain into a single executable ANSI SQL statement.""" + entity_ts_col = self.metadata.get("entity_ts_col") + if self.columns: + if ( + entity_ts_col + and entity_ts_col != ENTITY_TS_ALIAS + and ENTITY_TS_ALIAS in self.columns + ): + projected = [ + f"{quote_identifier(c)} AS {quote_identifier(entity_ts_col)}" + if c == ENTITY_TS_ALIAS + else c + for c in self.columns + if c != entity_ts_col + ] + cols = ", ".join(projected) + else: + cols = ", ".join(self.columns) + else: + cols = "*" + + if not self.ctes: + if self.current_from: + return f"SELECT {cols} FROM {self.current_from}" + return "SELECT 1" + + cte_parts = [f"{name} AS (\n{query}\n)" for name, query in self.ctes] + with_clause = "WITH\n" + ",\n".join(cte_parts) + return f"{with_clause}\nSELECT {cols}\nFROM {self.current_from}" + + def get_latest_cte_name(self) -> str: + """Return the identifier of the latest CTE in the plan.""" + return self.current_from diff --git a/sdk/python/feast/infra/compute_engines/trino/utils.py b/sdk/python/feast/infra/compute_engines/trino/utils.py new file mode 100644 index 00000000000..ad6d0bf18b1 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/trino/utils.py @@ -0,0 +1,330 @@ +from __future__ import annotations + +import concurrent.futures +import logging +from datetime import date, datetime, timezone +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Iterable, + Iterator, + Optional, + TypeVar, +) + +import numpy as np +import pandas as pd +import pyarrow as pa +from trino.exceptions import TrinoConnectionError, TrinoQueryError + +from feast.infra.offline_stores.contrib.trino_offline_store.trino_queries import ( + Trino, +) +from feast.infra.offline_stores.contrib.trino_offline_store.trino_type_map import ( + trino_to_pa_value_type, +) +from feast.types import ( + Array, + PrimitiveFeastType, + Set, + Struct, +) +from feast.utils import _convert_arrow_to_proto, _run_pyarrow_field_mapping + +if TYPE_CHECKING: + from feast.feature_view import FeatureView + from feast.infra.online_stores.online_store import OnlineStore + from feast.repo_config import RepoConfig + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + + +def unique_ordered(items: Iterable[T]) -> tuple[T, ...]: + """Return unique elements from an iterable, preserving first-seen insertion order.""" + return tuple(dict.fromkeys(items)) + + +def quote_identifier(name: str) -> str: + """Safely quote a Trino SQL identifier (table, column, or CTE alias).""" + escaped = name.replace('"', '""') + return f'"{escaped}"' + + +def _has_unquoted_semicolon(sql: str) -> bool: + """Return True if sql contains a semicolon outside of quotes.""" + in_quote = False + quote_char: Optional[str] = None + is_escaped = False + + for char in sql: + if is_escaped: + is_escaped = False + continue + + if char == "\\": + is_escaped = True + continue + + if char in ("'", '"'): + if not in_quote: + in_quote = True + quote_char = char + elif char == quote_char: + in_quote = False + quote_char = None + elif char == ";" and not in_quote: + return True + + return False + + +def from_feast_to_trino_type(feast_type: Any) -> Optional[str]: + """Convert a Feast data type to an explicit Trino SQL type string. + + Never infers NULL; returns strongly-typed Trino SQL types. + """ + if isinstance(feast_type, Struct): + inner_fields = [] + for name, ftype in feast_type.fields.items(): + trino_type = from_feast_to_trino_type(ftype) + if trino_type is None: + return None + inner_fields.append(f"{quote_identifier(name)} {trino_type}") + return f"ROW({', '.join(inner_fields)})" + + if isinstance(feast_type, PrimitiveFeastType): + mapping = { + PrimitiveFeastType.BYTES: "VARBINARY", + PrimitiveFeastType.STRING: "VARCHAR", + PrimitiveFeastType.INT32: "INTEGER", + PrimitiveFeastType.INT64: "BIGINT", + PrimitiveFeastType.FLOAT64: "DOUBLE", + PrimitiveFeastType.FLOAT32: "REAL", + PrimitiveFeastType.BOOL: "BOOLEAN", + PrimitiveFeastType.UNIX_TIMESTAMP: "TIMESTAMP", + PrimitiveFeastType.MAP: "MAP(VARCHAR, VARCHAR)", + PrimitiveFeastType.JSON: "VARCHAR", + } + return mapping.get(feast_type) + + if isinstance(feast_type, Array): + base_type = feast_type.base_type + if isinstance(base_type, Struct): + inner = from_feast_to_trino_type(base_type) + return f"ARRAY({inner})" if inner else None + if isinstance(base_type, PrimitiveFeastType): + if base_type == PrimitiveFeastType.MAP: + return "ARRAY(MAP(VARCHAR, VARCHAR))" + inner = from_feast_to_trino_type(base_type) + return f"ARRAY({inner})" if inner else None + + if isinstance(feast_type, Set): + inner = from_feast_to_trino_type(feast_type.base_type) + return f"ARRAY({inner})" if inner else None + + return None + + +def _is_nan(value: Any) -> bool: + """Check if value is NaN or null, guarded against container types.""" + if value is None: + return True + if isinstance(value, (list, tuple, np.ndarray)): + return False + try: + return bool(pd.isna(value)) + except (TypeError, ValueError): + return False + + +def _trino_sql_literal(val: Any) -> str: + """Format Python objects into valid, strongly-typed Trino SQL literals. + + Guarantees: + - Normalizes timezone-aware datetimes and timestamps to UTC. + - Handles containers before null checks to avoid array truth-value errors. + - Escapes single quotes in strings. + - Emits TRUE/FALSE for booleans. + - Emits NULL for None, pd.NaT, and NaN. + """ + if val is None or val is pd.NaT: + return "NULL" + + # Container types must be inspected before checking pd.isna() + if isinstance(val, list): + items = [_trino_sql_literal(item) for item in val] + return f"ARRAY[{', '.join(items)}]" + if isinstance(val, tuple): + items = [_trino_sql_literal(item) for item in val] + return f"ARRAY[{', '.join(items)}]" + if isinstance(val, np.ndarray): + items = [_trino_sql_literal(item) for item in val.tolist()] + return f"ARRAY[{', '.join(items)}]" + + if _is_nan(val): + return "NULL" + + if isinstance(val, (bool, np.bool_)): + return "TRUE" if val else "FALSE" + + if isinstance(val, (int, float, np.integer, np.floating)): + return str(val) + + if isinstance(val, (datetime, pd.Timestamp)): + if val.tzinfo is not None and val.tzinfo.utcoffset(val) is not None: + val = val.astimezone(timezone.utc) + return f"TIMESTAMP '{val.strftime('%Y-%m-%d %H:%M:%S.%f')}'" + + if isinstance(val, np.datetime64): + val = pd.Timestamp(val) + if val.tzinfo is not None and val.tzinfo.utcoffset(val) is not None: + val = val.astimezone(timezone.utc) + return f"TIMESTAMP '{val.strftime('%Y-%m-%d %H:%M:%S.%f')}'" + + if isinstance(val, date): + return f"DATE '{val.strftime('%Y-%m-%d')}'" + + if isinstance(val, str): + escaped = val.replace("'", "''") + return f"'{escaped}'" + + return str(val) + + +def get_or_create_trino_client(config: Any) -> Trino: + """Instantiate a Trino client with authentication credentials unwrapped.""" + auth = None + if getattr(config, "auth", None) is not None: + auth = config.auth.to_trino_auth() + + return Trino( + host=config.host, + port=config.port, + user=config.user, + catalog=config.catalog, + source=getattr(config, "source", "trino-python-client"), + http_scheme=getattr(config, "http_scheme", "http"), + verify=getattr(config, "verify", True), + extra_credential=getattr(config, "extra_credential", None), + auth=auth, + ) + + +def stream_trino_arrow_batches( + client: Trino, + query_text: str, + batch_size: int = 10000, +) -> Iterator[pa.RecordBatch]: + """Execute a query and yield PyArrow RecordBatches in streaming chunks. + + Maintains a strictly bounded memory footprint O(batch_size) without + accumulating the entire dataset in Python memory. + """ + cursor = client.get_cursor() + try: + cursor.execute(query_text) + + if not cursor.description: + return + + arrow_schema = pa.schema( + [ + pa.field(str(col[0]), trino_to_pa_value_type(str(col[1]))) + for col in cursor.description + ] + ) + + while rows := cursor.fetchmany(batch_size): + arrays = [ + pa.array( + [row[col_idx] for row in rows], + type=arrow_schema.field(col_idx).type, + ) + for col_idx in range(len(cursor.description)) + ] + yield pa.RecordBatch.from_arrays(arrays, schema=arrow_schema) + except TrinoConnectionError: + # Never swallow connection errors; bubble up immediately to caller + logger.error("Trino connection dropped while streaming batches.") + raise + except TrinoQueryError as e: + logger.debug("Trino query failed during streaming execution: %s", e) + raise + finally: + try: + cursor.close() + except Exception: + pass + + +def write_arrow_batches_to_online_store( + batches: Iterable[pa.RecordBatch], + feature_view: "FeatureView", + online_store: "OnlineStore", + repo_config: "RepoConfig", + concurrency: int = 4, + progress_callback: Optional[Callable[[int], None]] = None, +) -> None: + """Stream PyArrow record batches to the online store concurrently via threadpool.""" + join_key_to_value_type = { + entity.name: entity.dtype.to_value_type() + for entity in feature_view.entity_columns + } + batch_chunk_size = repo_config.materialization_config.online_write_batch_size + + def _process_single_batch(batch: pa.RecordBatch) -> int: + table = pa.Table.from_batches([batch]) + if ( + hasattr(feature_view, "batch_source") + and feature_view.batch_source is not None + and getattr(feature_view.batch_source, "field_mapping", None) is not None + ): + table = _run_pyarrow_field_mapping( + table, feature_view.batch_source.field_mapping + ) + + if ( + hasattr(feature_view, "batch_source") + and feature_view.batch_source is not None + and getattr(feature_view.batch_source, "created_timestamp_column", None) + and feature_view.batch_source.created_timestamp_column + not in table.column_names + ): + created_col = feature_view.batch_source.created_timestamp_column + table = table.append_column( + created_col, + pa.nulls(table.num_rows, type=pa.timestamp("us", tz="UTC")), + ) + + sub_batches = ( + [table] + if batch_chunk_size is None + else table.to_batches(max_chunksize=batch_chunk_size) + ) + rows_written = 0 + for sub_batch in sub_batches: + rows_to_write = _convert_arrow_to_proto( + sub_batch, feature_view, join_key_to_value_type + ) + online_store.online_write_batch( + config=repo_config, + table=feature_view, + data=rows_to_write, + progress=lambda x: None, + ) + rows_written += len(rows_to_write) + return rows_written + + with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = [] + for batch in batches: + if batch.num_rows > 0: + futures.append(executor.submit(_process_single_batch, batch)) + + for future in concurrent.futures.as_completed(futures): + written = future.result() + if progress_callback: + progress_callback(written) diff --git a/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino_queries.py b/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino_queries.py index ef23a5cf021..246fea2d066 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino_queries.py +++ b/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino_queries.py @@ -52,7 +52,8 @@ def __init__( self.auth = auth self._cursor: Optional[Cursor] = None - def _get_cursor(self) -> Cursor: + def get_cursor(self) -> Cursor: + """Return an active DB-API Cursor for Trino queries.""" if self._cursor is None: headers = ( {trino.constants.HEADER_EXTRA_CREDENTIAL: self.extra_credential} @@ -73,6 +74,9 @@ def _get_cursor(self) -> Cursor: return self._cursor + def _get_cursor(self) -> Cursor: + return self.get_cursor() + def create_query(self, query_text: str) -> Query: """ Create a Query object without executing it. diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index 728ec306153..61b1502130a 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -52,6 +52,8 @@ "ray.engine": "feast.infra.compute_engines.ray.compute.RayComputeEngine", "flink.engine": "feast.infra.compute_engines.flink.compute.FlinkComputeEngine", "spark_application": "feast.infra.compute_engines.spark_application.compute.SparkApplicationComputeEngine", + "trino.engine": "feast.infra.compute_engines.trino.compute.TrinoComputeEngine", + "trino": "feast.infra.compute_engines.trino.compute.TrinoComputeEngine", } LEGACY_ONLINE_STORE_CLASS_FOR_TYPE = { diff --git a/sdk/python/feast/transformation/mode.py b/sdk/python/feast/transformation/mode.py index bd6fdf22424..9c1e0be804b 100644 --- a/sdk/python/feast/transformation/mode.py +++ b/sdk/python/feast/transformation/mode.py @@ -9,4 +9,5 @@ class TransformationMode(Enum): FLINK = "flink" RAY = "ray" SQL = "sql" + TRINO_SQL = "trino_sql" SUBSTRAIT = "substrait" diff --git a/sdk/python/feast/transformation/trino_transformation.py b/sdk/python/feast/transformation/trino_transformation.py new file mode 100644 index 00000000000..ffc57a2dcfd --- /dev/null +++ b/sdk/python/feast/transformation/trino_transformation.py @@ -0,0 +1,106 @@ +from typing import Any, Dict, Optional, Union, cast + +from feast.transformation.base import Transformation +from feast.transformation.mode import TransformationMode + + +class TrinoTransformation(Transformation): + """ + TrinoTransformation defines a feature transformation executed as Trino SQL. + + Example: + ```python + trino_transformation = TrinoTransformation( + mode=TransformationMode.TRINO_SQL, + udf="SELECT driver_id, conv_rate * 2.0 AS conv_rate, acc_rate * 2.0 AS acc_rate FROM {}", + udf_string="SELECT driver_id, conv_rate * 2.0 AS conv_rate, acc_rate * 2.0 AS acc_rate FROM {}", + ) + ``` + or with a callable returning a SQL string: + ```python + def double_rates_sql(table_name: str) -> str: + return f"SELECT driver_id, conv_rate * 2.0 AS conv_rate, acc_rate * 2.0 AS acc_rate FROM {table_name}" + + trino_transformation = TrinoTransformation( + mode=TransformationMode.TRINO_SQL, + udf=double_rates_sql, + udf_string="double_rates_sql", + ) + ``` + """ + + def __new__( + cls, + mode: Union[TransformationMode, str] = TransformationMode.TRINO_SQL, + udf: Any = None, + udf_string: str = "", + name: Optional[str] = None, + tags: Optional[Dict[str, str]] = None, + description: str = "", + owner: str = "", + *args, + **kwargs, + ) -> "TrinoTransformation": + if isinstance(mode, str): + mode = TransformationMode(mode) + + resolved_name = name or getattr(udf, "__name__", None) or "trino_transformation" + + instance = super(TrinoTransformation, cls).__new__( + cls, + mode=mode, + udf=udf, + udf_string=udf_string, + name=resolved_name, + tags=tags, + description=description, + owner=owner, + ) + return cast(TrinoTransformation, instance) + + def __init__( + self, + mode: Union[TransformationMode, str] = TransformationMode.TRINO_SQL, + udf: Any = None, + udf_string: str = "", + name: Optional[str] = None, + tags: Optional[Dict[str, str]] = None, + description: str = "", + owner: str = "", + *args, + **kwargs, + ): + if isinstance(mode, str): + mode = TransformationMode(mode) + + resolved_name = name or getattr(udf, "__name__", None) or "trino_transformation" + + super().__init__( + mode=mode, + udf=udf, + name=resolved_name, + udf_string=udf_string, + tags=tags, + description=description, + owner=owner, + ) + + def transform(self, *inputs: str) -> str: + """ + Applies the SQL transformation by injecting the input table/CTE names + into the SQL query template or calling the UDF. + """ + if callable(self.udf): + return self.udf(*inputs) + elif isinstance(self.udf, str): + if inputs and "{}" in self.udf: + return self.udf.format(*inputs) + return self.udf + elif self.udf_string: + if inputs and "{}" in self.udf_string: + return self.udf_string.format(*inputs) + return self.udf_string + raise ValueError(f"Invalid TrinoTransformation definition: {self.name}") + + def infer_features(self, *args, **kwargs) -> Any: + pass diff --git a/sdk/python/tests/component/trino/test_compute.py b/sdk/python/tests/component/trino/test_compute.py new file mode 100644 index 00000000000..d6902f6ce09 --- /dev/null +++ b/sdk/python/tests/component/trino/test_compute.py @@ -0,0 +1,246 @@ +from datetime import datetime, timedelta, timezone +from typing import cast +from unittest.mock import MagicMock + +import pandas as pd +import pytest + +from feast import BatchFeatureView, Entity, Field +from feast.aggregation import Aggregation +from feast.infra.common.materialization_job import ( + MaterializationJobStatus, + MaterializationTask, +) +from feast.infra.common.retrieval_task import HistoricalRetrievalTask +from feast.infra.compute_engines.trino.compute import ( + TrinoComputeEngine, + TrinoComputeEngineConfig, +) +from feast.infra.compute_engines.trino.job import TrinoDAGRetrievalJob +from feast.infra.offline_stores.contrib.trino_offline_store.trino_queries import ( + Results, +) +from feast.infra.offline_stores.contrib.trino_offline_store.trino_source import ( + TrinoSource, +) +from feast.transformation.mode import TransformationMode +from feast.transformation.trino_transformation import TrinoTransformation +from feast.types import Float32, Int32 +from feast.value_type import ValueType + + +@pytest.fixture +def driver_entity(): + return Entity( + name="driver_id", + value_type=ValueType.INT32, + join_keys=["driver_id"], + ) + + +@pytest.fixture +def trino_source(): + return TrinoSource( + name="driver_hourly_stats_source", + table="iceberg.feast.driver_hourly_stats", + timestamp_field="event_timestamp", + created_timestamp_column="created_timestamp", + ) + + +def test_trino_compute_engine_get_historical_features(driver_entity, trino_source): + """Verifies historical feature retrieval via Trino SQL compilation and execution.""" + tf = TrinoTransformation( + mode=TransformationMode.TRINO_SQL, + udf="SELECT driver_id, event_timestamp, conv_rate * 2.0 AS conv_rate, acc_rate * 2.0 AS acc_rate FROM {}", + udf_string="SELECT driver_id, event_timestamp, conv_rate * 2.0 AS conv_rate, acc_rate * 2.0 AS acc_rate FROM {}", + ) + + fv = BatchFeatureView( + name="driver_hourly_stats", + entities=[driver_entity], + feature_transformation=tf, + aggregations=[ + Aggregation(column="conv_rate", function="sum"), + Aggregation(column="acc_rate", function="avg"), + ], + ttl=timedelta(days=3), + schema=[ + Field(name="conv_rate", dtype=Float32), + Field(name="acc_rate", dtype=Float32), + Field(name="driver_id", dtype=Int32), + ], + online=False, + offline=False, + source=trino_source, + ) + + entity_df = pd.DataFrame( + { + "driver_id": [1001, 1002], + "event_timestamp": [ + datetime(2025, 1, 15, 12, 0, 0, tzinfo=timezone.utc), + datetime(2025, 1, 15, 12, 0, 0, tzinfo=timezone.utc), + ], + } + ) + + mock_client = MagicMock() + # Mock query result + mock_client.execute_query.return_value = Results( + data=[ + [1001, 3.1, 1.4], + [1002, 2.0, 1.0], + ], + columns=[ + {"name": "driver_id", "type": "integer"}, + {"name": "sum_conv_rate", "type": "real"}, + {"name": "avg_acc_rate", "type": "real"}, + ], + ) + + repo_config = MagicMock() + repo_config.batch_engine = TrinoComputeEngineConfig( + host="localhost", + port=8080, + catalog="iceberg", + user="test_user", + ) + repo_config.offline_store = MagicMock() + + engine = TrinoComputeEngine( + repo_config=repo_config, + offline_store=MagicMock(), + online_store=MagicMock(), + ) + engine.client = mock_client + mock_registry = MagicMock() + mock_registry.get_entity.return_value = driver_entity + + task = HistoricalRetrievalTask( + project="test_project", + entity_df=entity_df, + feature_view=fv, + full_feature_name=False, + registry=mock_registry, + ) + + job = cast( + TrinoDAGRetrievalJob, engine.get_historical_features(mock_registry, task) + ) + compiled_sql = job.to_sql() + + # Assert SQL structure includes entity_df join and point-in-time condition + assert "WITH" in compiled_sql + assert "iceberg.feast.driver_hourly_stats" in compiled_sql + assert "GROUP BY" in compiled_sql + assert "conv_rate * 2.0" in compiled_sql + assert "_entity" in compiled_sql + assert "LEFT JOIN" in compiled_sql + assert "_pit_join_driver_hourly_stats_join" in compiled_sql + assert ( + '_join_driver_hourly_stats_join."driver_id" = _entity."driver_id"' + in compiled_sql + ) + assert '<= _entity."__entity_event_timestamp"' in compiled_sql + + # Assert entity table was created and rows uploaded to Trino + executed_queries = [call[0][0] for call in mock_client.execute_query.call_args_list] + assert any( + "CREATE TABLE IF NOT EXISTS iceberg.feast.feast_entity_df_" in q + for q in executed_queries + ) + assert any( + "INSERT INTO iceberg.feast.feast_entity_df_" in q for q in executed_queries + ) + + # Assert lazy evaluation works + df_out = job.to_df() + assert df_out["driver_id"].tolist() == [1001, 1002] + assert df_out["sum_conv_rate"].tolist() == [3.1, 2.0] + + +def test_trino_compute_engine_materialize(driver_entity, trino_source): + """Verifies materialization via streaming online writes and offline staging swap.""" + fv = BatchFeatureView( + name="driver_hourly_stats", + entities=[driver_entity], + aggregations=[ + Aggregation(column="conv_rate", function="sum"), + ], + schema=[ + Field(name="sum_conv_rate", dtype=Float32), + Field(name="driver_id", dtype=Int32), + ], + online=True, + offline=True, + source=trino_source, + ) + + mock_client = MagicMock() + # Mock cursor for streaming online writes + mock_cursor = MagicMock() + mock_cursor.description = [ + ("driver_id", "integer", None, None, None, None, None), + ("event_timestamp", "timestamp", None, None, None, None, None), + ("sum_conv_rate", "real", None, None, None, None, None), + ] + # Return 1 batch of rows then empty + mock_cursor.fetchmany.side_effect = [ + [ + (1001, datetime(2025, 1, 15, 12, 0, 0), 3.1), + (1002, datetime(2025, 1, 15, 12, 0, 0), 2.0), + ], + [], + ] + mock_client.get_cursor.return_value = mock_cursor + + repo_config = MagicMock() + repo_config.batch_engine = TrinoComputeEngineConfig( + host="localhost", + port=8080, + catalog="iceberg", + user="test_user", + batch_size=1000, + write_concurrency=2, + ) + repo_config.materialization_config.online_write_batch_size = 100 + repo_config.offline_store = MagicMock() + + mock_online_store = MagicMock() + + engine = TrinoComputeEngine( + repo_config=repo_config, + offline_store=MagicMock(), + online_store=mock_online_store, + ) + engine.client = mock_client + mock_registry = MagicMock() + mock_registry.get_entity.return_value = driver_entity + + task = MaterializationTask( + project="test_project", + feature_view=fv, + start_time=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2025, 1, 15, 0, 0, 0, tzinfo=timezone.utc), + tqdm_builder=MagicMock(), + ) + + jobs = engine.materialize(mock_registry, task) + assert len(jobs) == 1 + if jobs[0].error() is not None: + raise jobs[0].error() + assert jobs[0].status() == MaterializationJobStatus.SUCCEEDED + + # Verify online store was called with batch writes + assert mock_online_store.online_write_batch.called + + # Verify offline append query was executed (default offline_write_mode='append') + executed_queries = [call[0][0] for call in mock_client.execute_query.call_args_list] + assert any( + "INSERT INTO iceberg.feast.driver_hourly_stats" in q for q in executed_queries + ) + assert not any( + "DROP TABLE IF EXISTS iceberg.feast.driver_hourly_stats" == q + for q in executed_queries + ) diff --git a/sdk/python/tests/integration/compute_engines/test_trino_compute_engine.py b/sdk/python/tests/integration/compute_engines/test_trino_compute_engine.py new file mode 100644 index 00000000000..d04bcbb857d --- /dev/null +++ b/sdk/python/tests/integration/compute_engines/test_trino_compute_engine.py @@ -0,0 +1,629 @@ +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any, Dict +from unittest.mock import MagicMock + +import pandas as pd +import pytest + +from feast import BatchFeatureView, Entity, Field +from feast.aggregation import Aggregation +from feast.infra.common.materialization_job import ( + MaterializationJobStatus, + MaterializationTask, +) +from feast.infra.common.retrieval_task import HistoricalRetrievalTask +from feast.infra.compute_engines.trino.compute import ( + TrinoComputeEngine, + TrinoComputeEngineConfig, +) +from feast.infra.compute_engines.trino.job import TrinoDAGRetrievalJob +from feast.infra.offline_stores.contrib.trino_offline_store.connectors.upload import ( + upload_pandas_dataframe_to_trino, +) +from feast.infra.offline_stores.contrib.trino_offline_store.trino import ( + TrinoOfflineStoreConfig, +) +from feast.infra.offline_stores.contrib.trino_offline_store.trino_queries import Trino +from feast.infra.offline_stores.contrib.trino_offline_store.trino_source import ( + TrinoSource, +) +from feast.infra.online_stores.sqlite import SqliteOnlineStore, SqliteOnlineStoreConfig +from feast.transformation.mode import TransformationMode +from feast.transformation.trino_transformation import TrinoTransformation +from feast.types import Float32, Int32, Int64, String +from feast.value_type import ValueType + +docker_available = False +try: + import docker + from testcontainers.core.container import DockerContainer + from testcontainers.core.waiting_utils import wait_for_logs + + try: + _docker = docker.from_env() + _docker.ping() + docker_available = True + except Exception: + pass +except ImportError: + pass + +_requires_docker = pytest.mark.skipif( + not docker_available, + reason="Docker is not available or not running. Start Docker daemon to run Trino integration tests.", +) + + +@pytest.fixture(scope="module") +def trino_container(): + """Start a real Trino container using the memory catalog for integration tests.""" + catalog_dir = ( + Path(__file__).parents[3] + / "feast" + / "infra" + / "offline_stores" + / "contrib" + / "trino_offline_store" + / "tests" + / "catalog" + ) + container = ( + DockerContainer("trinodb/trino:483") + .with_volume_mapping(str(catalog_dir.resolve()), "/etc/catalog/") + .with_volume_mapping(str(catalog_dir.resolve()), "/etc/trino/catalog/") + .with_exposed_ports("8080") + ) + container.start() + try: + wait_for_logs(container=container, predicate="SERVER STARTED", timeout=60) + yield container + finally: + container.stop() + + +@pytest.fixture(scope="module") +def trino_server_info(trino_container) -> Dict[str, Any]: + host = trino_container.get_container_host_ip() + port = int(trino_container.get_exposed_port("8080")) + return { + "host": host, + "port": port, + "catalog": "memory", + "dataset": "feast_integration", + "user": "test_user", + } + + +@pytest.fixture(scope="module") +def live_trino_client(trino_server_info) -> Trino: + client = Trino( + host=trino_server_info["host"], + port=trino_server_info["port"], + user=trino_server_info["user"], + catalog=trino_server_info["catalog"], + source="feast-integration-tests", + http_scheme="http", + verify=False, + extra_credential=None, + auth=None, + ) + client.execute_query( + f"CREATE SCHEMA IF NOT EXISTS {trino_server_info['catalog']}.{trino_server_info['dataset']}" + ) + return client + + +@pytest.fixture +def driver_entity(): + return Entity( + name="driver_id", + value_type=ValueType.INT32, + join_keys=["driver_id"], + ) + + +@_requires_docker +@pytest.mark.integration +class TestTrinoComputeEngineIntegration: + """Integration test suite executing real queries against a live Trino cluster.""" + + def test_trino_historical_retrieval_pit_e2e( + self, trino_server_info, live_trino_client, driver_entity + ): + """Validates real Trino CTE compilation, entity_df upload, and PIT join execution.""" + catalog = trino_server_info["catalog"] + dataset = trino_server_info["dataset"] + source_table = f"{catalog}.{dataset}.source_driver_stats" + + # 1. Upload source features to memory catalog + live_trino_client.execute_query(f"DROP TABLE IF EXISTS {source_table}") + source_df = pd.DataFrame( + { + "driver_id": [1001, 1001, 1002], + "event_timestamp": [ + datetime(2025, 1, 10, 10, 0, 0, tzinfo=timezone.utc), + datetime(2025, 1, 14, 10, 0, 0, tzinfo=timezone.utc), + datetime(2025, 1, 12, 10, 0, 0, tzinfo=timezone.utc), + ], + "created_timestamp": [ + datetime(2025, 1, 10, 10, 0, 0, tzinfo=timezone.utc), + datetime(2025, 1, 14, 10, 0, 0, tzinfo=timezone.utc), + datetime(2025, 1, 12, 10, 0, 0, tzinfo=timezone.utc), + ], + "conv_rate": [0.2, 0.8, 0.5], + "acc_rate": [0.3, 0.9, 0.6], + } + ) + upload_pandas_dataframe_to_trino( + client=live_trino_client, + df=source_df, + table=source_table, + connector_args={"type": "memory"}, + ) + + trino_source = TrinoSource( + name="source_driver_stats", + table=source_table, + timestamp_field="event_timestamp", + created_timestamp_column="created_timestamp", + ) + + tf = TrinoTransformation( + mode=TransformationMode.TRINO_SQL, + udf="SELECT driver_id, event_timestamp, created_timestamp, conv_rate * 2.0 AS conv_rate, acc_rate FROM {}", + udf_string="SELECT driver_id, event_timestamp, created_timestamp, conv_rate * 2.0 AS conv_rate, acc_rate FROM {}", + ) + + fv = BatchFeatureView( + name="driver_hourly_stats", + entities=[driver_entity], + feature_transformation=tf, + aggregations=[ + Aggregation(column="conv_rate", function="sum"), + Aggregation(column="acc_rate", function="avg"), + ], + ttl=timedelta(days=3), + schema=[ + Field(name="conv_rate", dtype=Float32), + Field(name="acc_rate", dtype=Float32), + Field(name="driver_id", dtype=Int32), + ], + online=False, + offline=False, + source=trino_source, + ) + + # 2. Entity dataframe requesting point-in-time features + entity_df = pd.DataFrame( + { + "driver_id": [1001, 1002], + "event_timestamp": [ + datetime(2025, 1, 15, 12, 0, 0, tzinfo=timezone.utc), + datetime(2025, 1, 13, 12, 0, 0, tzinfo=timezone.utc), + ], + } + ) + + repo_config = MagicMock() + repo_config.batch_engine = TrinoComputeEngineConfig( + host=trino_server_info["host"], + port=trino_server_info["port"], + catalog=catalog, + dataset=dataset, + user=trino_server_info["user"], + connector={"type": "memory"}, + ) + repo_config.offline_store = TrinoOfflineStoreConfig( + host=trino_server_info["host"], + port=trino_server_info["port"], + catalog=catalog, + dataset=dataset, + user=trino_server_info["user"], + connector={"type": "memory"}, + ) + + engine = TrinoComputeEngine( + repo_config=repo_config, + offline_store=MagicMock(), + online_store=MagicMock(), + ) + + mock_registry = MagicMock() + mock_registry.get_entity.return_value = driver_entity + + task = HistoricalRetrievalTask( + project="test_project", + entity_df=entity_df, + feature_view=fv, + full_feature_name=False, + registry=mock_registry, + ) + + job = engine.get_historical_features(mock_registry, task) + assert isinstance(job, TrinoDAGRetrievalJob) + + # 3. Execute live query on Trino and verify returned PIT data + result_df = job.to_df() + assert len(result_df) == 2 + assert "driver_id" in result_df.columns + assert "sum_conv_rate" in result_df.columns + assert "avg_acc_rate" in result_df.columns + + row_1001 = result_df[result_df["driver_id"] == 1001].iloc[0] + # driver 1001 on 2025-01-15: row on 2025-01-14 (0.8 * 2 = 1.6, acc_rate = 0.9) within 3-day TTL, + # row on 2025-01-10 (5 days earlier) excluded by TTL + assert pytest.approx(row_1001["sum_conv_rate"], rel=1e-3) == 1.6 + assert pytest.approx(row_1001["avg_acc_rate"], rel=1e-3) == 0.9 + + row_1002 = result_df[result_df["driver_id"] == 1002].iloc[0] + # driver 1002 on 2025-01-13: row on 2025-01-12 (0.5 * 2 = 1.0, acc_rate = 0.6) within 3-day TTL + assert pytest.approx(row_1002["sum_conv_rate"], rel=1e-3) == 1.0 + assert pytest.approx(row_1002["avg_acc_rate"], rel=1e-3) == 0.6 + + def test_trino_materialize_online_streaming_e2e( + self, trino_server_info, live_trino_client, driver_entity, tmp_path + ): + """Validates real cursor fetchmany streaming and batch writes to online store.""" + catalog = trino_server_info["catalog"] + dataset = trino_server_info["dataset"] + source_table = f"{catalog}.{dataset}.stream_driver_stats" + + # 1. Upload source features + live_trino_client.execute_query(f"DROP TABLE IF EXISTS {source_table}") + source_df = pd.DataFrame( + { + "driver_id": [101, 102, 103, 104, 105], + "event_timestamp": [ + datetime(2025, 1, 10, 12, 0, 0, tzinfo=timezone.utc) + ] + * 5, + "created_timestamp": [ + datetime(2025, 1, 10, 12, 0, 0, tzinfo=timezone.utc) + ] + * 5, + "conv_rate": [0.1, 0.2, 0.3, 0.4, 0.5], + } + ) + upload_pandas_dataframe_to_trino( + client=live_trino_client, + df=source_df, + table=source_table, + connector_args={"type": "memory"}, + ) + + trino_source = TrinoSource( + name="stream_source", + table=source_table, + timestamp_field="event_timestamp", + created_timestamp_column="created_timestamp", + ) + + fv = BatchFeatureView( + name="stream_fv", + entities=[driver_entity], + schema=[ + Field(name="conv_rate", dtype=Float32), + Field(name="driver_id", dtype=Int32), + ], + online=True, + offline=False, + source=trino_source, + ) + + # 2. Setup real SqliteOnlineStore + sqlite_db_path = str(tmp_path / "online.db") + online_store_cfg = SqliteOnlineStoreConfig(path=sqlite_db_path) + online_store = SqliteOnlineStore() + + repo_config = MagicMock() + repo_config.project = "test_project" + repo_config.repo_path = str(tmp_path) + repo_config.online_store = online_store_cfg + repo_config.entity_key_serialization_version = 3 + repo_config.registry.enable_online_feature_view_versioning = False + online_store.update( + config=repo_config, + tables_to_delete=[], + tables_to_keep=[fv], + entities_to_delete=[], + entities_to_keep=[driver_entity], + partial=False, + ) + repo_config.batch_engine = TrinoComputeEngineConfig( + host=trino_server_info["host"], + port=trino_server_info["port"], + catalog=catalog, + dataset=dataset, + user=trino_server_info["user"], + batch_size=2, # Force multiple fetchmany chunk iterations + write_concurrency=1, + connector={"type": "memory"}, + ) + repo_config.materialization_config.online_write_batch_size = 2 + repo_config.offline_store = TrinoOfflineStoreConfig( + host=trino_server_info["host"], + port=trino_server_info["port"], + catalog=catalog, + dataset=dataset, + user=trino_server_info["user"], + connector={"type": "memory"}, + ) + + engine = TrinoComputeEngine( + repo_config=repo_config, + offline_store=MagicMock(), + online_store=online_store, + ) + + mock_registry = MagicMock() + mock_registry.get_entity.return_value = driver_entity + + task = MaterializationTask( + project="test_project", + feature_view=fv, + start_time=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2025, 1, 15, 0, 0, 0, tzinfo=timezone.utc), + tqdm_builder=MagicMock(), + ) + + # 3. Execute materialization and verify status + jobs = engine.materialize(mock_registry, task) + assert len(jobs) == 1 + if jobs[0].error() is not None: + raise jobs[0].error() + assert jobs[0].status() == MaterializationJobStatus.SUCCEEDED + + from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto + from feast.protos.feast.types.Value_pb2 import Value as ValueProto + + k101 = EntityKeyProto( + join_keys=["driver_id"], entity_values=[ValueProto(int32_val=101)] + ) + k105 = EntityKeyProto( + join_keys=["driver_id"], entity_values=[ValueProto(int32_val=105)] + ) + + online_features = online_store.online_read( + config=repo_config, + table=fv, + entity_keys=[k101, k105], + requested_features=["conv_rate"], + ) + assert len(online_features) == 2 + ts_101, feats_101 = online_features[0] + ts_105, feats_105 = online_features[1] + assert feats_101 is not None and feats_105 is not None + assert pytest.approx(feats_101["conv_rate"].float_val, rel=1e-3) == 0.1 + assert pytest.approx(feats_105["conv_rate"].float_val, rel=1e-3) == 0.5 + + def test_trino_materialize_offline_append_e2e( + self, trino_server_info, live_trino_client, driver_entity, tmp_path + ): + """Validates real Trino offline write in 'append' mode.""" + catalog = trino_server_info["catalog"] + dataset = trino_server_info["dataset"] + source_table = f"{catalog}.{dataset}.offline_src_stats" + target_table = f"{catalog}.{dataset}.offline_target_stats" + + # 1. Clean existing tables + live_trino_client.execute_query(f"DROP TABLE IF EXISTS {source_table}") + live_trino_client.execute_query(f"DROP TABLE IF EXISTS {target_table}") + + # 2. Upload source rows + source_df = pd.DataFrame( + { + "driver_id": [201, 202], + "event_timestamp": [ + datetime(2025, 1, 10, 12, 0, 0, tzinfo=timezone.utc) + ] + * 2, + "created_timestamp": [ + datetime(2025, 1, 10, 12, 0, 0, tzinfo=timezone.utc) + ] + * 2, + "conv_rate": [0.4, 0.7], + } + ) + upload_pandas_dataframe_to_trino( + client=live_trino_client, + df=source_df, + table=source_table, + connector_args={"type": "memory"}, + ) + + trino_source = TrinoSource( + name="offline_src", + table=source_table, + timestamp_field="event_timestamp", + created_timestamp_column="created_timestamp", + ) + + target_source = TrinoSource( + name="offline_target", + table=target_table, + timestamp_field="event_timestamp", + created_timestamp_column="created_timestamp", + ) + + source_fv = BatchFeatureView( + name="source_fv", + entities=[driver_entity], + schema=[ + Field(name="conv_rate", dtype=Float32), + Field(name="driver_id", dtype=Int32), + ], + source=trino_source, + ) + + fv = BatchFeatureView( + name="offline_fv", + entities=[driver_entity], + schema=[ + Field(name="conv_rate", dtype=Float32), + Field(name="driver_id", dtype=Int32), + ], + online=False, + offline=True, + source=[source_fv], + sink_source=target_source, + ) + + repo_config = MagicMock() + repo_config.repo_path = str(tmp_path) + repo_config.batch_engine = TrinoComputeEngineConfig( + host=trino_server_info["host"], + port=trino_server_info["port"], + catalog=catalog, + dataset=dataset, + user=trino_server_info["user"], + offline_write_mode="append", + connector={"type": "memory"}, + ) + repo_config.offline_store = TrinoOfflineStoreConfig( + host=trino_server_info["host"], + port=trino_server_info["port"], + catalog=catalog, + dataset=dataset, + user=trino_server_info["user"], + connector={"type": "memory"}, + ) + + engine = TrinoComputeEngine( + repo_config=repo_config, + offline_store=MagicMock(), + online_store=MagicMock(), + ) + + mock_registry = MagicMock() + mock_registry.get_entity.return_value = driver_entity + + task = MaterializationTask( + project="test_project", + feature_view=fv, + start_time=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2025, 1, 15, 0, 0, 0, tzinfo=timezone.utc), + tqdm_builder=MagicMock(), + ) + + # First materialization creates the target table + jobs = engine.materialize(mock_registry, task) + assert len(jobs) == 1 + if jobs[0].error() is not None: + raise jobs[0].error() + assert jobs[0].status() == MaterializationJobStatus.SUCCEEDED + + res = live_trino_client.execute_query(f"SELECT COUNT(*) FROM {target_table}") + assert res.data[0][0] == 2 + + # Second materialization appends rows + jobs2 = engine.materialize(mock_registry, task) + assert len(jobs2) == 1 + assert jobs2[0].status() == MaterializationJobStatus.SUCCEEDED + + res2 = live_trino_client.execute_query(f"SELECT COUNT(*) FROM {target_table}") + assert res2.data[0][0] == 4 + + def test_trino_data_types_roundtrip_e2e( + self, trino_server_info, live_trino_client, driver_entity + ): + """Validates end-to-end type mapping from Trino types to Python/Arrow results.""" + catalog = trino_server_info["catalog"] + dataset = trino_server_info["dataset"] + types_table = f"{catalog}.{dataset}.types_test_stats" + + live_trino_client.execute_query(f"DROP TABLE IF EXISTS {types_table}") + types_df = pd.DataFrame( + { + "driver_id": [3001], + "event_timestamp": [ + datetime(2025, 1, 15, 12, 0, 0, tzinfo=timezone.utc) + ], + "created_timestamp": [ + datetime(2025, 1, 15, 12, 0, 0, tzinfo=timezone.utc) + ], + "int_val": [9223372036854775807], # max int64 + "float_val": [3.1415926535], + "str_val": ["trino_roundtrip_test"], + } + ) + upload_pandas_dataframe_to_trino( + client=live_trino_client, + df=types_df, + table=types_table, + connector_args={"type": "memory"}, + ) + + trino_source = TrinoSource( + name="types_source", + table=types_table, + timestamp_field="event_timestamp", + created_timestamp_column="created_timestamp", + ) + + fv = BatchFeatureView( + name="types_fv", + entities=[driver_entity], + schema=[ + Field(name="int_val", dtype=Int64), + Field(name="float_val", dtype=Float32), + Field(name="str_val", dtype=String), + Field(name="driver_id", dtype=Int32), + ], + online=False, + offline=False, + source=trino_source, + ) + + entity_df = pd.DataFrame( + { + "driver_id": [3001], + "event_timestamp": [ + datetime(2025, 1, 16, 0, 0, 0, tzinfo=timezone.utc) + ], + } + ) + + repo_config = MagicMock() + repo_config.batch_engine = TrinoComputeEngineConfig( + host=trino_server_info["host"], + port=trino_server_info["port"], + catalog=catalog, + dataset=dataset, + user=trino_server_info["user"], + connector={"type": "memory"}, + ) + repo_config.offline_store = TrinoOfflineStoreConfig( + host=trino_server_info["host"], + port=trino_server_info["port"], + catalog=catalog, + dataset=dataset, + user=trino_server_info["user"], + connector={"type": "memory"}, + ) + + engine = TrinoComputeEngine( + repo_config=repo_config, + offline_store=MagicMock(), + online_store=MagicMock(), + ) + + mock_registry = MagicMock() + mock_registry.get_entity.return_value = driver_entity + + task = HistoricalRetrievalTask( + project="test_project", + entity_df=entity_df, + feature_view=fv, + full_feature_name=False, + registry=mock_registry, + ) + + job = engine.get_historical_features(mock_registry, task) + df = job.to_df() + assert len(df) == 1 + assert df["int_val"].iloc[0] == 9223372036854775807 + assert pytest.approx(df["float_val"].iloc[0], rel=1e-3) == 3.1415926535 + assert df["str_val"].iloc[0] == "trino_roundtrip_test" diff --git a/sdk/python/tests/unit/infra/compute_engines/trino/test_trino_compute_engine.py b/sdk/python/tests/unit/infra/compute_engines/trino/test_trino_compute_engine.py new file mode 100644 index 00000000000..ad2f2910938 --- /dev/null +++ b/sdk/python/tests/unit/infra/compute_engines/trino/test_trino_compute_engine.py @@ -0,0 +1,1872 @@ +from __future__ import annotations + +from datetime import date, datetime, timedelta, timezone +from unittest.mock import MagicMock, patch + +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest +from trino.exceptions import TrinoConnectionError, TrinoQueryError + +from feast.aggregation import Aggregation +from feast.infra.common.materialization_job import ( + MaterializationJobStatus, +) +from feast.infra.compute_engines.dag.context import ColumnInfo, ExecutionContext +from feast.infra.compute_engines.dag.model import DAGFormat +from feast.infra.compute_engines.dag.value import DAGValue +from feast.infra.compute_engines.trino.compute import ( + TrinoComputeEngine, + TrinoComputeEngineConfig, +) +from feast.infra.compute_engines.trino.job import ( + TrinoDAGRetrievalJob, +) +from feast.infra.compute_engines.trino.nodes import ( + TrinoAggregationNode, + TrinoDedupNode, + TrinoFilterNode, + TrinoJoinNode, + TrinoReadNode, + TrinoTransformationNode, + TrinoValidationNode, + TrinoWriteNode, +) +from feast.infra.compute_engines.trino.sql_builder import TrinoQueryPlan +from feast.infra.compute_engines.trino.utils import ( + _has_unquoted_semicolon, + _trino_sql_literal, + from_feast_to_trino_type, + quote_identifier, + unique_ordered, +) +from feast.infra.offline_stores.contrib.trino_offline_store.trino import ( + TrinoOfflineStoreConfig, +) +from feast.infra.offline_stores.contrib.trino_offline_store.trino_queries import ( + Results, +) +from feast.repo_config import get_batch_engine_config_from_type +from feast.transformation.mode import TransformationMode +from feast.transformation.trino_transformation import TrinoTransformation +from feast.types import Array, Float32, Float64, Int64, String + + +class TestTrinoComputeEngineConfig: + def test_default_config(self): + config = TrinoComputeEngineConfig() + assert config.type in ("trino.engine", "trino") + assert config.dataset == "feast" + assert config.port == 8080 + assert config.http_scheme == "http" + assert config.verify is True + assert config.batch_size == 10000 + assert config.write_concurrency == 4 + assert config.offline_write_mode == "append" + + def test_custom_config(self): + config = TrinoComputeEngineConfig( + host="trino.internal.net", + port=8443, + catalog="iceberg", + user="feast_user", + http_scheme="https", + verify=False, + dataset="feast_staging", + batch_size=5000, + write_concurrency=8, + ) + assert config.host == "trino.internal.net" + assert config.port == 8443 + assert config.catalog == "iceberg" + assert config.user == "feast_user" + assert config.http_scheme == "https" + assert config.verify is False + assert config.dataset == "feast_staging" + assert config.batch_size == 5000 + assert config.write_concurrency == 8 + + def test_get_batch_engine_config_from_type(self): + cls1 = get_batch_engine_config_from_type("trino.engine") + cls2 = get_batch_engine_config_from_type("trino") + assert cls1 is TrinoComputeEngineConfig + assert cls2 is TrinoComputeEngineConfig + + +class TestTrinoComputeEngine: + def test_engine_inherits_from_offline_store(self): + repo_config = MagicMock() + repo_config.batch_engine = TrinoComputeEngineConfig() # host=None + repo_config.offline_store = TrinoOfflineStoreConfig( + host="offline-trino.internal", + port=8443, + catalog="hive", + user="offline_user", + dataset="offline_dataset", + connector={"type": "memory"}, + auth=None, + ) + + engine = TrinoComputeEngine( + repo_config=repo_config, + offline_store=MagicMock(), + online_store=MagicMock(), + ) + + assert engine.config.host == "offline-trino.internal" + assert engine.config.port == 8443 + assert engine.config.catalog == "hive" + assert engine.config.user == "offline_user" + assert engine.config.dataset == "offline_dataset" + + def test_materialize_job_captures_connection_error(self): + mock_client = MagicMock() + mock_client.get_cursor.side_effect = TrinoConnectionError("Trino unreachable") + + repo_config = MagicMock() + repo_config.batch_engine = TrinoComputeEngineConfig( + host="localhost", catalog="iceberg", user="test" + ) + repo_config.offline_store = MagicMock() + + engine = TrinoComputeEngine( + repo_config=repo_config, + offline_store=MagicMock(), + online_store=MagicMock(), + ) + engine.client = mock_client + + task = MagicMock() + task.feature_view.name = "test_fv" + task.feature_view.online = True + task.feature_view.offline = False + task.feature_view.entities = [] + + with patch( + "feast.infra.compute_engines.trino.compute.TrinoFeatureBuilder" + ) as mock_builder_cls: + mock_builder = MagicMock() + mock_plan = MagicMock() + mock_plan.execute.side_effect = TrinoConnectionError("Trino unreachable") + mock_builder.build.return_value = mock_plan + mock_builder_cls.return_value = mock_builder + + mat_job = engine._materialize_one(MagicMock(), task) + assert mat_job.status() == MaterializationJobStatus.ERROR + assert isinstance(mat_job.error(), TrinoConnectionError) + + +class TestTrinoDAGRetrievalJob: + def test_retrieval_job_propagates_connection_error(self): + mock_client = MagicMock() + mock_client.execute_query.side_effect = TrinoConnectionError( + "Network disconnected" + ) + + job = TrinoDAGRetrievalJob( + client=mock_client, + plan=None, + context=None, + full_feature_names=False, + config=MagicMock(), + query="SELECT 1", + ) + + with pytest.raises(TrinoConnectionError, match="Network disconnected"): + job._to_arrow_internal() + + +class TestTrinoSqlFormattingAndTypes: + def test_quote_identifier(self): + assert quote_identifier("user_id") == '"user_id"' + assert quote_identifier('col"with"quote') == '"col""with""quote"' + + def test_has_unquoted_semicolon(self): + assert not _has_unquoted_semicolon('"conv_rate" > 0.5') + assert not _has_unquoted_semicolon("status = 'active;pending'") + assert not _has_unquoted_semicolon('status = "active;pending"') + assert not _has_unquoted_semicolon("status = 'escaped\\'quote;still_in_quotes'") + assert _has_unquoted_semicolon('"conv_rate" > 0.5;') + assert _has_unquoted_semicolon('"conv_rate" > 0.5; ') + assert _has_unquoted_semicolon('"conv_rate" > 0.5; DROP TABLE foo') + assert _has_unquoted_semicolon(";") + + def test_from_feast_to_trino_type(self): + assert from_feast_to_trino_type(Int64) == "BIGINT" + assert from_feast_to_trino_type(Float64) == "DOUBLE" + assert from_feast_to_trino_type(Float32) == "REAL" + assert from_feast_to_trino_type(String) == "VARCHAR" + assert from_feast_to_trino_type(Array(String)) == "ARRAY(VARCHAR)" + + def test_sql_literal_datetime_utc_normalization(self): + # 1. Naive datetime + dt = datetime(2025, 1, 15, 12, 0, 0) + assert _trino_sql_literal(dt) == "TIMESTAMP '2025-01-15 12:00:00.000000'" + + # 2. Timezone-aware datetime with +05:30 offset -> must convert to 06:30:00 UTC + dt_tz = datetime( + 2025, 1, 15, 12, 0, 0, tzinfo=timezone(timedelta(hours=5, minutes=30)) + ) + assert _trino_sql_literal(dt_tz) == "TIMESTAMP '2025-01-15 06:30:00.000000'" + + # 3. pd.Timestamp with timezone + ts_tz = pd.Timestamp("2025-01-15 12:00:00+05:30") + assert _trino_sql_literal(ts_tz) == "TIMESTAMP '2025-01-15 06:30:00.000000'" + + # 4. np.datetime64 + np_dt = np.datetime64("2025-01-15T12:00:00") + assert _trino_sql_literal(np_dt) == "TIMESTAMP '2025-01-15 12:00:00.000000'" + + # 5. pd.NaT and None + assert _trino_sql_literal(pd.NaT) == "NULL" + assert _trino_sql_literal(None) == "NULL" + + # 6. Date + d = date(2025, 1, 15) + assert _trino_sql_literal(d) == "DATE '2025-01-15'" + + def test_sql_literal_containers_and_escaping(self): + # Container arrays should not raise ValueError when checked for nan + assert _trino_sql_literal([1, 2, 3]) == "ARRAY[1, 2, 3]" + assert _trino_sql_literal((4, 5)) == "ARRAY[4, 5]" + assert _trino_sql_literal(np.array([6, 7])) == "ARRAY[6, 7]" + + # Strings with single quotes must be escaped + assert _trino_sql_literal("O'Connor") == "'O''Connor'" + + # Booleans + assert _trino_sql_literal(True) == "TRUE" + assert _trino_sql_literal(False) == "FALSE" + assert _trino_sql_literal(np.bool_(True)) == "TRUE" + + +class TestTrinoWriteNode: + def test_offline_write_staging_swap_failure_preserves_target(self): + mock_client = MagicMock() + executed_queries = [] + + def mock_exec(query): + executed_queries.append(query) + if "CREATE TABLE" in query: + raise TrinoQueryError( + {"message": "Disk full", "errorName": "GENERIC_INTERNAL_ERROR"} + ) + return Results(data=[], columns=[]) + + mock_client.execute_query.side_effect = mock_exec + + fv = MagicMock() + fv.name = "my_fv" + fv.online = False + fv.offline = True + fv.batch_source.get_table_query_string.return_value = ( + "iceberg.feast.my_fv_offline" + ) + + node = TrinoWriteNode( + name="output", + feature_view=fv, + client=mock_client, + ) + + plan = TrinoQueryPlan( + ctes=[("_source", "SELECT 1")], + current_from="_source", + ) + context = MagicMock() + context.repo_config.batch_engine.offline_write_mode = "overwrite" + context.node_outputs = {} + + input_node = MagicMock() + input_node.name = "upstream" + node.inputs = [input_node] + context.node_outputs["upstream"] = DAGValue(data=plan, format=DAGFormat.TRINO) + + with pytest.raises(TrinoQueryError): + node.execute(context) + + # Target table must NEVER have been dropped + assert not any( + "DROP TABLE IF EXISTS iceberg.feast.my_fv_offline\n" in q + for q in executed_queries + ) + assert not any( + q.strip() == "DROP TABLE IF EXISTS iceberg.feast.my_fv_offline" + for q in executed_queries + ) + # Staging table cleanup must have been called + assert any( + 'DROP TABLE IF EXISTS iceberg.feast."my_fv_offline__staging"' in q + for q in executed_queries + ) + + +class TestTrinoDAGCompilation: + """Verifies that all DAG nodes compile into Trino SQL CTEs.""" + + def test_trino_read_node_generates_sql(self): + mock_source = MagicMock() + mock_source.get_table_query_string.return_value = "iceberg.feast.driver_stats" + + col_info = ColumnInfo( + join_keys=["driver_id"], + feature_cols=["conv_rate", "acc_rate"], + ts_col="event_timestamp", + created_ts_col=None, + field_mapping={"conv_rate": "conversion_rate"}, + ) + + start_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc) + end_time = datetime(2025, 1, 2, 0, 0, 0, tzinfo=timezone.utc) + + node = TrinoReadNode( + name="test_read", + source=mock_source, + column_info=col_info, + client=MagicMock(), + start_time=start_time, + end_time=end_time, + ) + + context = MagicMock() + val = node.execute(context) + assert val.format == DAGFormat.TRINO + plan: TrinoQueryPlan = val.data + + sql = plan.to_sql() + assert "WITH\n_source_test_read AS (" in sql + assert '"driver_id"' in sql + assert '"conv_rate" AS "conversion_rate"' in sql + assert "TIMESTAMP '2025-01-01 00:00:00.000000'" in sql + assert "TIMESTAMP '2025-01-02 00:00:00.000000'" in sql + assert "FROM iceberg.feast.driver_stats" in sql + + def test_trino_filter_node_generates_sql(self): + plan = TrinoQueryPlan( + ctes=[("_source", 'SELECT * FROM "iceberg"."feast"."driver_stats"')], + current_from="_source", + columns=("driver_id", "event_timestamp", "__entity_event_timestamp"), + timestamp_col="event_timestamp", + metadata={"entity_joined": True}, + ) + input_node = MagicMock() + input_node.name = "source" + + col_info = ColumnInfo( + join_keys=["driver_id"], + feature_cols=[], + ts_col="event_timestamp", + created_ts_col=None, + ) + node = TrinoFilterNode( + name="filter", + column_info=col_info, + client=MagicMock(), + ttl=timedelta(days=1), + filter_condition='"conv_rate" > 0.5', + inputs=[input_node], + ) + + context = MagicMock() + context.node_outputs = {"source": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(context) + sql = val.data.to_sql() + + assert "_filter_filter AS (" in sql + assert '"event_timestamp" <= "__entity_event_timestamp"' in sql + assert "INTERVAL '86400' SECOND" in sql + assert '("conv_rate" > 0.5)' in sql + + def test_trino_filter_node_semicolon_handling(self): + plan = TrinoQueryPlan( + ctes=[("_source", 'SELECT * FROM "iceberg"."feast"."driver_stats"')], + current_from="_source", + columns=("driver_id", "conv_rate"), + ) + input_node = MagicMock() + input_node.name = "source" + + col_info = ColumnInfo( + join_keys=["driver_id"], + feature_cols=["conv_rate"], + ts_col=None, + created_ts_col=None, + ) + + # 1. Unquoted semicolon must raise ValueError + node_invalid = TrinoFilterNode( + name="filter", + column_info=col_info, + client=MagicMock(), + filter_condition='"conv_rate" > 0.5;', + inputs=[input_node], + ) + context = MagicMock() + context.node_outputs = {"source": DAGValue(data=plan, format=DAGFormat.TRINO)} + + with pytest.raises(ValueError, match="prohibited semicolon"): + node_invalid.execute(context) + + # 2. Semicolon inside string literal is permitted + node_valid = TrinoFilterNode( + name="filter", + column_info=col_info, + client=MagicMock(), + filter_condition="status = 'foo;bar'", + inputs=[input_node], + ) + val = node_valid.execute(context) + sql = val.data.to_sql() + + assert "WHERE (status = 'foo;bar')" in sql + + def test_trino_dedup_node_generates_row_number_sql(self): + plan = TrinoQueryPlan( + ctes=[("_source", 'SELECT * FROM "iceberg"."feast"."driver_stats"')], + current_from="_source", + join_keys=["driver_id"], + timestamp_col="event_timestamp", + created_timestamp_col="created_timestamp", + ) + input_node = MagicMock() + input_node.name = "source" + + col_info = ColumnInfo( + join_keys=["driver_id"], + feature_cols=[], + ts_col="event_timestamp", + created_ts_col="created_timestamp", + ) + node = TrinoDedupNode( + name="dedup", + column_info=col_info, + client=MagicMock(), + inputs=[input_node], + ) + + context = MagicMock() + context.node_outputs = {"source": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(context) + sql = val.data.to_sql() + + assert "ROW_NUMBER() OVER (" in sql + assert 'PARTITION BY "driver_id"' in sql + assert 'ORDER BY "event_timestamp" DESC, "created_timestamp" DESC' in sql + assert "_feast_rn = 1" in sql + # Assert _feast_rn column does NOT leak into output projection + assert "SELECT * FROM (" not in sql + assert ( + 'SELECT "driver_id", "event_timestamp", "created_timestamp" FROM (' in sql + ) + assert "_feast_rn" not in val.data.columns + + def test_trino_dedup_node_uses_upstream_plan_columns(self): + plan = TrinoQueryPlan( + ctes=[ + ( + "_source", + 'SELECT "driver_id", "event_timestamp", "conv_rate" FROM "driver_stats"', + ) + ], + current_from="_source", + columns=("driver_id", "event_timestamp", "conv_rate"), + join_keys=["driver_id"], + timestamp_col="event_timestamp", + ) + input_node = MagicMock() + input_node.name = "source" + + col_info = ColumnInfo( + join_keys=["driver_id"], + feature_cols=["conv_rate"], + ts_col="event_timestamp", + created_ts_col=None, + ) + node = TrinoDedupNode( + name="dedup", + column_info=col_info, + client=MagicMock(), + inputs=[input_node], + ) + + context = MagicMock() + context.node_outputs = {"source": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(context) + sql = val.data.to_sql() + + assert 'SELECT "driver_id", "event_timestamp", "conv_rate" FROM (' in sql + assert "_feast_rn" not in val.data.columns + assert val.data.columns == ("driver_id", "event_timestamp", "conv_rate") + + def test_trino_aggregation_node_generates_group_by_sql(self): + plan = TrinoQueryPlan( + ctes=[("_source", 'SELECT * FROM "iceberg"."feast"."driver_stats"')], + current_from="_source", + join_keys=["driver_id"], + timestamp_col="event_timestamp", + ) + input_node = MagicMock() + input_node.name = "source" + + aggs = [ + Aggregation(column="conv_rate", function="sum"), + Aggregation(column="driver_id", function="count_distinct"), + ] + + node = TrinoAggregationNode( + name="agg", + aggregations=aggs, + group_by_keys=["driver_id"], + timestamp_col="event_timestamp", + client=MagicMock(), + inputs=[input_node], + ) + + context = MagicMock() + context.node_outputs = {"source": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(context) + sql = val.data.to_sql() + + assert "GROUP BY" in sql + assert 'SUM("conv_rate") AS "sum_conv_rate"' in sql + assert 'approx_distinct("driver_id") AS "count_distinct_driver_id"' in sql + + def test_trino_transformation_node_generates_sql(self): + plan = TrinoQueryPlan( + ctes=[("_source", 'SELECT * FROM "iceberg"."feast"."driver_stats"')], + current_from="_source", + join_keys=["driver_id"], + ) + input_node = MagicMock() + input_node.name = "source" + + def double_sql(src_table): + return ( + f'SELECT "driver_id", "conv_rate" * 2.0 AS "conv_rate" FROM {src_table}' + ) + + node = TrinoTransformationNode( + name="transform", + udf=double_sql, + inputs=[input_node], + ) + + context = MagicMock() + context.node_outputs = {"source": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(context) + sql = val.data.to_sql() + + assert "_transform_transform AS (" in sql + assert ( + 'SELECT "driver_id", "conv_rate" * 2.0 AS "conv_rate" FROM _source' in sql + ) + + def test_trino_transformation_class(self): + tf = TrinoTransformation( + mode=TransformationMode.TRINO_SQL, + udf='SELECT "driver_id", "conv_rate" * 2.0 AS "conv_rate" FROM {}', + udf_string='SELECT "driver_id", "conv_rate" * 2.0 AS "conv_rate" FROM {}', + ) + res = tf.transform("upstream_step") + assert ( + 'SELECT "driver_id", "conv_rate" * 2.0 AS "conv_rate" FROM upstream_step' + in res + ) + + +class TestTrinoQueryPlanImmutabilityAndFallbacks: + def test_empty_query_plan_and_fallbacks(self): + empty_plan = TrinoQueryPlan() + assert empty_plan.to_sql() == "SELECT 1" + + plan_from = TrinoQueryPlan(current_from="my_table", columns=("id", "val")) + assert plan_from.to_sql() == "SELECT id, val FROM my_table" + assert plan_from.get_latest_cte_name() == "my_table" + + def test_query_plan_is_frozen(self): + from dataclasses import FrozenInstanceError + + plan = TrinoQueryPlan(current_from="t1") + with pytest.raises(FrozenInstanceError): + plan.current_from = "t2" # type: ignore + + +class TestTrinoJoinNodeExecution: + def test_join_multiple_plans_with_keys(self): + plan1 = TrinoQueryPlan(ctes=(("_p1", "SELECT 1"),), current_from="_p1") + plan2 = TrinoQueryPlan(ctes=(("_p2", "SELECT 2"),), current_from="_p2") + col_info = ColumnInfo( + join_keys=["entity_id"], feature_cols=[], ts_col="ts", created_ts_col=None + ) + + n1 = MagicMock() + n1.name = "n1" + n2 = MagicMock() + n2.name = "n2" + + join_node = TrinoJoinNode( + "join_test", col_info, client=MagicMock(), inputs=[n1, n2] + ) + ctx = ExecutionContext( + repo_config=MagicMock(), + project="p", + offline_store=MagicMock(), + online_store=MagicMock(), + entity_defs=[], + ) + ctx.node_outputs = { + "n1": DAGValue(data=plan1, format=DAGFormat.TRINO), + "n2": DAGValue(data=plan2, format=DAGFormat.TRINO), + } + + val = join_node.execute(ctx) + sql = val.data.to_sql() + assert 'LEFT JOIN _p2 ON _p1."entity_id" = _p2."entity_id"' in sql + + def test_join_without_keys_emits_cross_join(self): + plan1 = TrinoQueryPlan(ctes=(("_p1", "SELECT 1"),), current_from="_p1") + plan2 = TrinoQueryPlan(ctes=(("_p2", "SELECT 2"),), current_from="_p2") + col_info = ColumnInfo( + join_keys=[], feature_cols=[], ts_col=None, created_ts_col=None + ) + + n1 = MagicMock() + n1.name = "n1" + n2 = MagicMock() + n2.name = "n2" + + join_node = TrinoJoinNode( + "join_test", col_info, client=MagicMock(), inputs=[n1, n2] + ) + ctx = ExecutionContext( + repo_config=MagicMock(), + project="p", + offline_store=MagicMock(), + online_store=MagicMock(), + entity_defs=[], + ) + ctx.node_outputs = { + "n1": DAGValue(data=plan1, format=DAGFormat.TRINO), + "n2": DAGValue(data=plan2, format=DAGFormat.TRINO), + } + + val = join_node.execute(ctx) + sql = val.data.to_sql() + assert "CROSS JOIN _p2" in sql + + def test_join_with_entity_df_pandas(self): + plan1 = TrinoQueryPlan( + ctes=(("_p1", "SELECT 1"),), + current_from="_p1", + columns=("user_id", "feature_val", "ts"), + timestamp_col="ts", + ) + col_info = ColumnInfo( + join_keys=["user_id"], + feature_cols=["feature_val"], + ts_col="ts", + created_ts_col=None, + ) + n1 = MagicMock() + n1.name = "n1" + + client = MagicMock() + join_node = TrinoJoinNode("join_entity", col_info, client=client, inputs=[n1]) + + df_entity = pd.DataFrame( + {"user_id": [1, 2], "event_timestamp": [datetime.now(), datetime.now()]} + ) + ctx = ExecutionContext( + repo_config=MagicMock(), + project="p", + offline_store=MagicMock(), + online_store=MagicMock(), + entity_defs=[], + ) + ctx.entity_df = df_entity + ctx.node_outputs = {"n1": DAGValue(data=plan1, format=DAGFormat.TRINO)} + + with patch( + "feast.infra.compute_engines.trino.nodes.upload_pandas_dataframe_to_trino" + ) as mock_upload: + val = join_node.execute(ctx) + mock_upload.assert_called_once() + sql = val.data.to_sql() + assert ( + 'LEFT JOIN _join_join_entity ON _join_join_entity."user_id" = _entity."user_id"' + in sql + ) + assert '_join_join_entity."ts" <= _entity."__entity_event_timestamp"' in sql + # Verify join key 'user_id' is excluded from the feature side select to prevent duplicate columns + assert '_join_join_entity."user_id"' not in sql.split("FROM")[0] + assert '_join_join_entity."feature_val"' in sql + + +class TestTrinoTransformationNodeVariations: + def test_string_template_udf(self): + plan = TrinoQueryPlan(ctes=(("_src", "SELECT 1"),), current_from="_src") + n1 = MagicMock() + n1.name = "n1" + + node = TrinoTransformationNode( + "tf_str", udf="SELECT a * 2 FROM {}", inputs=[n1] + ) + ctx = ExecutionContext( + repo_config=MagicMock(), + project="p", + offline_store=MagicMock(), + online_store=MagicMock(), + entity_defs=[], + ) + ctx.node_outputs = {"n1": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(ctx) + assert "SELECT a * 2 FROM _src" in val.data.to_sql() + + def test_udf_string_fallback(self): + plan = TrinoQueryPlan(ctes=(("_src", "SELECT 1"),), current_from="_src") + n1 = MagicMock() + n1.name = "n1" + + node = TrinoTransformationNode( + "tf_fallback", udf=None, inputs=[n1], udf_string="SELECT x FROM {}" + ) + ctx = ExecutionContext( + repo_config=MagicMock(), + project="p", + offline_store=MagicMock(), + online_store=MagicMock(), + entity_defs=[], + ) + ctx.node_outputs = {"n1": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(ctx) + assert "SELECT x FROM _src" in val.data.to_sql() + + def test_missing_udf_raises(self): + plan = TrinoQueryPlan(ctes=(("_src", "SELECT 1"),), current_from="_src") + n1 = MagicMock() + n1.name = "n1" + + node = TrinoTransformationNode("tf_empty", udf=None, inputs=[n1], udf_string="") + ctx = ExecutionContext( + repo_config=MagicMock(), + project="p", + offline_store=MagicMock(), + online_store=MagicMock(), + entity_defs=[], + ) + ctx.node_outputs = {"n1": DAGValue(data=plan, format=DAGFormat.TRINO)} + + with pytest.raises(ValueError, match="requires SQL UDF string or callable"): + node.execute(ctx) + + +class TestTrinoValidationNode: + def test_validation_node_execute(self): + plan = TrinoQueryPlan(ctes=(("_src", "SELECT 1"),), current_from="_src") + n1 = MagicMock() + n1.name = "n1" + + node = TrinoValidationNode( + "val_node", + expected_columns={"a": "VARCHAR"}, + json_columns={"a"}, + inputs=[n1], + ) + ctx = ExecutionContext( + repo_config=MagicMock(), + project="p", + offline_store=MagicMock(), + online_store=MagicMock(), + entity_defs=[], + ) + ctx.node_outputs = {"n1": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(ctx) + assert ctx.node_outputs["val_node"] == val + assert val.data == plan + + +class TestTrinoAggregationNodeVariations: + def test_aggregation_without_time_window(self): + plan = TrinoQueryPlan(ctes=(("_src", "SELECT 1"),), current_from="_src") + n1 = MagicMock() + n1.name = "n1" + + aggs = [ + Aggregation(column="val", function="avg"), + Aggregation(column="val", function="min"), + Aggregation(column="val", function="max"), + Aggregation(column="val", function="count"), + Aggregation(column="val", function="custom_func"), + ] + node = TrinoAggregationNode( + name="agg_no_window", + aggregations=aggs, + group_by_keys=["user_id"], + timestamp_col="ts", + client=MagicMock(), + inputs=[n1], + ) + ctx = ExecutionContext( + repo_config=MagicMock(), + project="p", + offline_store=MagicMock(), + online_store=MagicMock(), + entity_defs=[], + ) + ctx.node_outputs = {"n1": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(ctx) + sql = val.data.to_sql() + assert 'MAX("ts") AS "ts"' in sql + assert 'AVG("val")' in sql + assert 'MIN("val")' in sql + assert 'MAX("val")' in sql + assert 'COUNT("val")' in sql + assert 'CUSTOM_FUNC("val")' in sql + + +class TestTrinoTypeMappingComprehensive: + def test_complex_types(self): + from feast.types import PrimitiveFeastType, Set, Struct + + assert from_feast_to_trino_type(PrimitiveFeastType.BYTES) == "VARBINARY" + assert from_feast_to_trino_type(PrimitiveFeastType.BOOL) == "BOOLEAN" + assert ( + from_feast_to_trino_type(PrimitiveFeastType.UNIX_TIMESTAMP) == "TIMESTAMP" + ) + assert ( + from_feast_to_trino_type(PrimitiveFeastType.MAP) == "MAP(VARCHAR, VARCHAR)" + ) + + struct_t = Struct( + {"x": PrimitiveFeastType.INT32, "y": PrimitiveFeastType.STRING} + ) + assert from_feast_to_trino_type(struct_t) == 'ROW("x" INTEGER, "y" VARCHAR)' + + arr_struct = Array(struct_t) + assert ( + from_feast_to_trino_type(arr_struct) + == 'ARRAY(ROW("x" INTEGER, "y" VARCHAR))' + ) + + arr_map = Array(PrimitiveFeastType.MAP) + assert from_feast_to_trino_type(arr_map) == "ARRAY(MAP(VARCHAR, VARCHAR))" + + set_str = Set(PrimitiveFeastType.STRING) + assert from_feast_to_trino_type(set_str) == "ARRAY(VARCHAR)" + + +class TestTrinoSqlLiteralComprehensive: + def test_literals(self): + assert _trino_sql_literal(None) == "NULL" + assert _trino_sql_literal(True) == "TRUE" + assert _trino_sql_literal(False) == "FALSE" + assert _trino_sql_literal(b"hello") == "b'hello'" + assert _trino_sql_literal(bytearray(b"abc")) == "bytearray(b'abc')" + assert _trino_sql_literal(date(2025, 5, 10)) == "DATE '2025-05-10'" + assert _trino_sql_literal(float("nan")) == "NULL" + assert _trino_sql_literal([1, 2, 3]) == "ARRAY[1, 2, 3]" + assert _trino_sql_literal({"k": "v"}) == "{'k': 'v'}" + + +class TestTrinoFeatureBuilderDirect: + def test_builder_methods(self): + from feast.infra.compute_engines.trino.feature_builder import ( + TrinoFeatureBuilder, + ) + from feast.types import PrimitiveFeastType + + client = MagicMock() + registry = MagicMock() + task = MagicMock() + task.start_time = datetime(2025, 1, 1, tzinfo=timezone.utc) + task.end_time = datetime(2025, 1, 2, tzinfo=timezone.utc) + task.project = "test_p" + + view = MagicMock() + view.name = "v1" + view.batch_source = MagicMock() + view.batch_source.timestamp_field = "ts" + view.batch_source.field_mapping = {} + view.stream_source = None + view.entities = [] + view.features = [MagicMock(name="f1", dtype=PrimitiveFeastType.STRING)] + view.features[0].name = "f1" + view.aggregations = [] + view.filter = None + view.ttl = None + + with patch( + "feast.infra.compute_engines.feature_builder._get_column_names", + return_value=([], ["f1"], "ts", None), + ): + builder = TrinoFeatureBuilder(registry, client, task) + src_node = builder.build_source_node(view) + assert src_node.name == "v1:source" + + filt_node = builder.build_filter_node(view, src_node) + assert filt_node.name == "v1:filter" + + dedup_node = builder.build_dedup_node(view, filt_node) + assert dedup_node.name == "v1:dedup" + + join_node = builder.build_join_node(view, [src_node]) + assert join_node.name == "v1:join" + + val_node = builder.build_validation_node(view, src_node) + assert val_node.name == "v1:validate" + + def test_should_join_entity_df(self): + from feast.infra.common.materialization_job import MaterializationTask + from feast.infra.common.retrieval_task import HistoricalRetrievalTask + from feast.infra.compute_engines.trino.feature_builder import ( + TrinoFeatureBuilder, + ) + + client = MagicMock() + registry = MagicMock() + + # 1. MaterializationTask -> False + mat_task = MagicMock(spec=MaterializationTask) + mat_task.feature_view = MagicMock() + with patch( + "feast.infra.compute_engines.feature_builder._get_column_names", + return_value=([], [], "ts", None), + ): + b_mat = TrinoFeatureBuilder(registry, client, mat_task) + assert b_mat._should_join_entity_df() is False + + # 2. HistoricalRetrievalTask with DataFrame -> True + ret_task_df = MagicMock(spec=HistoricalRetrievalTask) + ret_task_df.feature_view = MagicMock() + ret_task_df.entity_df = pd.DataFrame({"id": [1]}) + with patch( + "feast.infra.compute_engines.feature_builder._get_column_names", + return_value=([], [], "ts", None), + ): + b_df = TrinoFeatureBuilder(registry, client, ret_task_df) + assert b_df._should_join_entity_df() is True + + # 3. HistoricalRetrievalTask with SQL query string -> True + ret_task_sql = MagicMock(spec=HistoricalRetrievalTask) + ret_task_sql.feature_view = MagicMock() + ret_task_sql.entity_df = "SELECT * FROM entities" + with patch( + "feast.infra.compute_engines.feature_builder._get_column_names", + return_value=([], [], "ts", None), + ): + b_sql = TrinoFeatureBuilder(registry, client, ret_task_sql) + assert b_sql._should_join_entity_df() is True + + # 4. HistoricalRetrievalTask with empty / whitespace string -> False + ret_task_empty = MagicMock(spec=HistoricalRetrievalTask) + ret_task_empty.feature_view = MagicMock() + ret_task_empty.entity_df = " " + with patch( + "feast.infra.compute_engines.feature_builder._get_column_names", + return_value=([], [], "ts", None), + ): + b_empty = TrinoFeatureBuilder(registry, client, ret_task_empty) + assert b_empty._should_join_entity_df() is False + + # 5. HistoricalRetrievalTask with None -> False + ret_task_none = MagicMock(spec=HistoricalRetrievalTask) + ret_task_none.feature_view = MagicMock() + ret_task_none.entity_df = None + with patch( + "feast.infra.compute_engines.feature_builder._get_column_names", + return_value=([], [], "ts", None), + ): + b_none = TrinoFeatureBuilder(registry, client, ret_task_none) + assert b_none._should_join_entity_df() is False + + def test_build_single_feature_view_with_entity_df_injects_join_node(self): + from feast.infra.common.retrieval_task import HistoricalRetrievalTask + from feast.infra.compute_engines.trino.feature_builder import ( + TrinoFeatureBuilder, + ) + from feast.infra.compute_engines.trino.nodes import ( + TrinoDedupNode, + TrinoFilterNode, + TrinoJoinNode, + TrinoReadNode, + ) + + client = MagicMock() + registry = MagicMock() + + view = MagicMock() + view.name = "single_fv" + view.batch_source = MagicMock() + view.batch_source.timestamp_field = "ts" + view.batch_source.field_mapping = {} + view.stream_source = None + view.entities = ["id"] + view.features = [] + view.aggregations = [] + view.filter = None + view.ttl = None + view.feature_transformation = None + view.enable_validation = False + + # When entity_df is provided for HistoricalRetrievalTask: + task = MagicMock(spec=HistoricalRetrievalTask) + task.feature_view = view + task.entity_df = pd.DataFrame({"id": [1], "ts": [datetime.now()]}) + task.project = "test_p" + task.start_time = None + task.end_time = None + + with patch( + "feast.infra.compute_engines.feature_builder._get_column_names", + return_value=(["id"], [], "ts", None), + ): + builder = TrinoFeatureBuilder(registry, client, task) + last_node = builder._build(view, input_nodes=None) + + # Traverse the DAG inputs backwards from the final dedup node + assert isinstance(last_node, TrinoDedupNode) + filter_node = last_node.inputs[0] + assert isinstance(filter_node, TrinoFilterNode) + join_node = filter_node.inputs[0] + assert isinstance(join_node, TrinoJoinNode) + source_node = join_node.inputs[0] + assert isinstance(source_node, TrinoReadNode) + + +class TestTrinoJobsAndEngineMethods: + def test_materialization_job_properties(self): + from feast.infra.compute_engines.trino.job import ( + TrinoMaterializationJob, + ) + + job = TrinoMaterializationJob( + job_id="mat_1", status=MaterializationJobStatus.SUCCEEDED + ) + assert job.job_id() == "mat_1" + assert job.url() is None + assert job.status() == MaterializationJobStatus.SUCCEEDED + assert job.error() is None + assert job.should_be_retried() is False + + def test_retrieval_job_properties(self): + job = TrinoDAGRetrievalJob( + client=MagicMock(), + plan=MagicMock(), + context=MagicMock(), + full_feature_names=False, + config=MagicMock(), + ) + assert job.error() is None + assert job.full_feature_names is False + assert job.on_demand_feature_views == [] + assert job.metadata is None + + def test_engine_infrastructure_no_ops(self): + engine = TrinoComputeEngine( + repo_config=MagicMock(), + offline_store=MagicMock(), + online_store=MagicMock(), + ) + engine.update("p", [], [], [], []) + engine.teardown_infra("p", [], []) + + +class TestTrinoTransformationComprehensive: + def test_mode_as_string(self): + tf = TrinoTransformation( + mode="trino_sql", + udf="SELECT * FROM {}", + ) + assert tf.mode == TransformationMode.TRINO_SQL + assert tf.transform("source_tbl") == "SELECT * FROM source_tbl" + + def test_callable_udf(self): + def my_udf(t): + return f"SELECT id FROM {t}" + + tf = TrinoTransformation(udf=my_udf) + assert tf.transform("t1") == "SELECT id FROM t1" + + def test_static_sql_udf_without_template(self): + tf = TrinoTransformation(udf="SELECT 1") + assert tf.transform("t1") == "SELECT 1" + + def test_udf_string_fallback(self): + tf = TrinoTransformation(udf=None, udf_string="SELECT * FROM {}") + assert tf.transform("tbl") == "SELECT * FROM tbl" + + tf_static = TrinoTransformation(udf=None, udf_string="SELECT 42") + assert tf_static.transform("tbl") == "SELECT 42" + + def test_invalid_transformation_raises(self): + tf = TrinoTransformation(udf=None, udf_string="") + with pytest.raises(ValueError, match="Invalid TrinoTransformation"): + tf.transform("tbl") + + def test_infer_features_noop(self): + tf = TrinoTransformation(udf="SELECT 1") + assert tf.infer_features() is None + + +class TestTrinoDAGRetrievalJobComprehensive: + def test_to_sql_with_string_result(self): + mock_plan = MagicMock() + mock_plan.execute.return_value = MagicMock(data="SELECT 123") + job = TrinoDAGRetrievalJob( + client=MagicMock(), + plan=mock_plan, + context=MagicMock(), + full_feature_names=False, + config=MagicMock(), + ) + assert job.to_sql() == "SELECT 123" + + def test_ensure_executed_with_preexisting_error(self): + err = RuntimeError("Pre-existing retrieval error") + job = TrinoDAGRetrievalJob( + client=MagicMock(), + plan=None, + context=None, + full_feature_names=False, + config=MagicMock(), + error=err, + ) + with pytest.raises(RuntimeError, match="Pre-existing retrieval error"): + job.to_arrow() + + def test_ensure_executed_caching(self): + mock_client = MagicMock() + df = pd.DataFrame({"x": [1, 2, 3]}) + mock_results = MagicMock() + mock_results.to_dataframe.return_value = df + mock_client.execute_query.return_value = mock_results + + job = TrinoDAGRetrievalJob( + client=mock_client, + plan=None, + context=None, + full_feature_names=False, + config=MagicMock(), + query="SELECT x FROM tbl", + ) + t1 = job.to_arrow() + t2 = job.to_arrow() + assert t1 == t2 + assert mock_client.execute_query.call_count == 1 + + def test_ensure_executed_no_client_raises(self): + job = TrinoDAGRetrievalJob( + client=None, + plan=None, + context=None, + full_feature_names=False, + config=MagicMock(), + query="SELECT 1", + ) + with pytest.raises( + RuntimeError, match="No Trino client or SQL query available" + ): + job.to_arrow() + + def test_ensure_executed_trino_query_error(self): + mock_client = MagicMock() + mock_client.execute_query.side_effect = TrinoQueryError( + {"message": "syntax error"} + ) + job = TrinoDAGRetrievalJob( + client=mock_client, + plan=None, + context=None, + full_feature_names=False, + config=MagicMock(), + query="SELECT syntax error", + ) + with pytest.raises(TrinoQueryError): + job.to_arrow() + + def test_to_df(self): + mock_client = MagicMock() + df = pd.DataFrame({"a": [10, 20]}) + mock_results = MagicMock() + mock_results.to_dataframe.return_value = df + mock_client.execute_query.return_value = mock_results + + job = TrinoDAGRetrievalJob( + client=mock_client, + plan=None, + context=None, + full_feature_names=False, + config=MagicMock(), + query="SELECT a FROM tbl", + ) + result_df = job.to_df() + assert list(result_df["a"]) == [10, 20] + + def test_persist(self): + from feast.infra.offline_stores.contrib.trino_offline_store.trino_source import ( + SavedDatasetTrinoStorage, + ) + + mock_client = MagicMock() + job = TrinoDAGRetrievalJob( + client=mock_client, + plan=None, + context=None, + full_feature_names=False, + config=MagicMock(), + query="SELECT 1 AS a", + ) + with pytest.raises(ValueError, match="Expected `SavedDatasetTrinoStorage`"): + job.persist(MagicMock()) + + storage = SavedDatasetTrinoStorage(table="catalog.schema.persisted_tbl") + job.persist(storage, allow_overwrite=True) + assert mock_client.execute_query.call_count == 2 + + +class TestTrinoComputeEngineAdditional: + def test_engine_init_with_dict_and_str_config(self): + repo_conf_dict = MagicMock() + repo_conf_dict.batch_engine = {"type": "trino.engine", "port": 9090} + repo_conf_dict.offline_store = MagicMock() + engine_dict = TrinoComputeEngine( + repo_config=repo_conf_dict, + offline_store=MagicMock(), + online_store=MagicMock(), + ) + assert engine_dict.config.port == 9090 + + repo_conf_str = MagicMock() + repo_conf_str.batch_engine = "trino" + repo_conf_str.offline_store = MagicMock() + engine_str = TrinoComputeEngine( + repo_config=repo_conf_str, + offline_store=MagicMock(), + online_store=MagicMock(), + ) + assert engine_str.config.port == 8080 + + def test_materialize_one_trino_client_not_configured(self): + engine = TrinoComputeEngine( + repo_config=MagicMock(), + offline_store=MagicMock(), + online_store=MagicMock(), + ) + task = MagicMock() + task.feature_view.name = "err_fv" + task.start_time = datetime(2025, 1, 1, tzinfo=timezone.utc) + task.end_time = datetime(2025, 1, 2, tzinfo=timezone.utc) + registry = MagicMock() + + job = engine._materialize_one(registry, task) + assert job.status() == MaterializationJobStatus.ERROR + assert "Trino client is not configured." in str(job.error()) + + def test_materialize_from_offline_store(self): + repo_cfg = MagicMock() + repo_cfg.batch_engine = TrinoComputeEngineConfig() + repo_cfg.offline_store = MagicMock() + mock_offline = MagicMock() + mock_online = MagicMock() + + engine = TrinoComputeEngine( + repo_config=repo_cfg, + offline_store=mock_offline, + online_store=mock_online, + ) + engine.client = MagicMock() + + registry = MagicMock() + fv = MagicMock() + fv.name = "legacy_fv" + fv.entities = [] + fv.batch_source = MagicMock() + + start = datetime(2025, 1, 1, tzinfo=timezone.utc) + end = datetime(2025, 1, 2, tzinfo=timezone.utc) + + with ( + patch( + "feast.infra.compute_engines.trino.compute._get_column_names", + return_value=([], ["f1"], "ts", None), + ), + patch( + "feast.infra.compute_engines.trino.compute.stream_trino_arrow_batches", + return_value=[], + ), + patch( + "feast.infra.compute_engines.trino.compute.write_arrow_batches_to_online_store", + ), + ): + job = engine._materialize_from_offline_store( + registry=registry, + feature_view=fv, + start_date=start, + end_date=end, + project="proj", + ) + assert job.status() == MaterializationJobStatus.SUCCEEDED + + def test_materialize_from_offline_store_error(self): + repo_cfg = MagicMock() + repo_cfg.batch_engine = TrinoComputeEngineConfig() + repo_cfg.offline_store = MagicMock() + + engine = TrinoComputeEngine( + repo_config=repo_cfg, + offline_store=MagicMock(), + online_store=MagicMock(), + ) + engine.client = None + + registry = MagicMock() + fv = MagicMock() + fv.name = "legacy_fv_err" + fv.entities = [] + fv.batch_source = MagicMock() + + start = datetime(2025, 1, 1, tzinfo=timezone.utc) + end = datetime(2025, 1, 2, tzinfo=timezone.utc) + + with patch( + "feast.infra.compute_engines.trino.compute._get_column_names", + return_value=([], ["f1"], "ts", None), + ): + job = engine._materialize_from_offline_store( + registry=registry, + feature_view=fv, + start_date=start, + end_date=end, + project="proj", + ) + assert job.status() == MaterializationJobStatus.ERROR + + def test_get_historical_features_exception(self): + engine = TrinoComputeEngine( + repo_config=MagicMock(), + offline_store=MagicMock(), + online_store=MagicMock(), + ) + engine.client = MagicMock() + task = MagicMock() + task.full_feature_name = False + with patch( + "feast.infra.compute_engines.trino.compute.TrinoFeatureBuilder", + side_effect=ValueError("DAG build error"), + ): + job = engine.get_historical_features(MagicMock(), task) + assert job.error() is not None + with pytest.raises(ValueError, match="DAG build error"): + job.to_arrow() + + def test_get_historical_features_trino_client_not_configured(self): + engine = TrinoComputeEngine( + repo_config=MagicMock(), + offline_store=MagicMock(), + online_store=MagicMock(), + ) + task = MagicMock() + task.full_feature_name = False + job = engine.get_historical_features(MagicMock(), task) + assert job.error() is not None + assert "Trino client is not configured." in str(job.error()) + + +class TestTrinoNodesEdgeCases: + def test_read_node_without_columns_emits_wildcard(self): + src = MagicMock() + src.get_table_query_string.return_value = '"catalog"."schema"."tbl"' + src.field_mapping = {} + + node = TrinoReadNode( + name="read_all", + source=src, + column_info=ColumnInfo([], [], None, None), + client=MagicMock(), + ) + val = node.execute(MagicMock()) + assert "SELECT *" in val.data.to_sql() + assert 'FROM "catalog"."schema"."tbl"' in val.data.to_sql() + + def test_dedup_node_missing_keys_passthrough(self): + input_node = MagicMock() + input_node.name = "source" + plan = TrinoQueryPlan(current_from="tbl") + + col_info = ColumnInfo( + join_keys=[], feature_cols=[], ts_col=None, created_ts_col=None + ) + node = TrinoDedupNode( + name="dedup", + column_info=col_info, + client=MagicMock(), + inputs=[input_node], + ) + context = MagicMock() + context.node_outputs = {"source": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(context) + assert val.data == plan + + def test_aggregation_node_with_time_window(self): + plan = TrinoQueryPlan( + ctes=[("_source", 'SELECT * FROM "driver_stats"')], + current_from="_source", + join_keys=["driver_id"], + timestamp_col="event_timestamp", + ) + input_node = MagicMock() + input_node.name = "source" + + aggs = [ + Aggregation( + column="conv_rate", function="avg", time_window=timedelta(hours=1) + ), + ] + node = TrinoAggregationNode( + name="agg_window", + aggregations=aggs, + group_by_keys=["driver_id"], + timestamp_col="event_timestamp", + client=MagicMock(), + inputs=[input_node], + ) + context = MagicMock() + context.node_outputs = {"source": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(context) + sql = val.data.to_sql() + assert "from_unixtime(floor(to_unixtime(" in sql + assert "AVG(" in sql + + def test_join_node_with_string_entity_df(self): + plan = TrinoQueryPlan( + ctes=[("_source", 'SELECT * FROM "fv"')], + current_from="_source", + join_keys=["id"], + timestamp_col="ts", + ) + input_node = MagicMock() + input_node.name = "upstream" + + col_info = ColumnInfo( + join_keys=["id"], feature_cols=[], ts_col="ts", created_ts_col=None + ) + node = TrinoJoinNode( + name="join_str", + column_info=col_info, + client=MagicMock(), + inputs=[input_node], + ) + context = MagicMock() + context.entity_df = ( + "SELECT 1 AS id, TIMESTAMP '2025-01-01 00:00:00' AS event_timestamp" + ) + context.node_outputs = {"upstream": DAGValue(data=plan, format=DAGFormat.TRINO)} + + val = node.execute(context) + sql = val.data.to_sql() + assert ( + "FROM (SELECT 1 AS id, TIMESTAMP '2025-01-01 00:00:00' AS event_timestamp) AS _entity" + in sql + ) + assert ( + 'LEFT JOIN _join_join_str ON _join_join_str."id" = _entity."id" ' + 'AND _join_join_str."ts" <= _entity."__entity_event_timestamp"' in sql + ) + + def test_write_node_no_batch_source(self): + fv = MagicMock() + fv.online = False + fv.offline = True + fv.batch_source = None + node = TrinoWriteNode(name="out", feature_view=fv, client=MagicMock()) + input_node = MagicMock() + input_node.name = "src" + node.inputs = [input_node] + context = MagicMock() + context.node_outputs = { + "src": DAGValue(data=TrinoQueryPlan(), format=DAGFormat.TRINO) + } + # Should return silently without raising + node.execute(context) + + def test_write_node_offline_append_existing_table(self): + client = MagicMock() + client.execute_query.return_value = MagicMock() + + fv = MagicMock() + fv.online = False + fv.offline = True + batch_source = MagicMock() + batch_source.get_table_query_string.return_value = "feast.driver_features" + fv.batch_source = batch_source + + node = TrinoWriteNode(name="out", feature_view=fv, client=client) + input_node = MagicMock() + input_node.name = "src" + node.inputs = [input_node] + + plan = TrinoQueryPlan( + ctes=[("_src", "SELECT 1")], + current_from="_src", + columns=("driver_id", "feature_val"), + ) + context = MagicMock() + context.repo_config.batch_engine.offline_write_mode = "append" + context.node_outputs = {"src": DAGValue(data=plan, format=DAGFormat.TRINO)} + + node.execute(context) + + executed_queries = [call[0][0] for call in client.execute_query.call_args_list] + assert any( + "SELECT 1 FROM feast.driver_features WHERE 1 = 0" in q + for q in executed_queries + ) + assert any( + 'INSERT INTO feast.driver_features ("driver_id", "feature_val")' in q + for q in executed_queries + ) + assert not any("DROP TABLE" in q for q in executed_queries) + + def test_write_node_offline_append_new_table(self): + client = MagicMock() + + def mock_query(q): + if "SELECT 1 FROM" in q: + raise TrinoQueryError( + {"errorName": "TABLE_NOT_FOUND", "message": "Table not found"} + ) + return MagicMock() + + client.execute_query.side_effect = mock_query + + fv = MagicMock() + fv.online = False + fv.offline = True + batch_source = MagicMock() + batch_source.get_table_query_string.return_value = "feast.new_features" + fv.batch_source = batch_source + + node = TrinoWriteNode(name="out", feature_view=fv, client=client) + input_node = MagicMock() + input_node.name = "src" + node.inputs = [input_node] + + plan = TrinoQueryPlan( + ctes=[("_src", "SELECT 1")], + current_from="_src", + columns=("driver_id", "feature_val"), + ) + context = MagicMock() + context.repo_config.batch_engine.offline_write_mode = "append" + context.node_outputs = {"src": DAGValue(data=plan, format=DAGFormat.TRINO)} + + node.execute(context) + + executed_queries = [call[0][0] for call in client.execute_query.call_args_list] + assert any("CREATE TABLE feast.new_features AS" in q for q in executed_queries) + assert not any("INSERT INTO" in q for q in executed_queries) + + def test_write_node_offline_overwrite_safe_swap(self): + client = MagicMock() + client.execute_query.return_value = MagicMock() + + fv = MagicMock() + fv.online = False + fv.offline = True + batch_source = MagicMock() + batch_source.get_table_query_string.return_value = "feast.driver_features" + fv.batch_source = batch_source + + node = TrinoWriteNode(name="out", feature_view=fv, client=client) + input_node = MagicMock() + input_node.name = "src" + node.inputs = [input_node] + + plan = TrinoQueryPlan( + ctes=[("_src", "SELECT 1")], + current_from="_src", + ) + context = MagicMock() + context.repo_config.batch_engine.offline_write_mode = "overwrite" + context.node_outputs = {"src": DAGValue(data=plan, format=DAGFormat.TRINO)} + + node.execute(context) + + executed_queries = [call[0][0] for call in client.execute_query.call_args_list] + assert any( + 'CREATE TABLE feast."driver_features__staging" AS' in q + for q in executed_queries + ) + assert any( + 'ALTER TABLE feast.driver_features RENAME TO "driver_features__backup"' in q + for q in executed_queries + ) + assert any( + 'ALTER TABLE feast."driver_features__staging" RENAME TO driver_features' + in q + for q in executed_queries + ) + assert any( + 'DROP TABLE IF EXISTS feast."driver_features__backup"' in q + for q in executed_queries + ) + assert not any( + "DROP TABLE IF EXISTS feast.driver_features" == q for q in executed_queries + ) + + def test_write_node_offline_overwrite_rollback_on_failure(self): + client = MagicMock() + + def mock_query(q): + if 'ALTER TABLE feast."driver_features__staging" RENAME' in q: + raise TrinoQueryError( + {"errorName": "GENERIC_INTERNAL_ERROR", "message": "Rename failed"} + ) + return MagicMock() + + client.execute_query.side_effect = mock_query + + fv = MagicMock() + fv.online = False + fv.offline = True + batch_source = MagicMock() + batch_source.get_table_query_string.return_value = "feast.driver_features" + fv.batch_source = batch_source + + node = TrinoWriteNode(name="out", feature_view=fv, client=client) + input_node = MagicMock() + input_node.name = "src" + node.inputs = [input_node] + + plan = TrinoQueryPlan( + ctes=[("_src", "SELECT 1")], + current_from="_src", + ) + context = MagicMock() + context.repo_config.batch_engine.offline_write_mode = "overwrite" + context.node_outputs = {"src": DAGValue(data=plan, format=DAGFormat.TRINO)} + + with pytest.raises(TrinoQueryError): + node.execute(context) + + executed_queries = [call[0][0] for call in client.execute_query.call_args_list] + assert any( + 'ALTER TABLE feast."driver_features__backup" RENAME TO driver_features' in q + for q in executed_queries + ) + assert any( + 'DROP TABLE IF EXISTS feast."driver_features__staging"' in q + for q in executed_queries + ) + + def test_write_node_offline_merge_existing_table(self): + client = MagicMock() + client.execute_query.return_value = MagicMock() + + fv = MagicMock() + fv.online = False + fv.offline = True + batch_source = MagicMock() + batch_source.get_table_query_string.return_value = ( + "iceberg.feast.driver_features" + ) + fv.batch_source = batch_source + + node = TrinoWriteNode(name="out", feature_view=fv, client=client) + input_node = MagicMock() + input_node.name = "src" + node.inputs = [input_node] + + plan = TrinoQueryPlan( + ctes=[("_src", "SELECT 1")], + current_from="_src", + join_keys=("driver_id",), + timestamp_col="event_timestamp", + columns=("driver_id", "event_timestamp", "feature_val"), + ) + context = MagicMock() + context.repo_config.batch_engine.offline_write_mode = "merge" + context.node_outputs = {"src": DAGValue(data=plan, format=DAGFormat.TRINO)} + + node.execute(context) + + executed_queries = [call[0][0] for call in client.execute_query.call_args_list] + assert any( + "MERGE INTO iceberg.feast.driver_features AS target" in q + for q in executed_queries + ) + assert any("USING (\nWITH\n_src AS (" in q for q in executed_queries) + assert any( + 'ON target."driver_id" = stage."driver_id" AND target."event_timestamp" = stage."event_timestamp"' + in q + for q in executed_queries + ) + assert any( + 'UPDATE SET "feature_val" = stage."feature_val"' in q + for q in executed_queries + ) + assert any( + 'INSERT ("driver_id", "event_timestamp", "feature_val")' in q + for q in executed_queries + ) + assert not any("CREATE TABLE" in q for q in executed_queries) + assert not any("DROP TABLE" in q for q in executed_queries) + + def test_write_node_offline_merge_new_table(self): + client = MagicMock() + + def mock_query(q): + if "SELECT 1 FROM" in q: + raise TrinoQueryError( + {"errorName": "TABLE_NOT_FOUND", "message": "Table not found"} + ) + return MagicMock() + + client.execute_query.side_effect = mock_query + + fv = MagicMock() + fv.online = False + fv.offline = True + batch_source = MagicMock() + batch_source.get_table_query_string.return_value = "iceberg.feast.new_features" + fv.batch_source = batch_source + + node = TrinoWriteNode(name="out", feature_view=fv, client=client) + input_node = MagicMock() + input_node.name = "src" + node.inputs = [input_node] + + plan = TrinoQueryPlan( + ctes=[("_src", "SELECT 1")], + current_from="_src", + join_keys=("driver_id",), + timestamp_col="event_timestamp", + columns=("driver_id", "event_timestamp", "feature_val"), + ) + context = MagicMock() + context.repo_config.batch_engine.offline_write_mode = "merge" + context.node_outputs = {"src": DAGValue(data=plan, format=DAGFormat.TRINO)} + + node.execute(context) + + executed_queries = [call[0][0] for call in client.execute_query.call_args_list] + assert any( + "CREATE TABLE iceberg.feast.new_features AS" in q for q in executed_queries + ) + assert not any("MERGE INTO" in q for q in executed_queries) + + def test_write_node_offline_merge_fallback_on_missing_keys(self): + client = MagicMock() + client.execute_query.return_value = MagicMock() + + fv = MagicMock() + fv.online = False + fv.offline = True + batch_source = MagicMock() + batch_source.get_table_query_string.return_value = ( + "iceberg.feast.driver_features" + ) + fv.batch_source = batch_source + + node = TrinoWriteNode(name="out", feature_view=fv, client=client) + input_node = MagicMock() + input_node.name = "src" + node.inputs = [input_node] + + plan = TrinoQueryPlan( + ctes=[("_src", "SELECT 1")], + current_from="_src", + join_keys=(), + timestamp_col=None, + columns=("val1",), + ) + context = MagicMock() + context.repo_config.batch_engine.offline_write_mode = "merge" + context.node_outputs = {"src": DAGValue(data=plan, format=DAGFormat.TRINO)} + + node.execute(context) + + executed_queries = [call[0][0] for call in client.execute_query.call_args_list] + assert any( + 'INSERT INTO iceberg.feast.driver_features ("val1")' in q + for q in executed_queries + ) + assert not any("MERGE INTO" in q for q in executed_queries) + + +class TestTrinoUtilsStreamingAndOnlineWrite: + def test_stream_trino_arrow_batches_no_columns(self): + from feast.infra.compute_engines.trino.utils import stream_trino_arrow_batches + + mock_client = MagicMock() + mock_cursor = MagicMock() + mock_cursor.description = None + mock_cursor._query = None + mock_client.get_cursor.return_value = mock_cursor + + batches = list(stream_trino_arrow_batches(mock_client, "SELECT 1")) + assert len(batches) == 0 + + def test_stream_trino_arrow_batches_fetch(self): + from feast.infra.compute_engines.trino.utils import stream_trino_arrow_batches + + mock_client = MagicMock() + mock_cursor = MagicMock() + mock_cursor.description = [("id", "bigint", None, None, None, None, None)] + mock_cursor.fetchmany.side_effect = [[(1,), (2,)], []] + mock_client.get_cursor.return_value = mock_cursor + + batches = list( + stream_trino_arrow_batches(mock_client, "SELECT id FROM tbl", batch_size=2) + ) + assert len(batches) == 1 + assert batches[0].num_rows == 2 + mock_client.get_cursor.assert_called_once() + + def test_write_arrow_batches_to_online_store(self): + from feast.infra.compute_engines.trino.utils import ( + write_arrow_batches_to_online_store, + ) + + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2], type=pa.int64())], + names=["entity_id"], + ) + fv = MagicMock() + fv.entity_columns = [] + fv.batch_source = None + + repo_cfg = MagicMock() + repo_cfg.materialization_config.online_write_batch_size = 10 + + online_store = MagicMock() + progress = MagicMock() + + write_arrow_batches_to_online_store( + batches=[batch], + feature_view=fv, + online_store=online_store, + repo_config=repo_cfg, + concurrency=2, + progress_callback=progress, + ) + assert online_store.online_write_batch.called + assert progress.called + + +def test_unique_ordered(): + assert unique_ordered([1, 2, 2, 3, 1, 4]) == (1, 2, 3, 4) + assert unique_ordered(("a", "b", "a", "c")) == ("a", "b", "c") + assert unique_ordered([]) == () diff --git a/sdk/python/tests/unit/transformation/test_mode.py b/sdk/python/tests/unit/transformation/test_mode.py index 3b35edcba0f..7501001deb6 100644 --- a/sdk/python/tests/unit/transformation/test_mode.py +++ b/sdk/python/tests/unit/transformation/test_mode.py @@ -12,6 +12,7 @@ def test_all_modes_defined(self): "RAY", "SQL", "SUBSTRAIT", + "TRINO_SQL", } actual = {m.name for m in TransformationMode} assert actual == expected @@ -25,7 +26,9 @@ def test_mode_values(self): assert TransformationMode.RAY.value == "ray" assert TransformationMode.SQL.value == "sql" assert TransformationMode.SUBSTRAIT.value == "substrait" + assert TransformationMode.TRINO_SQL.value == "trino_sql" def test_mode_from_value(self): assert TransformationMode("python") == TransformationMode.PYTHON assert TransformationMode("pandas") == TransformationMode.PANDAS + assert TransformationMode("trino_sql") == TransformationMode.TRINO_SQL