(
entity_ds: Dataset,
feature_ds: Dataset,
join_keys: List[str],
timestamp_field: str,
requested_feats: List[str],
window_size: Optional[str] = None,
full_feature_names: bool = False,
feature_view_name: Optional[str] = None,
original_join_keys: Optional[List[str]] = None,
)
| 466 | |
| 467 | |
| 468 | def distributed_windowed_join( |
| 469 | entity_ds: Dataset, |
| 470 | feature_ds: Dataset, |
| 471 | join_keys: List[str], |
| 472 | timestamp_field: str, |
| 473 | requested_feats: List[str], |
| 474 | window_size: Optional[str] = None, |
| 475 | full_feature_names: bool = False, |
| 476 | feature_view_name: Optional[str] = None, |
| 477 | original_join_keys: Optional[List[str]] = None, |
| 478 | ) -> Dataset: |
| 479 | import pandas as pd |
| 480 | |
| 481 | def add_window_and_source(ds, timestamp_field, source_marker, window_size): |
| 482 | def add_window_and_source_batch(batch: pd.DataFrame) -> pd.DataFrame: |
| 483 | batch = batch.copy() |
| 484 | if timestamp_field in batch.columns: |
| 485 | batch["time_window"] = ( |
| 486 | pd.to_datetime(batch[timestamp_field]) |
| 487 | .dt.floor(window_size) |
| 488 | .astype("datetime64[ns, UTC]") |
| 489 | ) |
| 490 | batch["_data_source"] = source_marker |
| 491 | return batch |
| 492 | |
| 493 | return ds.map_batches(add_window_and_source_batch, batch_format="pandas") |
| 494 | |
| 495 | entity_windowed = add_window_and_source( |
| 496 | entity_ds, timestamp_field, "entity", window_size or "1H" |
| 497 | ) |
| 498 | feature_windowed = add_window_and_source( |
| 499 | feature_ds, timestamp_field, "feature", window_size or "1H" |
| 500 | ) |
| 501 | combined_ds = entity_windowed.union(feature_windowed) |
| 502 | |
| 503 | def windowed_point_in_time_logic(batch: pd.DataFrame) -> pd.DataFrame: |
| 504 | if len(batch) == 0: |
| 505 | return pd.DataFrame() |
| 506 | result_chunks = [] |
| 507 | group_keys = ["time_window"] + join_keys |
| 508 | for group_values, group_data in batch.groupby(group_keys): |
| 509 | entity_data = group_data[group_data["_data_source"] == "entity"].copy() |
| 510 | feature_data = group_data[group_data["_data_source"] == "feature"].copy() |
| 511 | if len(entity_data) > 0 and len(feature_data) > 0: |
| 512 | entity_clean = entity_data.drop(columns=["time_window", "_data_source"]) |
| 513 | feature_clean = feature_data.drop( |
| 514 | columns=["time_window", "_data_source"] |
| 515 | ) |
| 516 | if join_keys: |
| 517 | merged = pd.merge_asof( |
| 518 | entity_clean.sort_values(join_keys + [timestamp_field]), |
| 519 | feature_clean.sort_values(join_keys + [timestamp_field]), |
| 520 | on=timestamp_field, |
| 521 | by=join_keys, |
| 522 | direction="backward", |
| 523 | ) |
| 524 | else: |
| 525 | merged = pd.merge_asof( |
no test coverage detected