Skip to content

feat: Offline Store historical features retrieval based on datetime range in Ray #5738

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ntkathole merged 5 commits into from
Dec 30, 2025
Merged

feat: Offline Store historical features retrieval based on datetime range in Ray #5738

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1b5aebe
feat: Offline Store historical features retrieval based on datetime r…
aniketpalu Nov 25, 2025
de8b2c5
Reforamatted code to fix lint issues
aniketpalu Nov 25, 2025
c25280c
preserve event_timestamp in non-entity mode for correct point-in-time…
aniketpalu Dec 26, 2025
0c20ac0
Minor lint changes
aniketpalu Dec 26, 2025
f589956
Added test cases for datetime range based feature retrieval in Ray
aniketpalu Dec 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
preserve event_timestamp in non-entity mode for correct point-in-time… 
… joins

Signed-off-by: Aniket Paluskar <apaluska@redhat.com>
  • Loading branch information
@aniketpalu @ntkathole
aniketpalu authored and ntkathole committed Dec 30, 2025
commit c25280c6722cb64a63f685666405c2530f76a4eb
82 changes: 54 additions & 28 deletions sdk/python/feast/infra/offline_stores/contrib/ray_offline_store/ray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1237,15 +1237,24 @@ def _filter_range(batch: pd.DataFrame) -> pd.Series:
return _filter_range


def _make_select_distinct_keys(join_keys: List[str]):
# Why: factory function for distinct key projection in Ray map_batches
def _select_distinct_keys(batch: pd.DataFrame) -> pd.DataFrame:
def _make_select_distinct_entity_timestamps(
join_keys: List[str], timestamp_field: str
):
# Why: factory function for distinct (entity_keys, event_timestamp) projection in Ray map_batches
# This preserves multiple transactions per entity ID with different timestamps for proper PIT joins
def _select_distinct_entity_timestamps(batch: pd.DataFrame) -> pd.DataFrame:
cols = [c for c in join_keys if c in batch.columns]
if timestamp_field in batch.columns:
# Rename timestamp to standardized event_timestamp
batch = batch.copy()
if timestamp_field != "event_timestamp":
batch["event_timestamp"] = batch[timestamp_field]
cols = cols + ["event_timestamp"]
if not cols:
return pd.DataFrame(columns=join_keys)
return pd.DataFrame(columns=join_keys + ["event_timestamp"])
return batch[cols].drop_duplicates().reset_index(drop=True)

return _select_distinct_keys
return _select_distinct_entity_timestamps


