(ds, timestamp_field, source_marker, window_size)
| 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" |
no test coverage detected