From e2ea580595b27ce25cb99808f5c0753e1b3f6902 Mon Sep 17 00:00:00 2001 From: Dmitriy Vinogradov Date: Mon, 29 Jun 2026 09:47:23 +0300 Subject: [PATCH] feat: implement strict mode for max_records_per_run to prevent offset overshoot with composite keys --- .github/workflows/test.yaml | 2 + datapipe/index_metadata.py | 133 ++++ datapipe/meta/sql_meta.py | 18 +- datapipe/step/batch_transform.py | 140 +++- .../max_records_per_run_associativity_plan.md | 745 ++++++++++++++++++ ...atch_transform_with_offset_optimization.py | 15 +- tests/test_overshoot_composite_keys.py | 236 ++++++ 7 files changed, 1260 insertions(+), 29 deletions(-) create mode 100644 datapipe/index_metadata.py create mode 100644 docs/max_records_per_run_associativity_plan.md create mode 100644 tests/test_overshoot_composite_keys.py diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 9dca0fe4..e4bcbc5c 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -59,8 +59,10 @@ jobs: - python-version: "3.12" test-db-env: "postgres" + pip-extra: '"pandas<2.2"' - python-version: "3.12" test-db-env: "sqlite" + pip-extra: '"pandas<2.2"' services: # Label used to access the service container diff --git a/datapipe/index_metadata.py b/datapipe/index_metadata.py new file mode 100644 index 00000000..9e47b85b --- /dev/null +++ b/datapipe/index_metadata.py @@ -0,0 +1,133 @@ +""" +Wrapper для IndexDF с дополнительными метаданными. + +Используется для передачи исходного changed_idx (с полным PK) вместе +со сгруппированным idx (по transform_keys) для поддержки strict mode +при использовании max_records_per_run. +""" + +from typing import Optional + +from datapipe.types import IndexDF + + +class IndexDFWithMetadata: + """ + Wrapper для IndexDF с дополнительными метаданными. + + Позволяет передавать исходный changed_idx (с полным PK) вместе + со сгруппированным idx (по transform_keys) для поддержки strict mode + при использовании max_records_per_run. + + Ведёт себя как обычный IndexDF для обратной совместимости через делегирование. + + Example: + >>> idx_grouped = data_to_index(chunk_df, transform_keys) + >>> idx_with_meta = IndexDFWithMetadata( + ... idx=idx_grouped, + ... changed_idx=chunk_df # Исходный с полным PK + ... ) + >>> len(idx_with_meta) # Работает как обычный DataFrame + 3 + >>> idx_with_meta.has_changed_idx() + True + """ + + def __init__( + self, + idx: IndexDF, + changed_idx: Optional[IndexDF] = None + ): + """ + Args: + idx: Основной индекс (обычно сгруппированный по transform_keys) + changed_idx: Исходный индекс с полным PK из changed_idx query + (опционально, используется для strict mode) + """ + self._idx = idx + self._changed_idx = changed_idx + + @property + def idx(self) -> IndexDF: + """Основной индекс (сгруппированный по transform_keys)""" + return self._idx + + @property + def changed_idx(self) -> Optional[IndexDF]: + """Исходный индекс с полным PK (если есть)""" + return self._changed_idx + + def has_changed_idx(self) -> bool: + """Проверка наличия changed_idx""" + return self._changed_idx is not None + + # Делегирование к pandas DataFrame для обратной совместимости + def __getattr__(self, name): + """Делегируем все неизвестные атрибуты к idx""" + return getattr(self._idx, name) + + def __getitem__(self, key): + """Поддержка индексации df[key]""" + return self._idx[key] + + def __len__(self): + """len(idx_with_meta) возвращает длину основного индекса""" + return len(self._idx) + + def __repr__(self): + return ( + f"IndexDFWithMetadata(" + f"idx={len(self._idx)} rows, " + f"has_changed_idx={self.has_changed_idx()}" + f")" + ) + + def __str__(self): + return str(self._idx) + + # Поддержка итерации + def __iter__(self): + return iter(self._idx) + + # Поддержка pandas операций + @property + def columns(self): + return self._idx.columns + + @property + def shape(self): + return self._idx.shape + + def to_records(self, *args, **kwargs): + """Делегируем to_records к основному DataFrame""" + return self._idx.to_records(*args, **kwargs) + + +def unwrap_idx(idx_or_wrapped) -> tuple[IndexDF, Optional[IndexDF]]: + """ + Извлекает idx и changed_idx из обёртки или возвращает как есть. + + Этот хелпер позволяет функциям принимать как обычный IndexDF, + так и IndexDFWithMetadata без явной проверки типа. + + Args: + idx_or_wrapped: IndexDF или IndexDFWithMetadata + + Returns: + Кортеж (idx, changed_idx): + - idx: основной индекс (IndexDF) + - changed_idx: опциональный исходный индекс с полным PK + + Example: + >>> idx, changed_idx = unwrap_idx(some_idx) + >>> if changed_idx is not None: + ... # Используем strict mode + ... data = dt.get_data(changed_idx) + >>> else: + ... # Legacy mode + ... data = dt.get_data(idx) + """ + if isinstance(idx_or_wrapped, IndexDFWithMetadata): + return idx_or_wrapped.idx, idx_or_wrapped.changed_idx + else: + return idx_or_wrapped, None diff --git a/datapipe/meta/sql_meta.py b/datapipe/meta/sql_meta.py index b4cd9c76..53ecd04b 100644 --- a/datapipe/meta/sql_meta.py +++ b/datapipe/meta/sql_meta.py @@ -1642,6 +1642,7 @@ def _build_error_records_cte( transform_keys: List[str], all_select_keys: List[str], filters_idx: Optional[IndexDF], + input_dts: Optional[List["ComputeInput"]] = None, ) -> Any: """ Строит CTE для записей с ошибками из TransformMetaTable. @@ -1649,15 +1650,29 @@ def _build_error_records_cte( Returns: CTE object для error_records """ + # Получаем типы колонок из первой входной таблицы для дополнительных колонок + column_types = {} + if input_dts and len(input_dts) > 0: + first_input = input_dts[0] + if hasattr(first_input.dt, 'meta_table') and first_input.dt.meta_table: + # Используем sql_table.c для получения колонок + sql_table = first_input.dt.meta_table.sql_table + for col_name in all_select_keys: + if col_name in sql_table.c: + column_types[col_name] = sql_table.c[col_name].type + tr_tbl = meta_table.sql_table error_select_cols: List[Any] = [] for k in all_select_keys: if k in transform_keys: error_select_cols.append(sa.column(k)) else: - # Для дополнительных колонок (включая update_ts) используем NULL + # Для дополнительных колонок (включая update_ts) используем NULL с правильным типом if k == 'update_ts': error_select_cols.append(sa.cast(sa.literal(None), sa.Float).label(k)) + elif k in column_types: + # Используем тип из входной таблицы + error_select_cols.append(sa.cast(sa.literal(None), column_types[k]).label(k)) else: error_select_cols.append(sa.literal(None).label(k)) @@ -1923,6 +1938,7 @@ def build_changed_idx_sql_v2( transform_keys=transform_keys, all_select_keys=all_select_keys, filters_idx=filters_idx, + input_dts=input_dts, ) # 4. Объединить через UNION или CROSS JOIN diff --git a/datapipe/step/batch_transform.py b/datapipe/step/batch_transform.py index 26e5027f..ca14aebb 100644 --- a/datapipe/step/batch_transform.py +++ b/datapipe/step/batch_transform.py @@ -176,11 +176,14 @@ def _get_optimization_method_name(self, run_config: Optional[RunConfig] = None) def _get_additional_idx_columns(self) -> List[str]: """ - Собрать дополнительные колонки, необходимые для filtered join. + Собрать дополнительные колонки, необходимые для filtered join и strict mode. Возвращает список колонок из join_keys, которые нужно включить в idx для работы filtered join оптимизации. + Если max_records_per_run задан, также включает полный PK для поддержки + strict mode (чтобы можно было читать только конкретные записи, а не все события пар). + Returns: Список имен колонок (без дубликатов) """ @@ -194,6 +197,14 @@ def _get_additional_idx_columns(self) -> List[str]: if idx_col not in self.transform_keys and idx_col not in additional_columns: additional_columns.append(idx_col) + # Если max_records_per_run задан, добавляем полный PK для strict mode + if self.max_records_per_run is not None and len(self.input_dts) > 0: + # Берем primary_keys из первой входной таблицы + primary_keys = self.input_dts[0].dt.meta_table.primary_keys + for pk in primary_keys: + if pk not in self.transform_keys and pk not in additional_columns: + additional_columns.append(pk) + return additional_columns def _build_changed_idx_sql( @@ -402,7 +413,25 @@ def alter_res_df(): for k, v in extra_filters.items(): df[k] = v - yield cast(IndexDF, df) + # Если max_records_per_run задан, оборачиваем в IndexDFWithMetadata + if self.max_records_per_run is not None: + from datapipe.index_metadata import IndexDFWithMetadata + from datapipe.types import data_to_index + + # df содержит полный PK (благодаря _get_additional_idx_columns) + changed_idx_full_pk = cast(IndexDF, df) + + # Группируем по transform_keys для обратной совместимости + idx_grouped = data_to_index(df, self.transform_keys) + + # Оборачиваем + yield IndexDFWithMetadata( + idx=idx_grouped, + changed_idx=changed_idx_full_pk + ) + else: + # Legacy: просто группируем по transform_keys + yield cast(IndexDF, df) query_exec_time = time.time() - start_time logger.debug( @@ -478,6 +507,7 @@ def _get_max_update_ts_for_batch( ds: DataStore, compute_input: "ComputeInput", processed_idx: IndexDF, + changed_idx: Optional[IndexDF] = None, ) -> Optional[float]: """ Получить максимальный update_ts из входной таблицы для УСПЕШНО обработанного батча. @@ -485,6 +515,9 @@ def _get_max_update_ts_for_batch( Важно: используем processed_idx который содержит только успешно обработанные записи из output_dfs (result.index), а не весь батч idx. + Если передан changed_idx (при max_records_per_run), используем его для вычисления offset + вместо запроса к БД. Это предотвращает overshoot когда transform_keys не совпадают с PK. + Для JoinSpec таблиц (с join_keys) используем join_keys для фильтрации вместо primary_keys. Пример: profiles с join_keys={'user_id': 'id'} фильтруется по profiles.id IN (processed_idx.user_id) """ @@ -493,6 +526,38 @@ def _get_max_update_ts_for_batch( if len(processed_idx) == 0: return None + # НОВОЕ: Если есть changed_idx - используем его для вычисления offset (strict mode) + # Это предотвращает overshoot при composite keys с max_records_per_run + if changed_idx is not None and 'update_ts' in changed_idx.columns: + # Фильтруем changed_idx: только успешно обработанные пары + # Merge по transform_keys чтобы оставить только те записи changed_idx, + # для которых пары были успешно обработаны + import pandas as pd + + merged = changed_idx.merge( + processed_idx, + on=self.transform_keys, + how='inner' + ) + + if len(merged) == 0: + return None + + max_update = merged['update_ts'].max() + + # Учитываем delete_ts если есть + if 'delete_ts' in merged.columns: + max_delete = merged['delete_ts'].max() + if pd.notna(max_delete) and max_delete > max_update: + max_update = max_delete + + logger.info( + f"[{self.get_name()}] Offset from changed_idx (strict mode): {max_update}, " + f"input: {compute_input.dt.name}, records: {len(merged)}" + ) + + return float(max_update) + input_dt = compute_input.dt tbl = input_dt.meta_table.sql_table @@ -551,6 +616,11 @@ def store_batch_result( process_ts: float, run_config: Optional[RunConfig] = None, ) -> ChangeList: + from datapipe.index_metadata import unwrap_idx + + # Извлекаем idx и changed_idx если есть + idx, changed_idx = unwrap_idx(idx) + run_config = self._apply_filters_to_run_config(run_config) changes = ChangeList() @@ -615,7 +685,7 @@ def store_batch_result( for inp in self.input_dts: # Найти максимальный update_ts из УСПЕШНО обработанного батча - max_update_ts = self._get_max_update_ts_for_batch(ds, inp, processed_idx) + max_update_ts = self._get_max_update_ts_for_batch(ds, inp, processed_idx, changed_idx) if max_update_ts is not None: offset_key = (self.get_name(), inp.dt.name) @@ -685,11 +755,37 @@ def get_batch_input_dfs( Если у ComputeInput указаны join_keys, читаем только связанные записи для оптимизации производительности. + + Если max_records_per_run задан и idx содержит метаданные (IndexDFWithMetadata), + используется strict mode: читаются только конкретные записи из changed_idx + вместо всех событий для пар. """ + from datapipe.index_metadata import unwrap_idx + + # Извлекаем idx и changed_idx если есть + idx, changed_idx = unwrap_idx(idx) + + # Strict mode: если max_records_per_run задан и есть changed_idx + use_strict = ( + self.max_records_per_run is not None and + changed_idx is not None + ) + result = [] for inp in self.input_dts: - if inp.join_keys: + if use_strict: + # STRICT MODE: Читаем ТОЛЬКО конкретные записи из changed_idx + # Используется когда max_records_per_run задан для предотвращения overshoot. + # changed_idx содержит полный PK ограниченного батча, что позволяет + # читать именно те записи которые попали в LIMIT запроса. + assert changed_idx is not None # for mypy + data = inp.dt.get_data(changed_idx) + logger.debug( + f"[{self.get_name()}] Strict mode: reading {len(changed_idx)} records " + f"for input {inp.dt.name} (max_records_per_run={self.max_records_per_run})" + ) + elif inp.join_keys: # FILTERED JOIN: Читаем только связанные записи # Для composite key сохраняем реальные пары/кортежи ключей из idx, # а не независимые множества значений по каждой колонке. @@ -824,29 +920,25 @@ def run_full( # 1. "Полудвиги" окна при падении mid-batch # 2. Порчу offset'ов при вызове run_changelist (он вообще не должен трогать offset'ы) # 3. Потерю offset'ов в RayExecutor (т.к. они теперь возвращаются через результат) - # 4. Потерю данных при срабатывании max_records_per_run (offset прыгает, срезанные записи теряются) if changes.offsets: - if limit_hit: - # max_records_per_run лимит сработал - НЕ обновляем offset - # Это предотвращает потерю данных когда offset вычисляется по всем событиям - # обработанных пар (composite keys), а не только по событиям из батча + # Обновляем offset (даже если limit_hit, т.к. strict mode предотвращает overshoot) + try: + ds.offset_table.update_offsets_bulk(changes.offsets) + if limit_hit: + logger.info( + f"Updated offsets for {self.get_name()}: {changes.offsets}. " + f"Processed {idx_count} of {actual_changed_count} records " + f"(max_records_per_run={self.max_records_per_run})." + ) + else: + logger.info(f"Updated offsets for {self.get_name()}: {changes.offsets}") + except Exception as e: + # Таблица offset'ов может не существовать (create_meta_table=False) + # Логируем warning но не прерываем выполнение logger.warning( - f"[{self.get_name()}] Skipping offset update: max_records_per_run limit hit. " - f"Processed {idx_count} of {actual_changed_count} records. " - f"Offset will be updated after all records are processed." + f"Failed to update offsets for {self.get_name()}: {e}. " + "Offset table may not exist (create_meta_table=False)" ) - else: - # Обычный путь - обновляем offset - try: - ds.offset_table.update_offsets_bulk(changes.offsets) - logger.info(f"Updated offsets for {self.get_name()}: {changes.offsets}") - except Exception as e: - # Таблица offset'ов может не существовать (create_meta_table=False) - # Логируем warning но не прерываем выполнение - logger.warning( - f"Failed to update offsets for {self.get_name()}: {e}. " - "Offset table may not exist (create_meta_table=False)" - ) else: # Диагностика: логируем если offset не был вычислен ни для одного input logger.warning( diff --git a/docs/max_records_per_run_associativity_plan.md b/docs/max_records_per_run_associativity_plan.md new file mode 100644 index 00000000..4398fd97 --- /dev/null +++ b/docs/max_records_per_run_associativity_plan.md @@ -0,0 +1,745 @@ +# План реализации: Исправление overshoot для max_records_per_run + +## Проблема + +**Текущее поведение:** +- `max_records_per_run=5` ограничивает changed_idx query → 5 событий +- Эти события группируются по `transform_keys` → N пар +- `get_batch_input_dfs` читает **ВСЕ события** этих пар из БД (может быть 100+) +- Offset вычисляется по всем прочитанным данным → **overshoot** + +**Пример:** +``` +БД содержит: + (p1, post1): события id=1,2,3 (old_ts), id=6,7 (middle_ts), id=9,10,11 (new_ts) + (p2, post2): id=4,5 (old_ts) + (p3, post3): id=8 (middle_ts) + + id | profile | post | ts + ----+---------+-------+---- + 1 | p1 | post1 | old + 2 | p1 | post1 | old + 3 | p1 | post1 | old + 4 | p2 | post2 | old + 5 | p2 | post2 | old <-- LIMIT 5 + 6 | p1 | post1 | middle + 7 | p1 | post1 | middle + 8 | p3 | post3 | middle + 9 | p1 | post1 | new + 10 | p1 | post1 | new + 11 | p1 | post1 | new + +max_records_per_run=5: + changed_idx query LIMIT 5 → id=1-5 + Пары: (p1,post1), (p2,post2) + get_batch_input_dfs читает: id=1,2,3,4,5,6,7,9,10,11 (10 событий вместо 5!) + offset = new_ts → overshoot! + +Следующий запрос: + WHERE update_ts > new_ts → пропускает id=6,7,8 ❌ +``` + +## Корневая причина + +Для **ассоциативных** трансформаций (где `T(A) + T(B) = T(A+B)`) должно выполняться: +- Обработка id=1-5, затем id=6-11 = Обработка id=1-11 за раз + +Но `get_batch_input_dfs` читает "лишние" данные (id=6,7,9,10,11), которые **ещё не должны быть обработаны**. + +**Важно:** Мы предполагаем, что если разработчик использует `max_records_per_run`, то его трансформация ассоциативна. Проверку ассоциативности пока НЕ делаем. + +## Решение + +### Принцип: Имитация инкрементального поступления данных + +Когда используется `max_records_per_run`, нужно вести себя так, как будто данных после LIMIT **физически не существует**. + +--- + +## Часть 1: Strict режим для get_batch_input_dfs + +### 1.1. Текущее поведение + +```python +def get_batch_input_dfs(self, ds: DataStore, idx: IndexDF, ...) -> List[DataDF]: + """ + idx содержит transform_keys: [(p1, post1), (p2, post2)] + + Читает ВСЕ события этих пар: + SELECT * FROM events WHERE (profile_id, post_id) IN (...) + + Возвращает: все события пар (может быть намного больше чем в idx) + """ +``` + +### 1.2. Новое поведение (strict mode) + +Нужно добавить режим, где `get_batch_input_dfs` читает **только** конкретные записи из changed_idx. + +**Проблема:** `idx` после группировки по `transform_keys` теряет полный PK! + +**Пример:** +``` +Входная таблица PK: (id, profile_id, post_id) +transform_keys: (profile_id, post_id) + +changed_idx query вернул: + id=1, profile_id=p1, post_id=post1 + id=2, profile_id=p1, post_id=post1 + id=3, profile_id=p1, post_id=post1 + id=4, profile_id=p2, post_id=post2 + id=5, profile_id=p2, post_id=post2 + +idx передаваемый в process_batch (после группировки): + profile_id=p1, post_id=post1 ← потеряли id=1,2,3! + profile_id=p2, post_id=post2 ← потеряли id=4,5! +``` + +### 1.3. Решение: Сохранять исходный changed_idx + +Нужно передавать **два** индекса: +1. **`idx`** - после группировки по transform_keys (для обратной совместимости) +2. **`changed_idx`** - исходный с полным PK (для strict режима) + +**Способ передачи:** Через `RunConfig.labels` (минимальные breaking changes) + +```python +# В run_full +for chunk_df in idx_gen: + changed_idx = chunk_df # Полный PK + idx = data_to_index(chunk_df, self.transform_keys) # Группировка + + # Сохраняем changed_idx в run_config + run_config_with_changed = RunConfig.add_labels( + run_config, + {"_datapipe_changed_idx_full_pk": changed_idx} + ) + + changes = self.process_batch(ds, idx, run_config_with_changed) +``` + +--- + +## Часть 2: Изменения в get_batch_input_dfs + +```python +def get_batch_input_dfs( + self, + ds: DataStore, + idx: IndexDF, + run_config: Optional[RunConfig] = None, +) -> List[DataDF]: + # Извлекаем changed_idx если есть + changed_idx = None + if run_config and "_datapipe_changed_idx_full_pk" in run_config.labels: + changed_idx = run_config.labels["_datapipe_changed_idx_full_pk"] + + # Strict mode: если max_records_per_run задан и есть changed_idx + use_strict = ( + self.max_records_per_run is not None and + changed_idx is not None + ) + + result = [] + for inp in self.input_dts: + if use_strict: + # Читаем ТОЛЬКО записи из changed_idx (по полному PK) + data = inp.dt.get_data(changed_idx) + logger.info( + f"[{self.get_name()}] Strict mode: reading {len(changed_idx)} records " + f"(limited by max_records_per_run={self.max_records_per_run})" + ) + else: + # Legacy: читаем все данные для пар из idx + if inp.join_keys: + # Filtered join logic... + data = inp.dt.get_data(filtered_idx) + else: + data = inp.dt.get_data(idx) + + result.append(data) + + return result +``` + +--- + +## Часть 3: Изменения в offset calculation + +### 3.1. Текущая логика offset + +```python +def store_batch_result(..., idx: IndexDF, ...): + # idx = данные после get_batch_input_dfs (может содержать overshoot) + processed_idx = data_to_index(first_output, self.transform_keys) + max_update_ts = self._get_max_update_ts_for_batch(ds, inp, processed_idx) + # Запрос: SELECT MAX(update_ts) WHERE (transform_keys) IN (processed_idx) + # → overshoot! +``` + +### 3.2. Новая логика + +```python +def store_batch_result( + self, + ds: DataStore, + idx: IndexDF, + output_dfs: Optional[TransformResult], + process_ts: float, + run_config: Optional[RunConfig] = None, +) -> ChangeList: + # Извлекаем changed_idx + changed_idx = None + if run_config and "_datapipe_changed_idx_full_pk" in run_config.labels: + changed_idx = run_config.labels["_datapipe_changed_idx_full_pk"] + + # ... existing code для формирования processed_idx ... + + for inp in self.input_dts: + max_update_ts = self._get_max_update_ts_for_batch( + ds, inp, processed_idx, changed_idx + ) + ... +``` + +### 3.3. Обновленный _get_max_update_ts_for_batch + +```python +def _get_max_update_ts_for_batch( + self, + ds: DataStore, + compute_input: "ComputeInput", + processed_idx: IndexDF, + changed_idx: Optional[IndexDF] = None, +) -> Optional[float]: + if len(processed_idx) == 0: + return None + + # Если есть changed_idx - используем его для вычисления offset + if changed_idx is not None and 'update_ts' in changed_idx.columns: + # Фильтруем changed_idx: только успешно обработанные пары + merged = changed_idx.merge( + processed_idx, + on=self.transform_keys, + how='inner' + ) + + if len(merged) == 0: + return None + + max_update = merged['update_ts'].max() + + logger.info( + f"[{self.get_name()}] Offset from changed_idx: {max_update} " + f"({len(merged)} records, max_records_per_run={self.max_records_per_run})" + ) + + return float(max_update) + + # Fallback: текущая логика (запрос к БД) + input_dt = compute_input.dt + tbl = input_dt.meta_table.sql_table + # ... existing code ... +``` + +--- + +## Часть 4: Убрать limit_hit guard + +Теперь можно безопасно обновлять offset даже при `limit_hit=True`: +- Offset вычисляется по changed_idx (не overshoot) +- Strict mode гарантирует, что обработали только данные из changed_idx + +```python +# В run_full, УДАЛИТЬ: +# if limit_hit: +# logger.warning("Skipping offset update: max_records_per_run limit hit") +# else: +# ds.offset_table.update_offsets_bulk(changes.offsets) + +# ЗАМЕНИТЬ на: +if changes.offsets: + try: + ds.offset_table.update_offsets_bulk(changes.offsets) + logger.info(f"Updated offsets for {self.get_name()}: {changes.offsets}") + except Exception as e: + logger.warning(f"Failed to update offsets: {e}") +``` + +--- + +## Часть 5: Итоговый workflow + +### 5.1. Без max_records_per_run (текущее поведение) + +``` +1. changed_idx query → все новые записи +2. Группировка по transform_keys → idx +3. get_batch_input_dfs(idx) → читает все данные пар +4. transform(данные) +5. offset = MAX(update_ts) из БД для обработанных пар +``` + +### 5.2. С max_records_per_run (новое поведение) + +``` +1. changed_idx query LIMIT N → первые N записей (с полным PK) +2. Сохраняем changed_idx в RunConfig +3. Группировка по transform_keys → idx +4. get_batch_input_dfs(idx, changed_idx) → читает ТОЛЬКО changed_idx (strict mode) +5. transform(данные из changed_idx) +6. offset = MAX(update_ts) из changed_idx для обработанных пар +``` + +--- + +## Часть 6: Детальный план реализации + +### Этап 1: Сохранить changed_idx с полным PK + +**Файл:** `datapipe/step/batch_transform.py` → `run_full()` + +**Изменение:** +```python +# Текущий код в run_full: +for idx in idx_gen: + changes = executor.run_process_batch( + ..., + idx_gen=iter([idx]), + ... + ) + +# Новый код: +for chunk_df in idx_gen: + # НЕ группируем сразу! Сохраняем полный PK + changed_idx_full_pk = chunk_df + + # Группируем для обратной совместимости + idx = data_to_index(chunk_df, self.transform_keys) + + # Передаем через run_config + if self.max_records_per_run is not None: + run_config = RunConfig.add_labels( + run_config, + {"_datapipe_changed_idx_full_pk": changed_idx_full_pk} + ) + + changes = executor.run_process_batch( + ..., + idx_gen=iter([idx]), + run_config=run_config, + ... + ) +``` + +**Проблема:** Нужно проверить где происходит группировка в текущем коде. + +--- + +### Этап 2: Strict mode в get_batch_input_dfs + +**Файл:** `datapipe/step/batch_transform.py` → `get_batch_input_dfs()` + +```python +def get_batch_input_dfs( + self, + ds: DataStore, + idx: IndexDF, + run_config: Optional[RunConfig] = None, +) -> List[DataDF]: + # Извлекаем changed_idx если есть + changed_idx = None + if run_config and "_datapipe_changed_idx_full_pk" in run_config.labels: + changed_idx = run_config.labels["_datapipe_changed_idx_full_pk"] + + # Strict mode: если max_records_per_run задан и есть changed_idx + use_strict = ( + self.max_records_per_run is not None and + changed_idx is not None + ) + + result = [] + for inp in self.input_dts: + if use_strict: + # Читаем ТОЛЬКО записи из changed_idx (по полному PK) + data = inp.dt.get_data(changed_idx) + logger.debug( + f"[{self.get_name()}] Strict mode: reading {len(changed_idx)} records " + f"for input {inp.dt.name} (max_records_per_run={self.max_records_per_run})" + ) + else: + # Legacy: читаем все данные для пар из idx + if inp.join_keys: + # Существующий код для filtered join + filtered_idx = IndexDF( + idx[idx_cols] + .dropna() + .drop_duplicates() + .rename(columns=inp.join_keys) + .reset_index(drop=True) + ) + data = inp.dt.get_data(filtered_idx) + else: + data = inp.dt.get_data(idx) + + result.append(data) + + return result +``` + +--- + +### Этап 3: Offset calculation с changed_idx + +**Файл:** `datapipe/step/batch_transform.py` → `store_batch_result()` + +```python +def store_batch_result( + self, + ds: DataStore, + idx: IndexDF, + output_dfs: Optional[TransformResult], + process_ts: float, + run_config: Optional[RunConfig] = None, +) -> ChangeList: + # Извлекаем changed_idx + changed_idx = None + if run_config and "_datapipe_changed_idx_full_pk" in run_config.labels: + changed_idx = run_config.labels["_datapipe_changed_idx_full_pk"] + + # ... existing code для формирования processed_idx ... + + for inp in self.input_dts: + max_update_ts = self._get_max_update_ts_for_batch( + ds, inp, processed_idx, changed_idx + ) + + if max_update_ts is not None: + offset_key = (self.get_name(), inp.dt.name) + changes.offsets[offset_key] = max_update_ts +``` + +**Файл:** `datapipe/step/batch_transform.py` → `_get_max_update_ts_for_batch()` + +```python +def _get_max_update_ts_for_batch( + self, + ds: DataStore, + compute_input: "ComputeInput", + processed_idx: IndexDF, + changed_idx: Optional[IndexDF] = None, +) -> Optional[float]: + if len(processed_idx) == 0: + return None + + # НОВАЯ ЛОГИКА: Если есть changed_idx - используем его + if changed_idx is not None and 'update_ts' in changed_idx.columns: + # Фильтруем changed_idx: только успешно обработанные пары + merged = changed_idx.merge( + processed_idx, + on=self.transform_keys, + how='inner' + ) + + if len(merged) == 0: + return None + + max_update = merged['update_ts'].max() + + # Учитываем delete_ts если есть + if 'delete_ts' in merged.columns: + max_delete = merged['delete_ts'].max() + if pd.notna(max_delete) and max_delete > max_update: + max_update = max_delete + + logger.info( + f"[{self.get_name()}] Offset from changed_idx (strict mode): {max_update}, " + f"input: {compute_input.dt.name}, records: {len(merged)}" + ) + + return float(max_update) + + # СТАРАЯ ЛОГИКА (fallback): Запрос к БД + input_dt = compute_input.dt + tbl = input_dt.meta_table.sql_table + + # ... existing code с SELECT MAX(update_ts) ... +``` + +--- + +### Этап 4: Убрать limit_hit guard + +**Файл:** `datapipe/step/batch_transform.py` → `run_full()` + +```python +# УДАЛИТЬ эти строки (после 896): +# if limit_hit: +# logger.warning( +# f"[{self.get_name()}] Skipping offset update: max_records_per_run limit hit. " +# ... +# ) +# else: +# try: +# ds.offset_table.update_offsets_bulk(changes.offsets) +# ... + +# ЗАМЕНИТЬ на: +if changes.offsets: + try: + ds.offset_table.update_offsets_bulk(changes.offsets) + logger.info(f"Updated offsets for {self.get_name()}: {changes.offsets}") + except Exception as e: + logger.warning( + f"Failed to update offsets for {self.get_name()}: {e}. " + "Offset table may not exist (create_meta_table=False)" + ) + +# Также можно убрать вычисление limit_hit или оставить для логирования +``` + +--- + +## Часть 7: Тестирование + +### Тест 1: max_records_per_run с composite keys + +```python +def test_max_records_per_run_no_overshoot(dbconn): + """ + Проверяет, что max_records_per_run: + 1. Обрабатывает только changed_idx (strict mode) + 2. Offset не делает overshoot + 3. Все данные в итоге обрабатываются + """ + ds = DataStore(dbconn, create_meta_table=True) + + # ... setup как в test_overshoot_composite_keys ... + + step = BatchTransformStep( + ds=ds, + name="parse_post_deepview", + func=parse_post_deepview, + input_dts=[ComputeInput(dt=input_dt)], # БЕЗ join_type="full" + output_dts=[output_dt], + transform_keys=["profile_id", "post_id"], + use_offset_optimization=True, + max_records_per_run=5, + ) + + # Первый запуск + step.run_full(ds) + + offset1 = ds.offset_table.get_offsets_for_transformation(step.get_name())["amplitude_events"] + + # Offset должен быть old_ts (не overshoot!) + assert offset1 <= old_ts + 1, f"Offset {offset1} overshoot! Expected <= {old_ts}" + + # Второй запуск + step.run_full(ds) + + offset2 = ds.offset_table.get_offsets_for_transformation(step.get_name())["amplitude_events"] + assert offset2 <= middle_ts + 1 + + # Третий запуск + step.run_full(ds) + + offset3 = ds.offset_table.get_offsets_for_transformation(step.get_name())["amplitude_events"] + assert offset3 <= new_ts + 1 + + # Проверяем финальный результат + output = output_dt.get_data() + assert len(output) == 3 # Все 3 пары + + # Проверяем что все события обработались + assert output[output['profile_id'] == 'p1']['deepview_count'].iloc[0] == 8 + assert output[output['profile_id'] == 'p2']['deepview_count'].iloc[0] == 2 + assert output[output['profile_id'] == 'p3']['deepview_count'].iloc[0] == 1 +``` + +### Тест 2: Обратная совместимость + +```python +def test_backward_compatibility_without_max_records(dbconn): + """ + Проверяет, что без max_records_per_run всё работает как раньше + """ + step = BatchTransformStep( + ..., + # max_records_per_run=None - не указан + ) + + step.run_full(ds) + + # Проверяем что всё обработалось за один запуск + output = output_dt.get_data() + assert len(output) == 3 +``` + +### Тест 3: Strict mode действительно ограничивает данные + +```python +def test_strict_mode_limits_data(dbconn): + """ + Проверяет что в strict mode get_batch_input_dfs читает + только changed_idx, а не все события пар + """ + # ... setup ... + + # Мокаем get_data чтобы проверить что передаётся + original_get_data = input_dt.get_data + calls = [] + + def mock_get_data(idx): + calls.append(len(idx)) + return original_get_data(idx) + + input_dt.get_data = mock_get_data + + step = BatchTransformStep( + ..., + max_records_per_run=5, + ) + + step.run_full(ds) + + # Проверяем что get_data вызывался с 5 записями, а не с 11 + assert calls[0] == 5, f"Expected 5 records, got {calls[0]}" +``` + +--- + +## Часть 8: Edge cases + +### 8.1. Что если transform функция отфильтровала часть данных? + +```python +# changed_idx содержит 5 записей (3 пары) +# transform вернул только 2 пары (1 отфильтровалась) + +# offset считается только по обработанным парам: +processed_idx = data_to_index(output, transform_keys) # 2 пары +merged = changed_idx.merge(processed_idx, on=transform_keys) # 2 пары из changed_idx +offset = max(merged['update_ts']) # ✅ Правильно +``` + +### 8.2. Что если входная таблица НЕ имеет update_ts? + +```python +if changed_idx is not None and 'update_ts' not in changed_idx.columns: + logger.warning( + f"[{self.get_name()}] changed_idx doesn't have update_ts, " + "falling back to DB query for offset calculation" + ) + changed_idx = None # Fallback на старую логику +``` + +### 8.3. Что если несколько input_dts? + +```python +# changed_idx относится к результату JOIN +# Если есть _datapipe_input_source - фильтруем по ней + +if '_datapipe_input_source' in changed_idx.columns: + filtered = changed_idx[ + changed_idx['_datapipe_input_source'] == compute_input.dt.name + ] +else: + filtered = changed_idx + +# Затем вычисляем offset по filtered +``` + +### 8.4. Что если join_type="full"? + +С `join_type="full"` параметр `max_records_per_run` не имеет смысла (читаем всю таблицу). + +Можно добавить предупреждение: +```python +if self.max_records_per_run and any(inp.join_type == "full" for inp in self.input_dts): + logger.warning( + f"[{self.get_name()}] max_records_per_run has no effect with join_type='full'. " + "Consider using filtered join or removing max_records_per_run." + ) +``` + +--- + +## Часть 9: Документация + +### Обновить docstring BatchTransformStep + +```python +""" +... + +Parameters: + max_records_per_run: Optional[int] + Limit the number of records processed per run. This is useful for: + - Rate limiting external API calls + - Memory management for large datasets + - Incremental processing in development + + IMPORTANT: This parameter assumes your transformation is associative, + meaning T(A) + T(B) = T(A+B) for any disjoint data subsets A and B. + + When set, the step will: + 1. Use strict mode: read ONLY the limited records (not all events of pairs) + 2. Calculate offset based on processed records (prevents overshoot) + + Examples of associative transformations: + - Filtering: df[df['value'] > 10] + - Mapping: df['new'] = df['old'] * 2 + - Simple aggregations: groupby(['key']).agg({'col': 'sum'}) + + Examples of NON-associative transformations (don't use max_records_per_run): + - Sorting: df.sort_values() + - Ranking: df['rank'] = df['value'].rank() + - Top-N: df.head(10) + - Window functions with full window + + If your transformation is non-associative, do not use this parameter. +""" +``` + +--- + +## Итоговый чеклист реализации + +- [ ] **Этап 1:** Сохранить changed_idx с полным PK + - [ ] Найти где происходит группировка по transform_keys + - [ ] Сохранить полный PK до группировки + - [ ] Передать через RunConfig.labels["_datapipe_changed_idx_full_pk"] + +- [ ] **Этап 2:** Strict mode в `get_batch_input_dfs` + - [ ] Извлекать changed_idx из RunConfig + - [ ] Условие: `max_records_per_run is not None and changed_idx is not None` + - [ ] Читать `inp.dt.get_data(changed_idx)` вместо `idx` + - [ ] Добавить логирование + +- [ ] **Этап 3:** Offset calculation с changed_idx + - [ ] Передать changed_idx в `_get_max_update_ts_for_batch` + - [ ] Merge changed_idx с processed_idx + - [ ] MAX(update_ts) из merged + - [ ] Fallback на старую логику если changed_idx=None + +- [ ] **Этап 4:** Убрать limit_hit guard + - [ ] Удалить проверку `if limit_hit` в `run_full` + - [ ] Всегда обновлять offset + +- [ ] **Этап 5:** Тестирование + - [ ] Тест: max_records_per_run → offset правильный (не overshoot) + - [ ] Тест: обратная совместимость (без max_records_per_run) + - [ ] Тест: strict mode ограничивает данные + - [ ] Тест: фильтрация в transform + - [ ] Обновить существующий test_overshoot_composite_keys + +- [ ] **Этап 6:** Документация + - [ ] Обновить docstring BatchTransformStep + - [ ] Добавить примеры ассоциативных/неассоциативных трансформаций + - [ ] Предупреждение о требовании ассоциативности + +- [ ] **Этап 7:** Очистка + - [ ] Убрать отладочные print из batch_transform.py + - [ ] Проверить что все тесты проходят + - [ ] Убрать старые документы (solution4_implementation.py, offset_bug_solutions_v2.md) diff --git a/tests/test_batch_transform_with_offset_optimization.py b/tests/test_batch_transform_with_offset_optimization.py index a1601c86..49aaeea2 100644 --- a/tests/test_batch_transform_with_offset_optimization.py +++ b/tests/test_batch_transform_with_offset_optimization.py @@ -488,12 +488,19 @@ def transform_func(df): first_run_output = output_dt.get_data().sort_values("profile_id").reset_index(drop=True) assert first_run_output["profile_id"].tolist() == ["p1", "p2", "p3"] - # КРИТИЧЕСКАЯ ПРОВЕРКА: offset НЕ должен обновиться когда лимит сработал! + # КРИТИЧЕСКАЯ ПРОВЕРКА: со strict mode offset ДОЛЖЕН обновиться, + # но только до first_ts (не overshoot до second_ts)! # У нас 6 записей, но max_records_per_run=3, значит 3 остались необработанными offsets_after_first = ds.offset_table.get_offsets_for_transformation(step.get_name()) - # Offset не должен быть записан вообще (первый запуск с лимитом) - assert "test_input_max_records" not in offsets_after_first, ( - "Offset не должен обновляться при срабатывании max_records_per_run лимита" + # С strict mode: offset обновляется всегда, но без overshoot + assert "test_input_max_records" in offsets_after_first, ( + "Со strict mode offset должен обновляться даже при срабатывании max_records_per_run" + ) + offset_after_first = offsets_after_first["test_input_max_records"] + # Офсет должен быть first_ts (записи p1-p3), а не second_ts (записи p5-p6) + assert abs(offset_after_first - first_ts) < 0.1, ( + f"Offset должен быть {first_ts}, но был {offset_after_first}. " + "Strict mode должен предотвращать overshoot!" ) # Проверяем что в очереди остались необработанные записи assert step.get_changed_idx_count(ds) == 3, "Должны остаться 3 необработанные записи" diff --git a/tests/test_overshoot_composite_keys.py b/tests/test_overshoot_composite_keys.py new file mode 100644 index 00000000..0bea6cf5 --- /dev/null +++ b/tests/test_overshoot_composite_keys.py @@ -0,0 +1,236 @@ +""" +Failing test для демонстрации overshoot проблемы с composite keys +""" +import time +import pandas as pd +from sqlalchemy import Column, Integer, String, Float + +from datapipe.compute import ComputeInput +from datapipe.datatable import DataStore +from datapipe.step.batch_transform import BatchTransformStep +from datapipe.store.database import DBConn, TableStoreDB + + +def test_offset_overshoot_with_composite_keys(dbconn: DBConn): + """ + Демонстрирует проблему overshoot при использовании composite keys. + + Сценарий: + - PK входной таблицы: (id, profile_id, post_id) + - transform_keys: (profile_id, post_id) - без id + - Для пары (p1, post1) есть несколько событий с разными id + - max_records_per_run ограничивает батч + + Проблема: + - Офсет вычисляется по ВСЕМ событиям обработанных пар + - А не только по событиям из батча + - Результат: overshoot - офсет прыгает через необработанные данные + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Входная таблица с composite PK + input_store = TableStoreDB( + dbconn, + "amplitude_events", + [ + Column("id", Integer, primary_key=True), + Column("profile_id", String, primary_key=True), + Column("post_id", String, primary_key=True), + Column("event_type", String), + ], + create_table=True, + ) + input_dt = ds.create_table("amplitude_events", input_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "post_deepview", + [ + Column("profile_id", String, primary_key=True), + Column("post_id", String, primary_key=True), + Column("deepview_count", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("post_deepview", output_store) + + # Transform функция: считает deepview для каждой пары + def parse_post_deepview(df): + result = df.groupby(['profile_id', 'post_id']).agg({ + 'event_type': 'count' + }).reset_index() + result = result.rename(columns={'event_type': 'deepview_count'}) + return result + + step = BatchTransformStep( + ds=ds, + name="parse_post_deepview", + func=parse_post_deepview, + input_dts=[ComputeInput(dt=input_dt)], # БЕЗ join_type="full"! + output_dts=[output_dt], + transform_keys=["profile_id", "post_id"], # БЕЗ id! + use_offset_optimization=True, + max_records_per_run=5, # Ограничиваем батч + ) + + # ========== ДАННЫЕ ========== + + # Timestamp для разных батчей + old_ts = time.time() - 100 + middle_ts = old_ts + 10 + new_ts = middle_ts + 10 + + # БАТЧ 1: Старые события (update_ts = old_ts) + # 3 события для пары (p1, post1) - id=1,2,3 + # 2 события для пары (p2, post2) - id=4,5 + input_dt.store_chunk( + pd.DataFrame({ + "id": [1, 2, 3, 4, 5], + "profile_id": ["p1", "p1", "p1", "p2", "p2"], + "post_id": ["post1", "post1", "post1", "post2", "post2"], + "event_type": ["deepPostView"] * 5, + }), + now=old_ts, + ) + + # БАТЧ 2: Средние события (update_ts = middle_ts) + # 2 события для пары (p1, post1) - id=6,7 + # 1 событие для пары (p3, post3) - id=8 + input_dt.store_chunk( + pd.DataFrame({ + "id": [6, 7, 8], + "profile_id": ["p1", "p1", "p3"], + "post_id": ["post1", "post1", "post3"], + "event_type": ["deepPostView"] * 3, + }), + now=middle_ts, + ) + + # БАТЧ 3: Новые события (update_ts = new_ts) + # 3 события для пары (p1, post1) - id=9,10,11 + input_dt.store_chunk( + pd.DataFrame({ + "id": [9, 10, 11], + "profile_id": ["p1", "p1", "p1"], + "post_id": ["post1", "post1", "post1"], + "event_type": ["deepPostView"] * 3, + }), + now=new_ts, + ) + + print("\n" + "=" * 80) + print("ДАННЫЕ В amplitude_events:") + print("=" * 80) + all_data = input_dt.get_data().sort_values("id") + print(all_data.to_string()) + print(f"\nВсего событий: {len(all_data)}") + print(f"Для пары (p1, post1): {len(all_data[all_data['profile_id'] == 'p1'])} событий") + print(f" - id=1,2,3 с update_ts={old_ts}") + print(f" - id=6,7 с update_ts={middle_ts}") + print(f" - id=9,10,11 с update_ts={new_ts}") + + # ========== ПЕРВЫЙ ЗАПУСК ========== + print("\n" + "=" * 80) + print("ПЕРВЫЙ ЗАПУСК (max_records_per_run=5):") + print("=" * 80) + + # Проверяем changed_idx_count + changed_count = step.get_changed_idx_count(ds) + print(f"\nChanged idx count: {changed_count}") + + step.run_full(ds) + + # Проверяем что обработали + output = output_dt.get_data() + print(f"\nОбработанные пары:") + print(output.to_string()) + + # Проверяем офсет + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_value = offsets.get("amplitude_events") + + print(f"\nОфсет после первого запуска: {offset_value}") + if offset_value: + print(f"Офсет (datetime): {pd.Timestamp(offset_value, unit='s')}") + print(f"old_ts (datetime): {pd.Timestamp(old_ts, unit='s')}") + print(f"middle_ts (datetime): {pd.Timestamp(middle_ts, unit='s')}") + print(f"new_ts (datetime): {pd.Timestamp(new_ts, unit='s')}") + + # ========== ПРОВЕРКА ========== + print("\n" + "=" * 80) + print("ПРОВЕРКА:") + print("=" * 80) + + # Changed idx query вернул первые 5 событий (LIMIT 5): + # id=1,2,3,4,5 с update_ts=old_ts + # Это агрегируется в 2 пары: (p1, post1), (p2, post2) + + # Ожидаемое поведение (ПРАВИЛЬНО): + # offset = max(update_ts из батча) = old_ts + # Потому что в батче были только события с update_ts=old_ts + + # Текущее поведение (БАГ - OVERSHOOT): + # _get_max_update_ts_for_batch делает: + # SELECT MAX(update_ts) WHERE (profile_id, post_id) IN (('p1', 'post1'), ('p2', 'post2')) + # Для пары (p1, post1) возвращает MAX по ВСЕМ её событиям + # Включая id=9,10,11 с update_ts=new_ts которые НЕ были в батче! + # offset = new_ts (OVERSHOOT!) + + if offset_value: + # Проверяем наличие overshoot + has_overshoot = offset_value > old_ts + 1 + + if has_overshoot: + print(f"\n❌ БАГ ПОДТВЕРЖДЁН!") + print(f"Офсет = {offset_value}") + print(f"Это близко к new_ts = {new_ts}") + print(f"\nПроблема:") + print(f" - Changed idx query вернул события id=1-5 (update_ts={old_ts})") + print(f" - Но офсет взял MAX по ВСЕМ событиям пар (p1, post1) и (p2, post2)") + print(f" - Включая события id=9,10,11 (update_ts={new_ts}) которые НЕ были в батче!") + print(f"\nПоследствия:") + print(f" - События id=6,7,8 (update_ts={middle_ts}) ПОТЕРЯНЫ") + print(f" - Следующий changed_idx query: WHERE update_ts > {offset_value}") + print(f" - События с update_ts={middle_ts} не попадут в запрос!") + else: + print(f"\n✅ БАГ ИСПРАВЛЕН!") + print(f"Офсет = {offset_value}") + print(f"Это соответствует old_ts = {old_ts}") + print(f"\nПравильное поведение (strict mode):") + print(f" - Changed idx query вернул события id=1-5 (update_ts={old_ts})") + print(f" - Офсет вычислен только по этим 5 событиям: offset = {offset_value}") + print(f" - События id=6,7,8 с update_ts={middle_ts} будут обработаны в следующем запуске") + print(f" - Никакие данные не потеряны!") + + # Это failing тест - он должен упасть при overshoot + assert offset_value <= old_ts + 1, ( + f"Офсет {offset_value} слишком большой (overshoot)! " + f"Должен быть <= {old_ts + 1} (max update_ts из обработанного батча)" + ) + else: + print(f"\n✅ Офсет не обновился (limit_hit guard сработал)") + print(f"Это текущее поведение после коммита 10f7326") + print(f"Предотвращает overshoot, но создаёт deadlock") + + +if __name__ == "__main__": + # Для локального запуска + from datapipe.store.database import DBConn + import os + + db_env = os.getenv("TEST_DB_ENV", "sqlite") + if db_env == "sqlite": + try: + import pysqlite3 # noqa: F401 + DBCONNSTR = "sqlite+pysqlite3:///:memory:" + except ImportError: + DBCONNSTR = "sqlite+pysqlite:///:memory:" + dbconn = DBConn(DBCONNSTR, None) + else: + pg_host = os.getenv("POSTGRES_HOST", "localhost") + pg_port = os.getenv("POSTGRES_PORT", "5432") + DBCONNSTR = f"postgresql://postgres:password@{pg_host}:{pg_port}/postgres" + dbconn = DBConn(DBCONNSTR, "test") + + test_offset_overshoot_with_composite_keys(dbconn)