def _distinct_entities_for_feature_view_ray(
Expand All @@ -1257,7 +1266,8 @@ def _distinct_entities_for_feature_view_ray(
start_date: datetime,
end_date: datetime,
) -> Tuple[Dataset, List[str]]:
# Why: read minimal columns, filter by time, and project distinct join keys per FeatureView
# Why: read minimal columns, filter by time, and project distinct (join_keys, event_timestamp) per FeatureView
# This preserves multiple transactions per entity ID for proper point-in-time joins
ray_wrapper = get_ray_wrapper()
entities = fv.entities or []
entity_objs = [registry.get_entity(e, project) for e in entities]
Expand Down Expand Up @@ -1287,62 +1297,73 @@ def _distinct_entities_for_feature_view_ray(

ds = ensure_timestamp_compatibility(ds, [timestamp_field])
ds = ds.filter(_make_filter_range(timestamp_field, start_date, end_date))
ds = ds.map_batches(_make_select_distinct_keys(join_keys), batch_format="pandas")
# Extract distinct (entity_keys, event_timestamp) combinations - not just entity_keys
ds = ds.map_batches(
_make_select_distinct_entity_timestamps(join_keys, timestamp_field),
batch_format="pandas",
)
return ds, join_keys


def _make_align_columns(all_join_keys: List[str]):
def _make_align_columns(all_join_keys: List[str], include_timestamp: bool = False):
# Why: factory function for schema alignment in Ray map_batches
# When include_timestamp=True, also aligns event_timestamp column for proper PIT joins
def _align_columns(batch: pd.DataFrame) -> pd.DataFrame:
batch = batch.copy()
for k in all_join_keys:
output_cols = list(all_join_keys)
if include_timestamp:
output_cols = output_cols + ["event_timestamp"]
for k in output_cols:
if k not in batch.columns:
batch[k] = pd.NA
return batch[all_join_keys]
return batch[output_cols]

return _align_columns


def _make_distinct_by_keys(keys: List[str]):
def _make_distinct_by_keys(keys: List[str], include_timestamp: bool = False):
# Why: factory function for deduplication in Ray map_batches
# When include_timestamp=True, deduplicates on (keys + event_timestamp) for proper PIT joins
def _distinct(batch: pd.DataFrame) -> pd.DataFrame:
return batch.drop_duplicates(subset=keys).reset_index(drop=True)
subset = list(keys)
if include_timestamp and "event_timestamp" in batch.columns:
subset = subset + ["event_timestamp"]
return batch.drop_duplicates(subset=subset).reset_index(drop=True)

return _distinct


def _align_and_union_entities_ray(
datasets: List[Dataset],
all_join_keys: List[str],
include_timestamp: bool = False,
) -> Dataset:
# Why: align schemas across FeatureViews and union to a unified entity set
# When include_timestamp=True, preserves distinct (entity_keys, event_timestamp) combinations
# for proper point-in-time joins with multiple transactions per entity
ray_wrapper = get_ray_wrapper()
output_cols = list(all_join_keys)
if include_timestamp:
output_cols = output_cols + ["event_timestamp"]
if not datasets:
return ray_wrapper.from_pandas(pd.DataFrame(columns=all_join_keys))
return ray_wrapper.from_pandas(pd.DataFrame(columns=output_cols))

aligned = [
ds.map_batches(_make_align_columns(all_join_keys), batch_format="pandas")
ds.map_batches(
_make_align_columns(all_join_keys, include_timestamp=include_timestamp),
batch_format="pandas",
)
for ds in datasets
]
entity_ds = aligned[0]
for ds in aligned[1:]:
entity_ds = entity_ds.union(ds)
return entity_ds.map_batches(
_make_distinct_by_keys(all_join_keys), batch_format="pandas"
_make_distinct_by_keys(all_join_keys, include_timestamp=include_timestamp),
batch_format="pandas",
)


def _add_asof_ts_ray(ds: Dataset, end_date: datetime) -> Dataset:
# Why: use a stable as-of timestamp for PIT joins when deriving entities
def _add_asof_ts(batch: pd.DataFrame) -> pd.DataFrame:
batch = batch.copy()
batch["event_timestamp"] = end_date
return batch

ds = ds.map_batches(_add_asof_ts, batch_format="pandas")
return ensure_timestamp_compatibility(ds, ["event_timestamp"])


class RayOfflineStore(OfflineStore):
def __init__(self) -> None:
self._staging_location: Optional[str] = None
Expand Down Expand Up @@ -2033,6 +2054,8 @@ def get_historical_features(
ray_wrapper = get_ray_wrapper()
if entity_df is None:
# Non-entity mode: derive entity set from feature sources within a bounded time window
# Preserves distinct (entity_keys, event_timestamp) combinations for proper PIT joins
# This handles cases where multiple transactions per entity ID exist
start_date, end_date = _compute_non_entity_dates_ray(
feature_views, kwargs.get("start_date"), kwargs.get("end_date")
)
Expand All @@ -2046,8 +2069,11 @@ def get_historical_features(
for k in join_keys:
if k not in all_join_keys:
all_join_keys.append(k)
entity_ds = _align_and_union_entities_ray(per_view_entity_ds, all_join_keys)
entity_ds = _add_asof_ts_ray(entity_ds, end_date)
# Use include_timestamp=True to preserve actual event_timestamp from data
# instead of assigning a fixed end_date to all entities
entity_ds = _align_and_union_entities_ray(
per_view_entity_ds, all_join_keys, include_timestamp=True
)
entity_df_sample = entity_ds.limit(1000).to_pandas()
elif isinstance(entity_df, str):
entity_ds = ray_wrapper.read_csv(entity_df)
Expand Down