Fetching pre-filtered table from db

Signed-off-by: Aniket Paluskar <apaluska@redhat.com>
ntkathole · Mar 14, 2026 · Feb 24, 2026 · Feb 24, 2026 · Feb 24, 2026 · Feb 24, 2026
commit b99128868018e11fa9113765179354562b5cd8ee
@@ -1,6 +1,6 @@
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
-from typing import Any, Callable, List, Literal, Optional, Union
+from typing import Any, Callable, List, Literal, Optional, Tuple, Union
 
 import ibis
 import pandas as pd
@@ -76,19 +76,44 @@ def _read_data_source(data_source: DataSource, repo_path: str = "") -> Table:
     return _read_data_source
 
 
-def _build_data_source_reader_for_retrieval(config: RepoConfig, con=None):
+def _build_data_source_reader_for_retrieval(
+    config: RepoConfig,
+    con=None,
+    entity_ts_range: Optional[Tuple[datetime, datetime]] = None,
+    max_ttl: Optional[timedelta] = None,
+):
     """Build a reader that materializes Oracle data into an in-memory table.
 
     Used by ``get_historical_features`` which joins feature tables with an
     in-memory entity table (``ibis.memtable``).  Both sides must be on the
     same backend for computed columns like ``entity_row_id`` to survive the
     join — converting to memtable ensures this.
+
+    When ``entity_ts_range`` is provided the reader pre-filters the Oracle
+    table by timestamp before materializing, so only the rows relevant to
+    the point-in-time join are pulled into memory.
     """
     if con is None:
         con = get_ibis_connection(config)
 
     def _read_data_source(data_source: DataSource, repo_path: str = "") -> Table:
         table = _read_oracle_table(con, data_source)
+
+        # Pre-filter by entity timestamp range to avoid reading the
+        # entire table into memory.  The bounds mirror what
+        # point_in_time_join applies later:
+        #   feature_ts <= max(entity_ts)
+        #   feature_ts >= min(entity_ts) - max_ttl
+        if entity_ts_range and data_source.timestamp_field:
+            ts_field = data_source.timestamp_field
+            if ts_field in table.columns:
+                lower = entity_ts_range[0] - (max_ttl or timedelta(0))
+                upper = entity_ts_range[1]
+                table = table.filter(
+                    (table[ts_field] >= ibis.literal(lower))
+                    & (table[ts_field] <= ibis.literal(upper))
+                )
+
         return ibis.memtable(table.execute())
 
     return _read_data_source
@@ -254,6 +279,25 @@ def get_historical_features(
         if isinstance(entity_df, str):
             entity_df = con.sql(entity_df).execute()
 
+        # Compute entity timestamp range so the reader can pre-filter
+        # Oracle tables instead of loading them entirely into memory.
+        entity_ts_range: Optional[tuple] = None
+        if isinstance(entity_df, pd.DataFrame) and "event_timestamp" in entity_df.columns:
+            ts_series = pd.to_datetime(entity_df["event_timestamp"], utc=True)
+            entity_ts_range = (
+                ts_series.min().to_pydatetime(),
+                ts_series.max().to_pydatetime(),
+            )
+
+        max_ttl = max(
+            (
+                fv.ttl
+                for fv in feature_views
+                if fv.ttl and isinstance(fv.ttl, timedelta)
+            ),
+            default=timedelta(0),
+        )
+
         # Use the retrieval reader which materializes Oracle data into
         # in-memory tables so the point-in-time join with the entity
         # memtable happens on the same backend.
@@ -265,7 +309,12 @@ def get_historical_features(
             registry=registry,
             project=project,
             full_feature_names=full_feature_names,
-            data_source_reader=_build_data_source_reader_for_retrieval(config, con=con),
+            data_source_reader=_build_data_source_reader_for_retrieval(
+                config,
+                con=con,
+                entity_ts_range=entity_ts_range,
+                max_ttl=max_ttl,
+            ),
             data_source_writer=_build_data_source_writer(config, con=con),
         )