fix(trino): Clean up temporary entity tables after retrieval (feast-dev#6381)

Jwrede · rpathade · commit 78fcd7257677 · 2026-05-20T20:37:09.000-07:00
* fix(trino): Clean up temporary entity tables after retrieval TrinoOfflineStore.get_historical_features() creates a temporary table for the entity DataFrame but never drops it, leaking tables indefinitely. Apply the same context manager pattern used by BigQuery, Redshift, and Athena offline stores: wrap the query in a generator that issues DROP TABLE IF EXISTS in a finally block. Fixes feast-dev#6306 Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> * fix: sort imports for ruff compliance Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> * fix: decouple temp table cleanup from query access Avoid dropping the temporary entity table on to_sql() calls. Previously, every method used a context manager that dropped the table on exit, so calling to_sql() before to_df() would destroy the table and cause subsequent queries to fail. Now the query is stored as a plain string and cleanup is handled by a dedicated _drop_temp_table() method called only after query execution (to_df, to_trino). A __del__ fallback ensures cleanup if execution methods are never called. The _cleaned_up flag makes the drop idempotent. Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> --------- Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> Signed-off-by: RutujaPathade <73137503+RutujaPathade@users.noreply.github.com>
diff --git a/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino.py b/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino.py
@@ -1,6 +1,15 @@
+import logging
 import uuid
 from datetime import date, datetime
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import numpy as np
 import pandas as pd
@@ -37,6 +46,8 @@
 from feast.repo_config import FeastConfigBaseModel, RepoConfig
 from feast.saved_dataset import SavedDatasetStorage
 
+logger = logging.getLogger(__name__)
+
 
 class BasicAuthModel(FeastConfigBaseModel):
     username: StrictStr
@@ -183,13 +194,16 @@ def __init__(
         full_feature_names: bool,
         on_demand_feature_views: Optional[List[OnDemandFeatureView]] = None,
         metadata: Optional[RetrievalMetadata] = None,
+        temp_table: Optional[str] = None,
     ):
         self._query = query
         self._client = client
         self._config = config
         self._full_feature_names = full_feature_names
         self._on_demand_feature_views = on_demand_feature_views or []
         self._metadata = metadata
+        self._temp_table = temp_table
+        self._cleaned_up = False
 
     @property
     def full_feature_names(self) -> bool:
@@ -199,11 +213,29 @@ def full_feature_names(self) -> bool:
     def on_demand_feature_views(self) -> List[OnDemandFeatureView]:
         return self._on_demand_feature_views
 
+    def _drop_temp_table(self) -> None:
+        if self._cleaned_up or not self._temp_table:
+            return
+        self._cleaned_up = True
+        try:
+            self._client.execute_query(f"DROP TABLE IF EXISTS {self._temp_table}")
+        except Exception:
+            logger.exception(
+                "Failed to drop temporary entity table %s",
+                self._temp_table,
+            )
+
+    def __del__(self) -> None:
+        self._drop_temp_table()
+
     def _to_df_internal(self, timeout: Optional[int] = None) -> pd.DataFrame:
         """Return dataset as Pandas DataFrame synchronously including on demand transforms"""
-        results = self._client.execute_query(query_text=self._query)
-        self.pyarrow_schema = results.pyarrow_schema
-        return results.to_dataframe()
+        try:
+            results = self._client.execute_query(query_text=self._query)
+            self.pyarrow_schema = results.pyarrow_schema
+            return results.to_dataframe()
+        finally:
+            self._drop_temp_table()
 
     def _to_arrow_internal(self, timeout: Optional[int] = None) -> pyarrow.Table:
         """Return payrrow dataset as synchronously including on demand transforms"""
@@ -234,8 +266,11 @@ def to_trino(
             destination_table = f"{self._client.catalog}.{self._config.offline_store.dataset}.historical_{today}_{rand_id}"
 
         # TODO: Implement the timeout logic
-        query = f"CREATE TABLE {destination_table} AS ({self._query})"
-        self._client.execute_query(query_text=query)
+        try:
+            create_query = f"CREATE TABLE {destination_table} AS ({self._query})"
+            self._client.execute_query(query_text=create_query)
+        finally:
+            self._drop_temp_table()
         return destination_table
 
     def persist(
@@ -372,11 +407,12 @@ def get_historical_features(
         )
 
         # Generate the Trino SQL query from the query context
+        entity_table_ref = table_reference
         if type(entity_df) is str:
-            table_reference = f"({entity_df})"
+            entity_table_ref = f"({entity_df})"
         query = offline_utils.build_point_in_time_query(
             query_context,
-            left_table_query_string=table_reference,
+            left_table_query_string=entity_table_ref,
             entity_df_event_timestamp_col=entity_df_event_timestamp_col,
             entity_df_columns=entity_schema.keys(),
             query_template=MULTIPLE_FEATURE_VIEW_POINT_IN_TIME_JOIN,
@@ -385,6 +421,7 @@ def get_historical_features(
 
         return TrinoRetrievalJob(
             query=query,
+            temp_table=table_reference if isinstance(entity_df, pd.DataFrame) else None,
             client=client,
             config=config,
             full_feature_names=full_feature_names,
@@ -483,8 +520,6 @@ def _upload_entity_df_and_get_entity_schema(
     else:
         raise InvalidEntityType(type(entity_df))
 
-    # TODO: Ensure that the table expires after some time
-
 
 def _get_trino_client(config: RepoConfig) -> Trino:
     auth = None
