Flush a single buffer. Extracts data under lock, then does I/O without lock.
(self, key: _OfflineBatchKey)
| 1113 | return dfs |
| 1114 | |
| 1115 | def _flush(self, key: _OfflineBatchKey) -> None: |
| 1116 | """ |
| 1117 | Flush a single buffer. Extracts data under lock, then does I/O without lock. |
| 1118 | """ |
| 1119 | while True: |
| 1120 | with self._lock: |
| 1121 | dfs = self._drain_locked(key) |
| 1122 | |
| 1123 | if not dfs: |
| 1124 | return |
| 1125 | |
| 1126 | batch_df = pd.concat(dfs, ignore_index=True) |
| 1127 | |
| 1128 | # NOTE: offline writes are currently synchronous only, so we call directly |
| 1129 | try: |
| 1130 | self._store.push( |
| 1131 | push_source_name=key.push_source_name, |
| 1132 | df=batch_df, |
| 1133 | allow_registry_cache=key.allow_registry_cache, |
| 1134 | to=PushMode.OFFLINE, |
| 1135 | transform_on_write=key.transform_on_write, |
| 1136 | ) |
| 1137 | except Exception: |
| 1138 | logger.exception("Error flushing offline batch for %s", key) |
| 1139 | with self._lock: |
| 1140 | self._buffers[key] = dfs + self._buffers[key] |
| 1141 | self._inflight.discard(key) |
| 1142 | return |
| 1143 | |
| 1144 | logger.debug( |
| 1145 | "Flushing offline batch for push_source=%s with %s rows", |
| 1146 | key.push_source_name, |
| 1147 | len(batch_df), |
| 1148 | ) |
| 1149 | |
| 1150 | with self._lock: |
| 1151 | self._last_flush[key] = time.time() |
| 1152 | self._inflight.discard(key) |
| 1153 | pending_rows = sum(len(d) for d in self._buffers.get(key, [])) |
| 1154 | should_flush = pending_rows >= self._cfg.batch_size |
| 1155 | |
| 1156 | if not should_flush: |
| 1157 | return |
| 1158 | |
| 1159 | logger.debug( |
| 1160 | "OfflineWriteBatcher size threshold reached for %s: %s rows", |
| 1161 | key, |
| 1162 | pending_rows, |
| 1163 | ) |