In-process offline write batcher for /push requests. - Buffers DataFrames per (push_source_name, allow_registry_cache, transform_on_write) - Flushes when either: * total rows in a buffer >= batch_size, or * time since last flush >= batch_interval_seconds - Flush run
| 963 | |
| 964 | |
| 965 | class OfflineWriteBatcher: |
| 966 | """ |
| 967 | In-process offline write batcher for /push requests. |
| 968 | |
| 969 | - Buffers DataFrames per (push_source_name, allow_registry_cache, transform_on_write) |
| 970 | - Flushes when either: |
| 971 | * total rows in a buffer >= batch_size, or |
| 972 | * time since last flush >= batch_interval_seconds |
| 973 | - Flush runs in a dedicated background thread so the HTTP event loop stays unblocked. |
| 974 | """ |
| 975 | |
| 976 | def __init__(self, store: "feast.FeatureStore", cfg: Any): |
| 977 | self._store = store |
| 978 | self._cfg = cfg |
| 979 | |
| 980 | # Buffers and timestamps keyed by batch key |
| 981 | self._buffers: DefaultDict[_OfflineBatchKey, List[pd.DataFrame]] = defaultdict( |
| 982 | list |
| 983 | ) |
| 984 | self._last_flush: DefaultDict[_OfflineBatchKey, float] = defaultdict(time.time) |
| 985 | self._inflight: Set[_OfflineBatchKey] = set() |
| 986 | |
| 987 | self._lock = threading.Lock() |
| 988 | self._stop_event = threading.Event() |
| 989 | |
| 990 | # Start background flusher thread |
| 991 | self._thread = threading.Thread( |
| 992 | target=self._run, name="offline_write_batcher", daemon=True |
| 993 | ) |
| 994 | self._thread.start() |
| 995 | |
| 996 | logger.debug( |
| 997 | "OfflineWriteBatcher initialized: batch_size=%s, batch_interval_seconds=%s", |
| 998 | getattr(cfg, "batch_size", None), |
| 999 | getattr(cfg, "batch_interval_seconds", None), |
| 1000 | ) |
| 1001 | |
| 1002 | # ---------- Public API ---------- |
| 1003 | |
| 1004 | def enqueue( |
| 1005 | self, |
| 1006 | push_source_name: str, |
| 1007 | df: pd.DataFrame, |
| 1008 | allow_registry_cache: bool, |
| 1009 | transform_on_write: bool, |
| 1010 | ) -> None: |
| 1011 | """ |
| 1012 | Enqueue a dataframe for offline write, grouped by push source + flags. |
| 1013 | Cheap and non-blocking; heavy I/O happens in background thread. |
| 1014 | """ |
| 1015 | key = _OfflineBatchKey( |
| 1016 | push_source_name=push_source_name, |
| 1017 | allow_registry_cache=allow_registry_cache, |
| 1018 | transform_on_write=transform_on_write, |
| 1019 | ) |
| 1020 | |
| 1021 | with self._lock: |
| 1022 | self._buffers[key].append(df) |