diff --git a/datapipe/cli.py b/datapipe/cli.py index ab5e2ab8..11918292 100644 --- a/datapipe/cli.py +++ b/datapipe/cli.py @@ -498,6 +498,73 @@ def migrate_transform_tables(ctx: click.Context, labels: str, name: str) -> None return migrations_v013.migrate_transform_tables(app, batch_transforms_steps) +@cli.command() +@click.option("--step", type=click.STRING, help="Step name to initialize offsets for (optional)") +@click.pass_context +def init_offsets(ctx, step: Optional[str]): + """ + Инициализировать таблицу offset'ов из существующих данных TransformMetaTable. + + Команда сканирует уже обработанные данные и устанавливает начальные значения offset'ов, + чтобы обеспечить плавную миграцию на оптимизацию через offset'ы (метод v2). + + Если указан --step, инициализирует только этот шаг. Иначе инициализирует + все экземпляры BatchTransformStep в пайплайне. + """ + from datapipe.meta.sql_meta import initialize_offsets_from_transform_meta + + app: DatapipeApp = ctx.obj["app"] + + # Collect all BatchTransformStep instances + transform_steps = [] + for compute_step in app.steps: + if isinstance(compute_step, BaseBatchTransformStep): + if step is None or compute_step.get_name() == step: + transform_steps.append(compute_step) + + if not transform_steps: + if step: + rprint(f"[red]Step '{step}' not found or is not a BatchTransformStep[/red]") + else: + rprint("[yellow]No BatchTransformStep instances found in pipeline[/yellow]") + return + + rprint(f"[cyan]Found {len(transform_steps)} transform step(s) to initialize[/cyan]") + + # Initialize offsets for each step + results = {} + for transform_step in transform_steps: + step_name = transform_step.get_name() + rprint(f"\n[cyan]Initializing offsets for: {step_name}[/cyan]") + + try: + offsets = initialize_offsets_from_transform_meta(app.ds, transform_step) + + if offsets: + rprint(f"[green]✓ Initialized {len(offsets)} offset(s):[/green]") + for input_name, offset_value in offsets.items(): + rprint(f" - {input_name}: {offset_value}") + results[step_name] = offsets + else: + rprint("[yellow]No offsets initialized (no processed data found)[/yellow]") + results[step_name] = {} + + except Exception as e: + rprint(f"[red]✗ Failed to initialize: {e}[/red]") + results[step_name] = {} + + # Summary + rprint("\n[cyan]═══ Summary ═══[/cyan]") + success_count = sum(1 for v in results.values() if v is not None and len(v) > 0) + empty_count = sum(1 for v in results.values() if v is not None and len(v) == 0) + failed_count = sum(1 for v in results.values() if v is None) + + rprint(f"[green]Successful: {success_count}[/green]") + rprint(f"[yellow]Empty (no data): {empty_count}[/yellow]") + if failed_count > 0: + rprint(f"[red]Failed: {failed_count}[/red]") + + try: entry_points = metadata.entry_points(group="datapipe.cli") # type: ignore except TypeError: diff --git a/datapipe/compute.py b/datapipe/compute.py index d5e8c314..2a10e7d1 100644 --- a/datapipe/compute.py +++ b/datapipe/compute.py @@ -85,6 +85,9 @@ class StepStatus: class ComputeInput: dt: DataTable join_type: Literal["inner", "full"] = "full" + # Filtered join optimization: mapping from idx columns to dt columns + # Example: {"user_id": "id"} means filter dt by dt.id IN (idx.user_id) + join_keys: Optional[Dict[str, str]] = None class ComputeStep: diff --git a/datapipe/datatable.py b/datapipe/datatable.py index d7c4a6c2..9dbd58ec 100644 --- a/datapipe/datatable.py +++ b/datapipe/datatable.py @@ -5,7 +5,7 @@ from opentelemetry import trace from datapipe.event_logger import EventLogger -from datapipe.meta.sql_meta import MetaTable +from datapipe.meta.sql_meta import MetaTable, TransformInputOffsetTable from datapipe.run_config import RunConfig from datapipe.store.database import DBConn from datapipe.store.table_store import TableStore @@ -170,6 +170,9 @@ def __init__( self.create_meta_table = create_meta_table + # Создать таблицу offset'ов (используем тот же флаг create_meta_table) + self.offset_table = TransformInputOffsetTable(meta_dbconn, create_table=create_meta_table) + def create_table(self, name: str, table_store: TableStore) -> DataTable: assert name not in self.tables diff --git a/datapipe/meta/sql_meta.py b/datapipe/meta/sql_meta.py index 767dca39..1fb945fa 100644 --- a/datapipe/meta/sql_meta.py +++ b/datapipe/meta/sql_meta.py @@ -1,4 +1,5 @@ import itertools +import logging import math import time from dataclasses import dataclass @@ -36,7 +37,14 @@ if TYPE_CHECKING: from datapipe.compute import ComputeInput from datapipe.datatable import DataStore + from datapipe.store.database import TableStoreDB +logger = logging.getLogger("datapipe.meta.sql_meta") + +# Эпсилон для offset optimization: смещаем offset на эту величину назад +# для захвата записей с одинаковыми timestamps при использовании строгого > +# Это предотвращает потерю данных (Hypothesis 1) при сохранении производительности +OFFSET_EPSILON_SECONDS = 0.01 # 10 миллисекунд TABLE_META_SCHEMA: List[sa.Column] = [ sa.Column("hash", sa.Integer), @@ -251,8 +259,16 @@ def get_changes_for_store_chunk( changed_meta_df - строки метаданных, которые нужно изменить """ + current_time = time.time() if now is None: - now = time.time() + now = current_time + elif now < current_time - 1.0: # Порог 1 секунда - игнорируем микросекундные различия + # Предупреждение: использование timestamp из прошлого может привести к потере данных + # при использовании offset optimization (Hypothesis 4: delayed records) + logger.warning( + f"store_chunk called with now={now:.2f} which is {current_time - now:.2f}s in the past. " + f"This may cause data loss with offset optimization if offset > now." + ) # получить meta по чанку existing_meta_df = self.get_metadata(hash_to_index(hash_df, self.primary_keys), include_deleted=True) @@ -707,7 +723,7 @@ def _make_agg_of_agg( return sql.cte(name=f"all__{agg_col}") -def build_changed_idx_sql( +def build_changed_idx_sql_v1( ds: "DataStore", meta_table: "TransformMetaTable", input_dts: List["ComputeInput"], @@ -716,15 +732,27 @@ def build_changed_idx_sql( order_by: Optional[List[str]] = None, order: Literal["asc", "desc"] = "asc", run_config: Optional[RunConfig] = None, # TODO remove + additional_columns: Optional[List[str]] = None, ) -> Tuple[Iterable[str], Any]: + """ + Args: + additional_columns: Дополнительные колонки для включения в результат (для filtered join) + """ + if additional_columns is None: + additional_columns = [] + + # Полный список колонок для SELECT (transform_keys + additional_columns) + all_select_keys = list(transform_keys) + additional_columns + all_input_keys_counts: Dict[str, int] = {} for col in itertools.chain(*[inp.dt.primary_schema for inp in input_dts]): all_input_keys_counts[col.name] = all_input_keys_counts.get(col.name, 0) + 1 inp_ctes = [] for inp in input_dts: + # Используем all_select_keys для включения дополнительных колонок keys, cte = inp.dt.meta_table.get_agg_cte( - transform_keys=transform_keys, + transform_keys=all_select_keys, filters_idx=filters_idx, run_config=run_config, ) @@ -732,7 +760,7 @@ def build_changed_idx_sql( agg_of_aggs = _make_agg_of_agg( ds=ds, - transform_keys=transform_keys, + transform_keys=all_select_keys, ctes=inp_ctes, agg_col="update_ts", ) @@ -757,12 +785,14 @@ def build_changed_idx_sql( else: # len(transform_keys) > 1: join_onclause_sql = sa.and_(*[agg_of_aggs.c[key] == out.c[key] for key in transform_keys]) + # Важно: Включаем все колонки (transform_keys + additional_columns) sql = ( sa.select( # Нам нужно выбирать хотя бы что-то, чтобы не было ошибки при # пустом transform_keys sa.literal(1).label("_datapipe_dummy"), - *[sa.func.coalesce(agg_of_aggs.c[key], out.c[key]).label(key) for key in transform_keys], + *[sa.func.coalesce(agg_of_aggs.c[key], out.c[key]).label(key) if key in transform_keys + else agg_of_aggs.c[key].label(key) for key in all_select_keys if key in agg_of_aggs.c], ) .select_from(agg_of_aggs) .outerjoin( @@ -797,4 +827,1340 @@ def build_changed_idx_sql( *[sa.asc(sa.column(k)) for k in order_by], out.c.priority.desc().nullslast(), ) - return (transform_keys, sql) + return (all_select_keys, sql) + + +# ============================================================================ +# OFFSET OPTIMIZATION V2 - IMPLEMENTATION +# ============================================================================ +# +# Эта секция содержит функции для построения SQL запросов с offset optimization. +# Использует UNION вместо FULL OUTER JOIN для объединения changed records. +# +# Структура функций (от низкого к высокому уровню): +# +# 1. БАЗОВЫЕ УТИЛИТЫ +# - Создание WHERE условий, фильтров, JOIN conditions +# +# 2. ФУНКЦИИ ПРИНЯТИЯ РЕШЕНИЯ +# - Определяют какой путь использовать (meta vs data) +# +# 3. ФУНКЦИИ ПОСТРОЕНИЯ CTE +# - Forward join: таблица изменилась сама +# - Reverse join: справочная таблица изменилась, находим зависимые записи +# +# 4. КООРДИНАТОРЫ +# - Объединяют CTE от всех таблиц, добавляют error records, сортировку +# +# 5. ГЛАВНАЯ ФУНКЦИЯ +# - build_changed_idx_sql_v2: entry point для offset optimization +# ============================================================================ + + +# Обратная совместимость: алиас для старой версии +def build_changed_idx_sql( + ds: "DataStore", + meta_table: "TransformMetaTable", + input_dts: List["ComputeInput"], + transform_keys: List[str], + filters_idx: Optional[IndexDF] = None, + order_by: Optional[List[str]] = None, + order: Literal["asc", "desc"] = "asc", + run_config: Optional[RunConfig] = None, +) -> Tuple[Iterable[str], Any]: + """ + Обёртка для обратной совместимости. По умолчанию использует v1 (старую версию). + """ + return build_changed_idx_sql_v1( + ds=ds, + meta_table=meta_table, + input_dts=input_dts, + transform_keys=transform_keys, + filters_idx=filters_idx, + order_by=order_by, + order=order, + run_config=run_config, + ) + + +# ---------------------------------------------------------------------------- +# 1. БАЗОВЫЕ УТИЛИТЫ +# ---------------------------------------------------------------------------- + +def _generate_unique_cte_name(table_name: str, suffix: str, usage_count: Dict[str, int]) -> str: + """ + Генерирует уникальное имя CTE для таблицы. + + При первом использовании: "{table_name}_{suffix}" + При повторном: "{table_name}_{suffix}_{N}" + + Args: + table_name: Имя таблицы + suffix: Суффикс для CTE (например "changes", "deleted") + usage_count: Словарь счётчиков использования (мутируется!) + + Returns: + Уникальное имя CTE + """ + if table_name not in usage_count: + usage_count[table_name] = 0 + return f"{table_name}_{suffix}" + else: + usage_count[table_name] += 1 + return f"{table_name}_{suffix}_{usage_count[table_name]}" + + +def _build_offset_where_clause(tbl: Any, offset: float) -> Any: + """ + Строит WHERE условие для фильтрации по offset. + + Применяет epsilon-сдвиг для захвата записей с timestamp близкими к offset. + Это предотвращает потерю данных при одинаковых update_ts (Hypothesis 1) + при сохранении производительности (offset остаётся = MAX(update_ts)). + + WHERE update_ts > (offset - epsilon) эквивалентно >= для практических целей, + но читает только новые записи из БД. + """ + # Вычитаем epsilon при использовании offset для захвата записей + # с update_ts близкими к offset (включая равные) + adjusted_offset = offset - OFFSET_EPSILON_SECONDS + + return sa.or_( + tbl.c.update_ts > adjusted_offset, + sa.and_( + tbl.c.delete_ts.isnot(None), + tbl.c.delete_ts > adjusted_offset + ) + ) + + +def _build_join_condition(conditions: List[Any]) -> Any: + """Объединяет условия JOIN через AND.""" + if len(conditions) == 0: + return None + if len(conditions) == 1: + return conditions[0] + return sa.and_(*conditions) + + +def _apply_sql_filters( + sql: Any, + keys: List[str], + filters_idx: Optional[IndexDF], + tbl: Any, + primary_keys: List[str], + run_config: Optional[RunConfig], +) -> Any: + """Применяет filters_idx и run_config фильтры к SQL запросу.""" + sql = sql_apply_filters_idx_to_subquery(sql, keys, filters_idx) + sql = sql_apply_runconfig_filter(sql, tbl, primary_keys, run_config) + return sql + + +# ---------------------------------------------------------------------------- +# 2. ФУНКЦИИ ПРИНЯТИЯ РЕШЕНИЯ +# ---------------------------------------------------------------------------- + +def _can_use_only_meta( + required_keys: List[str], + meta_cols: List[str], + data_cols: List[str], + join_keys: Optional[Dict[str, str]] = None, +) -> bool: + """ + Базовая функция: проверяет можно ли обойтись только meta table без JOIN с data table. + + Args: + required_keys: Колонки которые нужно выбрать + meta_cols: Доступные колонки в meta table + data_cols: Доступные колонки в data table + join_keys: Опционально - FK (Dict[primary_col, ref_col]) которые должны быть в meta + + Returns: + True - можно использовать только meta (не нужен JOIN с data) + False - нужен JOIN с data (есть колонки или FK только в data) + """ + # Если есть join_keys, проверяем что все FK в meta + if join_keys: + for primary_col in join_keys.keys(): + if primary_col not in meta_cols: + return False + + # Проверяем все ли required_keys доступны без data table + for k in required_keys: + if k not in meta_cols: + # Ключ не в meta - проверяем есть ли он в data + if k in data_cols: + # Ключ в data - нужен JOIN с data + return False + # Ключ ни в meta, ни в data - будет NULL (допустимо) + + return True + + +def _should_use_forward_meta( + inp: "ComputeInput", + all_select_keys: List[str], +) -> bool: + """ + Определяет какой путь использовать для forward join. + + Returns: + True - ПУТЬ А: только meta (все ключи в meta, 1 CTE) + False - ПУТЬ Б: meta + data (нужен JOIN с data, 2 CTE) + """ + tbl = inp.dt.meta_table.sql_table + meta_cols = [c.name for c in tbl.columns] + + # Fallback: нет data_table + if not hasattr(inp.dt.table_store, 'data_table'): + return True + + data_tbl = inp.dt.table_store.data_table + data_cols = [c.name for c in data_tbl.columns] + + return _can_use_only_meta(all_select_keys, meta_cols, data_cols) + + +def _should_use_reverse_meta( + ref_join_keys: Dict[str, str], + all_select_keys: List[str], + primary_meta_tbl: Any, + primary_store: "TableStoreDB", +) -> bool: + """ + Определяет какой путь использовать для reverse join. + + Args: + ref_join_keys: join_keys из справочной таблицы + all_select_keys: Колонки которые нужно выбрать + primary_meta_tbl: Meta table основной таблицы + primary_store: TableStoreDB основной таблицы + + Returns: + True - ПУТЬ А: JOIN meta с meta (FK и все ключи в primary_meta) + False - ПУТЬ Б: JOIN meta с data (FK или ключи в primary_data) + """ + primary_data_tbl = primary_store.data_table + primary_meta_cols = [c.name for c in primary_meta_tbl.columns] + primary_data_cols = [c.name for c in primary_data_tbl.columns] + + return _can_use_only_meta(all_select_keys, primary_meta_cols, primary_data_cols, ref_join_keys) + + +# ---------------------------------------------------------------------------- +# 3. ФУНКЦИИ ПОСТРОЕНИЯ CTE +# ---------------------------------------------------------------------------- + +# 3.0 Внутренние helper'ы + +def _meta_sql_helper( + tbl: Any, + keys_in_meta: List[str], + offset: float, + filters_idx: Optional[IndexDF], + primary_keys: List[str], + run_config: Optional[RunConfig], +) -> Any: + """ + Строит SQL: SELECT из meta table с WHERE по offset. + + Использует UNION двух SELECT вместо OR для использования индексов: + - SELECT WHERE update_ts > offset (использует индекс на update_ts) + - UNION + - SELECT WHERE delete_ts > offset (использует индекс на delete_ts) + + Это позволяет PostgreSQL использовать Index Scan вместо Sequential Scan. + """ + select_cols = [sa.column(k) for k in keys_in_meta] + adjusted_offset = offset - OFFSET_EPSILON_SECONDS + + # Часть 1: Измененные записи (update_ts > offset) + updated_sql = sa.select(*select_cols).select_from(tbl).where( + tbl.c.update_ts > adjusted_offset + ) + updated_sql = _apply_sql_filters(updated_sql, keys_in_meta, filters_idx, tbl, primary_keys, run_config) + if len(select_cols) > 0: + updated_sql = updated_sql.group_by(*select_cols) + + # Часть 2: Удаленные записи (delete_ts > offset) + # Примечание: IS NOT NULL не нужен - NULL > offset всегда FALSE + deleted_sql = sa.select(*select_cols).select_from(tbl).where( + tbl.c.delete_ts > adjusted_offset + ) + deleted_sql = _apply_sql_filters(deleted_sql, keys_in_meta, filters_idx, tbl, primary_keys, run_config) + if len(select_cols) > 0: + deleted_sql = deleted_sql.group_by(*select_cols) + + # UNION двух частей для использования отдельных индексов + sql = sa.union(updated_sql, deleted_sql) + + return sql + + +def _meta_data_sql_helper( + tbl: Any, + data_tbl: Any, + keys_in_meta: List[str], + keys_in_data_available: List[str], + primary_keys: List[str], + offset: float, + filters_idx: Optional[IndexDF], + run_config: Optional[RunConfig], +) -> Tuple[Any, Any]: + """ + Строит SQL при JOIN meta с data table (2 CTE: changed + deleted). + + Returns: + Tuple[changed_cte, deleted_cte] + """ + join_conditions = [tbl.c[pk] == data_tbl.c[pk] for pk in primary_keys] + join_condition = _build_join_condition(join_conditions) + + select_cols = [tbl.c[k] for k in keys_in_meta] + [data_tbl.c[k] for k in keys_in_data_available] + all_keys = keys_in_meta + keys_in_data_available + + # CTE для измененных записей + changed_sql = sa.select(*select_cols).select_from( + tbl.join(data_tbl, join_condition) + ).where( + sa.and_( + tbl.c.update_ts >= offset, + tbl.c.delete_ts.is_(None) + ) + ) + + changed_sql = _apply_sql_filters(changed_sql, all_keys, filters_idx, tbl, primary_keys, run_config) + + if len(select_cols) > 0: + changed_sql = changed_sql.group_by(*select_cols) + + # CTE для удаленных записей + deleted_select_cols = [ + tbl.c[k] if k in keys_in_meta else sa.literal(None).label(k) + for k in all_keys + ] + + deleted_sql = sa.select(*deleted_select_cols).select_from(tbl).where( + sa.and_( + tbl.c.delete_ts.isnot(None), + tbl.c.delete_ts >= offset + ) + ) + + deleted_sql = _apply_sql_filters(deleted_sql, keys_in_meta, filters_idx, tbl, primary_keys, run_config) + + if len(deleted_select_cols) > 0: + deleted_sql = deleted_sql.group_by(*[tbl.c[k] for k in keys_in_meta]) + + return changed_sql, deleted_sql + + +# 3.1 Forward join + +def _build_forward_meta_cte( + meta_tbl: Any, + primary_keys: List[str], + all_select_keys: List[str], + offset: float, + filters_idx: Optional[IndexDF], + run_config: Optional[RunConfig], +) -> Any: + """ + ПУТЬ А для forward join: Строит 1 CTE когда все ключи в meta table. + + Deleted records включены в WHERE через _build_offset_where_clause: + WHERE (update_ts >= offset) OR (delete_ts >= offset) + + Returns: + Один CTE с changed + deleted records + """ + meta_cols = [c.name for c in meta_tbl.columns] + keys_in_meta = [k for k in all_select_keys if k in meta_cols] + + return _meta_sql_helper(meta_tbl, keys_in_meta, offset, filters_idx, primary_keys, run_config) + + +def _build_forward_data_cte( + meta_tbl: Any, + store: "TableStoreDB", + primary_keys: List[str], + all_select_keys: List[str], + offset: float, + filters_idx: Optional[IndexDF], + run_config: Optional[RunConfig], +) -> Tuple[Any, Any]: + """ + ПУТЬ Б для forward join: Строит 2 CTE когда нужны FK из data table. + + Changed CTE: INNER JOIN с data_table WHERE delete_ts IS NULL + Deleted CTE: SELECT FROM meta WHERE delete_ts >= offset (FK заменяем на NULL) + + Почему 2 CTE? Для deleted records НЕ нужен JOIN - они уже удалены из data table. + + Returns: + Tuple[changed_cte, deleted_cte] + """ + data_tbl = store.data_table + + meta_cols = [c.name for c in meta_tbl.columns] + keys_in_meta = [k for k in all_select_keys if k in meta_cols] + keys_in_data_only = [k for k in all_select_keys if k not in meta_cols] + + data_cols_available = [c.name for c in data_tbl.columns] + keys_in_data_available = [k for k in keys_in_data_only if k in data_cols_available] + + return _meta_data_sql_helper( + tbl=meta_tbl, + data_tbl=data_tbl, + keys_in_meta=keys_in_meta, + keys_in_data_available=keys_in_data_available, + primary_keys=primary_keys, + offset=offset, + filters_idx=filters_idx, + run_config=run_config, + ) + + +# 3.2 Reverse join + +def _build_reverse_meta_cte( + ref_meta_tbl: Any, + ref_join_keys: Dict[str, str], + ref_primary_keys: List[str], + primary_meta_tbl: Any, + all_select_keys: List[str], + offset: float, + filters_idx: Optional[IndexDF], + run_config: Optional[RunConfig], +) -> Optional[Any]: + """ + ПУТЬ А для reverse join: JOIN meta справочника с meta основной таблицы. + Используется когда FK в primary_meta. + + Returns: + SQL query или None если не удалось построить JOIN + """ + meta_cols = [c.name for c in ref_meta_tbl.columns] + primary_meta_cols = [c.name for c in primary_meta_tbl.columns] + + select_cols = [] + group_by_cols = [] + for k in all_select_keys: + if k == 'update_ts': + # update_ts всегда из meta справочника (ref_meta_tbl) + select_cols.append(ref_meta_tbl.c.update_ts) + group_by_cols.append(ref_meta_tbl.c.update_ts) + elif k in primary_meta_cols: + select_cols.append(primary_meta_tbl.c[k]) + group_by_cols.append(primary_meta_tbl.c[k]) + elif k in meta_cols: + # Берём колонки из meta справочника + select_cols.append(ref_meta_tbl.c[k]) + group_by_cols.append(ref_meta_tbl.c[k]) + else: + select_cols.append(sa.literal(None).label(k)) + + join_conditions = [ + primary_meta_tbl.c[primary_col] == ref_meta_tbl.c[ref_col] + for primary_col, ref_col in ref_join_keys.items() + if primary_col in primary_meta_cols and ref_col in meta_cols + ] + + join_condition = _build_join_condition(join_conditions) + if join_condition is None: + return None + + adjusted_offset = offset - OFFSET_EPSILON_SECONDS + + # Часть 1: update_ts > offset (использует индекс на update_ts) + updated_sql = sa.select(*select_cols).select_from( + ref_meta_tbl.join(primary_meta_tbl, join_condition) + ).where(ref_meta_tbl.c.update_ts > adjusted_offset) + updated_sql = _apply_sql_filters(updated_sql, all_select_keys, filters_idx, ref_meta_tbl, ref_primary_keys, run_config) + if len(group_by_cols) > 0: + updated_sql = updated_sql.group_by(*group_by_cols) + + # Часть 2: delete_ts > offset (использует индекс на delete_ts) + deleted_sql = sa.select(*select_cols).select_from( + ref_meta_tbl.join(primary_meta_tbl, join_condition) + ).where(ref_meta_tbl.c.delete_ts > adjusted_offset) + deleted_sql = _apply_sql_filters(deleted_sql, all_select_keys, filters_idx, ref_meta_tbl, ref_primary_keys, run_config) + if len(group_by_cols) > 0: + deleted_sql = deleted_sql.group_by(*group_by_cols) + + # UNION для использования отдельных индексов + sql = sa.union(updated_sql, deleted_sql) + + return sql + + +def _build_reverse_data_cte( + ref_meta_tbl: Any, + ref_join_keys: Dict[str, str], + ref_primary_keys: List[str], + primary_meta_tbl: Any, + primary_store: "TableStoreDB", + all_select_keys: List[str], + offset: float, + filters_idx: Optional[IndexDF], + run_config: Optional[RunConfig], +) -> Tuple[Optional[Any], Optional[Any]]: + """ + ПУТЬ Б для reverse join: JOIN meta справочника с data основной таблицы. + Используется когда FK в primary_data. + + Changed CTE: INNER JOIN с primary_data WHERE delete_ts IS NULL + Deleted CTE: SELECT FROM primary_meta WHERE delete_ts >= offset (FK заменяем на NULL) + + Returns: + Tuple[changed_cte, deleted_cte] или (None, None) если не удалось построить JOIN + """ + primary_data_tbl = primary_store.data_table + meta_cols = [c.name for c in ref_meta_tbl.columns] + primary_data_cols = [c.name for c in primary_data_tbl.columns] + primary_meta_cols = [c.name for c in primary_meta_tbl.columns] + + # CHANGED CTE: JOIN meta с data + select_cols = [] + group_by_cols = [] + for k in all_select_keys: + if k == 'update_ts': + # update_ts всегда из meta справочника (ref_meta_tbl) + select_cols.append(ref_meta_tbl.c.update_ts) + group_by_cols.append(ref_meta_tbl.c.update_ts) + elif k in primary_data_cols: + select_cols.append(primary_data_tbl.c[k]) + group_by_cols.append(primary_data_tbl.c[k]) + elif k in meta_cols: + # Берём колонки из meta справочника + select_cols.append(ref_meta_tbl.c[k]) + group_by_cols.append(ref_meta_tbl.c[k]) + else: + select_cols.append(sa.literal(None).label(k)) + + join_conditions = [ + primary_data_tbl.c[primary_col] == ref_meta_tbl.c[ref_col] + for primary_col, ref_col in ref_join_keys.items() + if primary_col in primary_data_cols and ref_col in meta_cols + ] + + join_condition = _build_join_condition(join_conditions) + if join_condition is None: + return None, None + + # Используем UNION вместо OR для использования индексов на update_ts и delete_ts + adjusted_offset = offset - OFFSET_EPSILON_SECONDS + + # Часть 1: update_ts > offset (использует индекс на update_ts) + updated_sql = sa.select(*select_cols).select_from( + ref_meta_tbl.join(primary_data_tbl, join_condition) + ).where(ref_meta_tbl.c.update_ts > adjusted_offset) + + updated_sql = _apply_sql_filters( + updated_sql, all_select_keys, filters_idx, ref_meta_tbl, ref_primary_keys, run_config + ) + + if len(group_by_cols) > 0: + updated_sql = updated_sql.group_by(*group_by_cols) + + # Часть 2: delete_ts > offset (использует индекс на delete_ts) + deleted_part_sql = sa.select(*select_cols).select_from( + ref_meta_tbl.join(primary_data_tbl, join_condition) + ).where(ref_meta_tbl.c.delete_ts > adjusted_offset) + + deleted_part_sql = _apply_sql_filters( + deleted_part_sql, all_select_keys, filters_idx, ref_meta_tbl, ref_primary_keys, run_config + ) + + if len(group_by_cols) > 0: + deleted_part_sql = deleted_part_sql.group_by(*group_by_cols) + + # UNION для использования отдельных индексов + changed_sql = sa.union(updated_sql, deleted_part_sql) + + # DELETED CTE: Только из primary_meta (FK заменяем на NULL) + deleted_select_cols = [] + for k in all_select_keys: + if k == 'update_ts': + deleted_select_cols.append(ref_meta_tbl.c.update_ts) + elif k in primary_meta_cols: + deleted_select_cols.append(primary_meta_tbl.c[k]) + elif k in meta_cols: + deleted_select_cols.append(ref_meta_tbl.c[k]) + else: + # FK из primary_data недоступен для deleted records + deleted_select_cols.append(sa.literal(None).label(k)) + + # JOIN primary_meta с ref_meta_tbl для deleted records + deleted_join_conditions = [ + primary_meta_tbl.c[primary_col] == ref_meta_tbl.c[ref_col] + for primary_col, ref_col in ref_join_keys.items() + if primary_col in primary_meta_cols and ref_col in meta_cols + ] + + deleted_join_condition = _build_join_condition(deleted_join_conditions) + if deleted_join_condition is None: + return changed_sql, None + + deleted_sql = sa.select(*deleted_select_cols).select_from( + ref_meta_tbl.join(primary_meta_tbl, deleted_join_condition) + ).where(ref_meta_tbl.c.delete_ts > adjusted_offset) + + deleted_sql = _apply_sql_filters( + deleted_sql, all_select_keys, filters_idx, ref_meta_tbl, ref_primary_keys, run_config + ) + + return changed_sql, deleted_sql + + +# ---------------------------------------------------------------------------- +# 4. КООРДИНАТОРЫ +# ---------------------------------------------------------------------------- + +def _build_forward_input_cte( + inp: "ComputeInput", + all_select_keys: List[str], + offset: float, + filters_idx: Optional[IndexDF], + run_config: Optional[RunConfig], + cte_name: str, +) -> List[Any]: + """ + Строит CTE для forward join (таблица без join_keys). + + Находит изменения в самой таблице inp (прямой путь). + + Returns: + List of CTE objects (1 CTE для пути А, 2 CTE для пути Б) + """ + meta_tbl = inp.dt.meta_table.sql_table + primary_keys = inp.dt.primary_keys + + # Определяем какой путь использовать + use_meta_path = _should_use_forward_meta(inp, all_select_keys) + + # Выполняем соответствующий путь + if use_meta_path: + # ПУТЬ А: строим 1 CTE из meta table (changed + deleted в одном WHERE) + changed_sql = _build_forward_meta_cte(meta_tbl, primary_keys, all_select_keys, offset, filters_idx, run_config) + deleted_sql = None + else: + # ПУТЬ Б: строим 2 CTE (changed с JOIN, deleted без JOIN) + from datapipe.store.database import TableStoreDB + + store = inp.dt.table_store + assert isinstance(store, TableStoreDB) + + changed_sql, deleted_sql = _build_forward_data_cte( + meta_tbl, store, primary_keys, all_select_keys, offset, filters_idx, run_config + ) + + ctes = [changed_sql.cte(name=cte_name)] + if deleted_sql is not None: + deleted_cte_name = cte_name.replace('_changes', '_deleted') + ctes.append(deleted_sql.cte(name=deleted_cte_name)) + + return ctes + + +def _build_reverse_input_cte( + ref_meta_tbl: Any, + ref_join_keys: Dict[str, str], + ref_primary_keys: List[str], + primary_meta_tbl: Any, + primary_store: "TableStoreDB", + all_select_keys: List[str], + offset: float, + filters_idx: Optional[IndexDF], + run_config: Optional[RunConfig], + cte_name: str, +) -> List[Any]: + """ + Строит CTE для reverse join (справочная таблица с join_keys). + + Когда изменяется справочная таблица (users), находит зависимые записи + в основной таблице (subscriptions) через обратный JOIN. + + Args: + ref_meta_tbl: Meta table справочной таблицы + ref_join_keys: join_keys связывающие таблицы + ref_primary_keys: Primary keys справочной таблицы + primary_meta_tbl: Meta table основной таблицы + primary_store: TableStoreDB основной таблицы + all_select_keys: Колонки для выборки + offset: Временная метка offset + filters_idx: Дополнительные фильтры + run_config: Конфигурация выполнения + cte_name: Имя для CTE + + Returns: + List of CTE objects (1 CTE для пути А, 2 CTE для пути Б) + """ + # Определяем какой путь использовать для reverse_join + use_meta_path = _should_use_reverse_meta( + ref_join_keys, all_select_keys, primary_meta_tbl, primary_store + ) + + # Выполняем соответствующий путь + if use_meta_path: + # ПУТЬ А: JOIN meta с meta (FK в primary_meta) + changed_sql = _build_reverse_meta_cte( + ref_meta_tbl, ref_join_keys, ref_primary_keys, + primary_meta_tbl, all_select_keys, offset, filters_idx, run_config + ) + deleted_sql = None + else: + # ПУТЬ Б: JOIN meta с data (FK в primary_data) - возвращает 2 CTE + changed_sql, deleted_sql = _build_reverse_data_cte( + ref_meta_tbl, ref_join_keys, ref_primary_keys, + primary_meta_tbl, primary_store, + all_select_keys, offset, filters_idx, run_config + ) + + if changed_sql is None: + return [] + + ctes = [changed_sql.cte(name=cte_name)] + if deleted_sql is not None: + deleted_cte_name = cte_name.replace('_changes', '_deleted') + ctes.append(deleted_sql.cte(name=deleted_cte_name)) + + return ctes + + +def _build_input_ctes( + input_dts: List["ComputeInput"], + offsets: Dict[str, float], + all_select_keys: List[str], + filters_idx: Optional[IndexDF], + run_config: Optional[RunConfig], +) -> List[Any]: + """ + Строит CTE для каждой входной таблицы с фильтром по offset. + + Для таблиц с join_keys делает обратный JOIN к основной таблице. + Для таблиц с join_keys также создаёт отдельный CTE для deleted records. + + Returns: + List of CTE objects + """ + changed_ctes = [] + + # Сначала находим "основную" таблицу - первую без join_keys + primary_inp = None + for inp in input_dts: + if not inp.join_keys: + primary_inp = inp + break + + # Отслеживаем использование таблиц для генерации уникальных имен CTE + # Ключ: имя таблицы, значение: количество использований + table_usage_count: Dict[str, int] = {} + + for inp in input_dts: + tbl = inp.dt.meta_table.sql_table + + # Генерируем уникальное имя CTE для текущей таблицы + cte_name = _generate_unique_cte_name(inp.dt.name, "changes", table_usage_count) + + # Проверяем есть ли ключи в meta table + meta_cols = [c.name for c in tbl.columns] + keys_in_meta = [k for k in all_select_keys if k in meta_cols] + + if len(keys_in_meta) == 0: + continue + + offset = offsets[inp.dt.name] + + # Два взаимоисключающих пути: + # 1. Reverse join: справочная таблица с join_keys + # 2. Forward join: таблица без join_keys + if inp.join_keys and primary_inp and hasattr(primary_inp.dt.table_store, 'data_table'): + from datapipe.store.database import TableStoreDB + + # Reverse join: находим зависимые записи в primary таблице от изменений в referense таблице + primary_store = primary_inp.dt.table_store + assert isinstance(primary_store, TableStoreDB), "primary table must be TableStoreDB for reverse join" + + ctes = _build_reverse_input_cte( + ref_meta_tbl=inp.dt.meta_table.sql_table, + ref_join_keys=inp.join_keys, + ref_primary_keys=inp.dt.primary_keys, + primary_meta_tbl=primary_inp.dt.meta_table.sql_table, + primary_store=primary_store, + all_select_keys=all_select_keys, + offset=offset, + filters_idx=filters_idx, + run_config=run_config, + cte_name=cte_name, + ) + else: + # Forward join: находим изменения в самой primary таблице и зависимые referense записи + ctes = _build_forward_input_cte( + inp, all_select_keys, offset, filters_idx, run_config, cte_name + ) + + changed_ctes.extend(ctes) + + return changed_ctes + + +def _build_error_records_cte( + meta_table: "TransformMetaTable", + transform_keys: List[str], + all_select_keys: List[str], + filters_idx: Optional[IndexDF], +) -> Any: + """ + Строит CTE для записей с ошибками из TransformMetaTable. + + Returns: + CTE object для error_records + """ + 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 + if k == 'update_ts': + error_select_cols.append(sa.cast(sa.literal(None), sa.Float).label(k)) + else: + error_select_cols.append(sa.literal(None).label(k)) + + error_records_sql: Any = sa.select(*error_select_cols).select_from(tr_tbl).where( + sa.or_( + tr_tbl.c.is_success != True, # noqa + tr_tbl.c.process_ts.is_(None) + ) + ) + + error_records_sql = sql_apply_filters_idx_to_subquery( + error_records_sql, transform_keys, filters_idx + ) + + if len(transform_keys) > 0: + error_records_sql = error_records_sql.group_by(*[sa.column(k) for k in transform_keys]) + + return error_records_sql.cte(name="error_records") + + +def _union_or_cross_join_ctes( + ds: "DataStore", + changed_ctes: List[Any], + error_records_cte: Any, + transform_keys: List[str], + additional_columns: List[str], + all_select_keys: List[str], +) -> Tuple[Any, bool]: + """ + Объединяет CTE через UNION или CROSS JOIN в зависимости от структуры ключей. + + Returns: + Tuple[CTE, needs_cross_join_flag] + """ + needs_cross_join = False + + if len(changed_ctes) == 0: + # Если нет входных таблиц с изменениями, используем только ошибки + union_sql: Any = sa.select( + *[error_records_cte.c[k] for k in all_select_keys] + ).select_from(error_records_cte) + else: + # Собираем информацию о том, какие transform_keys есть в каждом CTE + cte_transform_keys_sets = [] + for cte in changed_ctes: + cte_transform_keys = set(k for k in transform_keys if k in cte.c) + cte_transform_keys_sets.append(cte_transform_keys) + + # Проверяем все ли CTE содержат одинаковый набор transform_keys + unique_key_sets = set(map(frozenset, cte_transform_keys_sets)) + needs_cross_join = len(unique_key_sets) > 1 + + if needs_cross_join: + # Cross join сценарий: используем _make_agg_of_agg для FULL OUTER JOIN + compute_input_ctes = [] + keys_for_join = transform_keys + additional_columns # БЕЗ update_ts + for i, cte in enumerate(changed_ctes): + cte_keys = [k for k in keys_for_join if k in cte.c] + compute_input_ctes.append( + ComputeInputCTE(cte=cte, keys=cte_keys, join_type='full') + ) + + cross_join_raw = _make_agg_of_agg( + ds=ds, + transform_keys=transform_keys + additional_columns, + agg_col='update_ts', + ctes=compute_input_ctes + ) + + cross_join_clean_select = sa.select( + *[sa.column(k) for k in all_select_keys] + ).select_from(cross_join_raw) + + cross_join_cte = cross_join_clean_select.cte(name="cross_join_clean") + + # Объединяем cross join результат с error_records через UNION + union_sql = sa.union( + sa.select(*[cross_join_cte.c[k] for k in all_select_keys]).select_from(cross_join_cte), + sa.select(*[error_records_cte.c[k] for k in all_select_keys]).select_from(error_records_cte) + ) + else: + # Обычный UNION - все CTE содержат одинаковый набор transform_keys + union_parts = [] + for cte in changed_ctes: + select_cols = [ + cte.c[k] if k in cte.c else sa.literal(None).label(k) + for k in all_select_keys + ] + union_parts.append(sa.select(*select_cols).select_from(cte)) + + union_parts.append( + sa.select( + *[error_records_cte.c[k] for k in all_select_keys] + ).select_from(error_records_cte) + ) + + union_sql = sa.union(*union_parts) + + return union_sql.cte(name="changed_union"), needs_cross_join + + +def _deduplicate_if_needed( + union_cte: Any, + input_dts: List["ComputeInput"], + transform_keys: List[str], + additional_columns: List[str], + needs_cross_join: bool, +) -> Any: + """ + Применяет дедупликацию при использовании join_keys (если нужно). + + Returns: + Deduplicated CTE или исходный union_cte + """ + has_join_keys = any(inp.join_keys for inp in input_dts) + if has_join_keys and len(transform_keys) > 0 and not needs_cross_join: + group_by_cols = transform_keys + additional_columns + + select_cols = [] + for k in group_by_cols: + select_cols.append(union_cte.c[k]) + select_cols.append(sa.func.max(union_cte.c.update_ts).label('update_ts')) + + deduplicated_sql = ( + sa.select(*select_cols) + .select_from(union_cte) + .group_by(*[union_cte.c[k] for k in group_by_cols]) + ) + return deduplicated_sql.cte(name="changed_union_deduplicated") + + return union_cte + + +def _apply_final_filters_and_sort( + union_cte: Any, + meta_table: "TransformMetaTable", + transform_keys: List[str], + all_select_keys: List[str], + order_by: Optional[List[str]], + order: Literal["asc", "desc"], +) -> Any: + """ + Применяет финальные фильтры и сортировку к результату. + + Returns: + Final SQL query + """ + tr_tbl = meta_table.sql_table + + if len(transform_keys) == 0: + join_onclause_sql: Any = sa.literal(True) + elif len(transform_keys) == 1: + join_onclause_sql = union_cte.c[transform_keys[0]] == tr_tbl.c[transform_keys[0]] + else: + join_onclause_sql = sa.and_(*[union_cte.c[key] == tr_tbl.c[key] for key in transform_keys]) + + is_error_record = union_cte.c.update_ts.is_(None) + + out = ( + sa.select( + sa.literal(1).label("_datapipe_dummy"), + *[union_cte.c[k] for k in all_select_keys if k in union_cte.c], + ) + .select_from(union_cte) + .outerjoin(tr_tbl, onclause=join_onclause_sql) + .where( + sa.or_( + is_error_record, + tr_tbl.c.process_ts.is_(None), + sa.and_( + tr_tbl.c.is_success == True, # noqa + union_cte.c.update_ts > tr_tbl.c.process_ts + ) + ) + ) + ) + + if order_by is None: + out = out.order_by( + tr_tbl.c.priority.desc().nullslast(), + union_cte.c.update_ts.asc().nullslast(), + *[union_cte.c[k] for k in transform_keys], + ) + else: + if order == "desc": + out = out.order_by( + tr_tbl.c.priority.desc().nullslast(), + union_cte.c.update_ts.asc().nullslast(), + *[sa.desc(union_cte.c[k]) for k in order_by], + ) + elif order == "asc": + out = out.order_by( + tr_tbl.c.priority.desc().nullslast(), + union_cte.c.update_ts.asc().nullslast(), + *[sa.asc(union_cte.c[k]) for k in order_by], + ) + + return out + + +# ---------------------------------------------------------------------------- +# 5. ГЛАВНАЯ ФУНКЦИЯ +# ---------------------------------------------------------------------------- + +def build_changed_idx_sql_v2( + ds: "DataStore", + meta_table: "TransformMetaTable", + input_dts: List["ComputeInput"], + transform_keys: List[str], + offset_table: "TransformInputOffsetTable", + transformation_id: str, + filters_idx: Optional[IndexDF] = None, + order_by: Optional[List[str]] = None, + order: Literal["asc", "desc"] = "asc", + run_config: Optional[RunConfig] = None, + additional_columns: Optional[List[str]] = None, +) -> Tuple[Iterable[str], Any]: + """ + Новая версия build_changed_idx_sql, использующая offset'ы для оптимизации. + + Вместо FULL OUTER JOIN всех входных таблиц, выбираем только записи с + update_ts >= offset для каждой входной таблицы, затем объединяем через UNION. + + Args: + additional_columns: Дополнительные колонки для включения в результат (для filtered join) + """ + if additional_columns is None: + additional_columns = [] + + # Полный список колонок для SELECT (transform_keys + additional_columns) + all_select_keys = list(transform_keys) + additional_columns + + # Добавляем update_ts для ORDER BY если его еще нет + if 'update_ts' not in all_select_keys: + all_select_keys.append('update_ts') + + # 1. Получить все offset'ы одним запросом для избежания N+1 + offsets = offset_table.get_offsets_for_transformation(transformation_id) + for inp in input_dts: + if inp.dt.name not in offsets: + offsets[inp.dt.name] = 0.0 + + # 2. Построить CTE для каждой входной таблицы с фильтром по offset + changed_ctes = _build_input_ctes( + input_dts=input_dts, + offsets=offsets, + all_select_keys=all_select_keys, + filters_idx=filters_idx, + run_config=run_config, + ) + + # 3. Построить CTE для error_records + error_records_cte = _build_error_records_cte( + meta_table=meta_table, + transform_keys=transform_keys, + all_select_keys=all_select_keys, + filters_idx=filters_idx, + ) + + # 4. Объединить через UNION или CROSS JOIN + union_cte, needs_cross_join = _union_or_cross_join_ctes( + ds=ds, + changed_ctes=changed_ctes, + error_records_cte=error_records_cte, + transform_keys=transform_keys, + additional_columns=additional_columns, + all_select_keys=all_select_keys, + ) + + # 5. Дедупликация если нужно + union_cte = _deduplicate_if_needed( + union_cte=union_cte, + input_dts=input_dts, + transform_keys=transform_keys, + additional_columns=additional_columns, + needs_cross_join=needs_cross_join, + ) + + # 6. Применить финальные фильтры и сортировку + out = _apply_final_filters_and_sort( + union_cte=union_cte, + meta_table=meta_table, + transform_keys=transform_keys, + all_select_keys=all_select_keys, + order_by=order_by, + order=order, + ) + + return (all_select_keys, out) + +# ---------------------------------------------------------------------------- + +TRANSFORM_INPUT_OFFSET_SCHEMA: DataSchema = [ + sa.Column("transformation_id", sa.String, primary_key=True), + sa.Column("input_table_name", sa.String, primary_key=True), + sa.Column("update_ts_offset", sa.Float), +] + + +class TransformInputOffsetTable: + """ + Таблица для хранения offset'ов (последних обработанных update_ts) для каждой + входной таблицы каждой трансформации. + + Используется для оптимизации поиска измененных данных: вместо FULL OUTER JOIN + всех входных таблиц, выбираем только записи с update_ts > offset. + """ + + def __init__(self, dbconn: DBConn, create_table: bool = False): + self.dbconn = dbconn + self.sql_table = sa.Table( + "transform_input_offsets", + dbconn.sqla_metadata, + *[col._copy() for col in TRANSFORM_INPUT_OFFSET_SCHEMA], + ) + + if create_table: + # Создать таблицу если её нет (аналогично MetaTable) + self.sql_table.create(dbconn.con, checkfirst=True) + + # Создать индекс на transformation_id для быстрого поиска + idx = sa.Index( + "ix_transform_input_offsets_transformation_id", + self.sql_table.c.transformation_id, + ) + idx.create(dbconn.con, checkfirst=True) + + def get_offset(self, transformation_id: str, input_table_name: str) -> Optional[float]: + """Получить последний offset для трансформации и источника""" + sql = sa.select(self.sql_table.c.update_ts_offset).where( + sa.and_( + self.sql_table.c.transformation_id == transformation_id, + self.sql_table.c.input_table_name == input_table_name, + ) + ) + with self.dbconn.con.begin() as con: + result = con.execute(sql).scalar() + return result + + def get_offsets_for_transformation(self, transformation_id: str) -> Dict[str, float]: + """ + Получить все offset'ы для трансформации одним запросом. + + Returns: {input_table_name: update_ts_offset} + """ + try: + sql = sa.select( + self.sql_table.c.input_table_name, + self.sql_table.c.update_ts_offset, + ).where(self.sql_table.c.transformation_id == transformation_id) + + with self.dbconn.con.begin() as con: + results = con.execute(sql).fetchall() + + return {row[0]: row[1] for row in results} + except Exception: + # Таблица может не существовать если create_table=False + # Возвращаем пустой словарь - все offset'ы будут 0.0 (обработаем все данные) + return {} + + def _build_max_offset_expression(self, insert_sql: Any) -> Any: + """ + Создать выражение для атомарного выбора максимального offset. + + Использует CASE WHEN для гарантии что offset только растет, + работает и в SQLite и в PostgreSQL. + """ + return sa.case( + (self.sql_table.c.update_ts_offset > insert_sql.excluded.update_ts_offset, + self.sql_table.c.update_ts_offset), + else_=insert_sql.excluded.update_ts_offset + ) + + def update_offset( + self, transformation_id: str, input_table_name: str, update_ts_offset: float + ) -> None: + """Обновить offset после успешной обработки""" + insert_sql = self.dbconn.insert(self.sql_table).values( + transformation_id=transformation_id, + input_table_name=input_table_name, + update_ts_offset=update_ts_offset, + ) + + max_offset = self._build_max_offset_expression(insert_sql) + + sql = insert_sql.on_conflict_do_update( + index_elements=["transformation_id", "input_table_name"], + set_={"update_ts_offset": max_offset}, + ) + with self.dbconn.con.begin() as con: + con.execute(sql) + + def update_offsets_bulk(self, offsets: Dict[Tuple[str, str], float]) -> None: + """ + Обновить несколько offset'ов за одну транзакцию + offsets: {(transformation_id, input_table_name): update_ts_offset} + """ + if not offsets: + return + + values = [ + { + "transformation_id": trans_id, + "input_table_name": table_name, + "update_ts_offset": offset, + } + for (trans_id, table_name), offset in offsets.items() + ] + + insert_sql = self.dbconn.insert(self.sql_table).values(values) + + max_offset = self._build_max_offset_expression(insert_sql) + + sql = insert_sql.on_conflict_do_update( + index_elements=["transformation_id", "input_table_name"], + set_={"update_ts_offset": max_offset}, + ) + with self.dbconn.con.begin() as con: + con.execute(sql) + + def reset_offset(self, transformation_id: str, input_table_name: Optional[str] = None) -> None: + """Удалить offset (для полной переобработки)""" + sql = self.sql_table.delete().where(self.sql_table.c.transformation_id == transformation_id) + if input_table_name is not None: + sql = sql.where(self.sql_table.c.input_table_name == input_table_name) + + with self.dbconn.con.begin() as con: + con.execute(sql) + + def get_statistics(self, transformation_id: Optional[str] = None) -> List[Dict[str, Any]]: + """ + Получить статистику по offset'ам для мониторинга + + Returns: [{ + 'transformation_id': str, + 'input_table_name': str, + 'update_ts_offset': float, + 'offset_age_seconds': float # time.time() - update_ts_offset + }] + """ + sql = sa.select( + self.sql_table.c.transformation_id, + self.sql_table.c.input_table_name, + self.sql_table.c.update_ts_offset, + ) + + if transformation_id is not None: + sql = sql.where(self.sql_table.c.transformation_id == transformation_id) + + with self.dbconn.con.begin() as con: + results = con.execute(sql).fetchall() + + now = time.time() + return [ + { + "transformation_id": row[0], + "input_table_name": row[1], + "update_ts_offset": row[2], + "offset_age_seconds": now - row[2] if row[2] else None, + } + for row in results + ] + + def get_offset_count(self) -> int: + """Получить общее количество offset'ов в таблице""" + sql = sa.select(sa.func.count()).select_from(self.sql_table) + + with self.dbconn.con.begin() as con: + result = con.execute(sql).scalar() + return result if result is not None else 0 + + +def initialize_offsets_from_transform_meta( + ds: "DataStore", # type: ignore # noqa: F821 + transform_step: "BaseBatchTransformStep", # type: ignore # noqa: F821 +) -> Dict[str, float]: + """ + Инициализировать offset'ы для существующей трансформации на основе TransformMetaTable. + + Логика: + 1. Находим MIN(process_ts) из успешно обработанных записей в TransformMetaTable + 2. Для каждой входной таблицы находим MAX(update_ts) где update_ts <= min_process_ts + 3. Устанавливаем эти значения как начальные offset'ы + 4. Это гарантирует что мы не пропустим данные которые еще не были обработаны + + Args: + ds: DataStore с мета-подключением + transform_step: Шаг трансформации для которого инициализируются offset'ы + + Returns: + Dict с установленными offset'ами {input_table_name: update_ts_offset} + """ + meta_tbl = transform_step.meta_table.sql_table + + # Найти минимальный process_ts среди успешно обработанных записей + sql = sa.select(sa.func.min(meta_tbl.c.process_ts)).where( + meta_tbl.c.is_success == True # noqa: E712 + ) + + with ds.meta_dbconn.con.begin() as con: + min_process_ts = con.execute(sql).scalar() + + if min_process_ts is None: + # Нет успешно обработанных записей → offset не устанавливаем + return {} + + # Для каждой входной таблицы найти максимальный update_ts <= min_process_ts + offsets = {} + for inp in transform_step.input_dts: + input_tbl = inp.dt.meta_table.sql_table + + sql = sa.select(sa.func.max(input_tbl.c.update_ts)).where( + sa.and_( + input_tbl.c.update_ts <= min_process_ts, + input_tbl.c.delete_ts.is_(None), + ) + ) + + with ds.meta_dbconn.con.begin() as con: + max_update_ts = con.execute(sql).scalar() + + if max_update_ts is not None: + offsets[(transform_step.get_name(), inp.dt.name)] = max_update_ts + + # Установить offset'ы + if offsets: + try: + ds.offset_table.update_offsets_bulk(offsets) + except Exception as e: + # Таблица offset'ов может не существовать + logger.warning( + f"Failed to initialize offsets for {transform_step.get_name()}: {e}. " + "Offset table may not exist (create_meta_table=False)" + ) + return {} + + return {inp_name: offset for (_, inp_name), offset in offsets.items()} diff --git a/datapipe/step/batch_transform.py b/datapipe/step/batch_transform.py index a231ff36..465828b0 100644 --- a/datapipe/step/batch_transform.py +++ b/datapipe/step/batch_transform.py @@ -22,7 +22,7 @@ import pandas as pd from opentelemetry import trace -from sqlalchemy import alias, func +from sqlalchemy import alias, case, func from sqlalchemy.sql.expression import select from tqdm_loggable.auto import tqdm @@ -35,7 +35,11 @@ ) from datapipe.datatable import DataStore, DataTable, MetaTable from datapipe.executor import Executor, ExecutorConfig, SingleThreadExecutor -from datapipe.meta.sql_meta import TransformMetaTable, build_changed_idx_sql +from datapipe.meta.sql_meta import ( + TransformMetaTable, + build_changed_idx_sql_v1, + build_changed_idx_sql_v2, +) from datapipe.run_config import LabelDict, RunConfig from datapipe.types import ( ChangeList, @@ -90,6 +94,7 @@ def __init__( filters: Optional[Union[LabelDict, Callable[[], LabelDict]]] = None, order_by: Optional[List[str]] = None, order: Literal["asc", "desc"] = "asc", + use_offset_optimization: bool = False, ) -> None: # Support both old API (List[DataTable]) and new API (List[ComputeInput]) # Convert to new API format @@ -112,6 +117,8 @@ def __init__( ) self.chunk_size = chunk_size + self.ds = ds # Сохраняем ссылку на DataStore для доступа к offset_table + self.use_offset_optimization = use_offset_optimization # Force transform_keys to be a list, otherwise Pandas will not be happy if transform_keys is not None and not isinstance(transform_keys, list): @@ -133,6 +140,117 @@ def __init__( self.order_by = order_by self.order = order + def _get_use_offset_optimization(self, run_config: Optional[RunConfig] = None) -> bool: + """ + Определить, использовать ли оптимизацию через offset'ы. + + Проверяет флаг self.use_offset_optimization с возможностью переопределения + через RunConfig.labels["use_offset_optimization"]. + + Args: + run_config: Конфигурация запуска, может содержать переопределение флага + + Returns: + True если нужно использовать offset-оптимизацию, False иначе + """ + use_offset = self.use_offset_optimization + if run_config is not None and run_config.labels is not None: + label_override = run_config.labels.get("use_offset_optimization") + if label_override is not None: + use_offset = bool(label_override) + return use_offset + + def _get_optimization_method_name(self, run_config: Optional[RunConfig] = None) -> str: + """ + Получить имя используемого метода оптимизации для логирования. + + Args: + run_config: Конфигурация запуска + + Returns: + "v2_offset" если используется оптимизация, "v1_join" иначе + """ + return "v2_offset" if self._get_use_offset_optimization(run_config) else "v1_join" + + def _get_additional_idx_columns(self) -> List[str]: + """ + Собрать дополнительные колонки, необходимые для filtered join. + + Возвращает список колонок из join_keys, которые нужно включить в idx + для работы filtered join оптимизации. + + Returns: + Список имен колонок (без дубликатов) + """ + additional_columns = [] + + for inp in self.input_dts: + if inp.join_keys: + # Добавляем колонки из ключей join_keys (левая часть маппинга) + # Например, для {"user_id": "id"} добавляем "user_id" + for idx_col in inp.join_keys.keys(): + if idx_col not in self.transform_keys and idx_col not in additional_columns: + additional_columns.append(idx_col) + + return additional_columns + + def _build_changed_idx_sql( + self, + ds: DataStore, + filters_idx: Optional[IndexDF] = None, + order_by: Optional[List[str]] = None, + order: Literal["asc", "desc"] = "asc", + run_config: Optional[RunConfig] = None, + ): + """ + Вспомогательный метод для выбора версии build_changed_idx_sql. + Переключается между v1 (FULL OUTER JOIN) и v2 (offset-based) на основе флага. + + Флаг можно переопределить через RunConfig.labels["use_offset_optimization"]. + """ + use_offset = self._get_use_offset_optimization(run_config) + method = self._get_optimization_method_name(run_config) + + # Получить дополнительные колонки для filtered join + additional_columns = self._get_additional_idx_columns() + + with tracer.start_as_current_span(f"build_changed_idx_sql_{method}"): + start_time = time.time() + + if use_offset: + keys, sql = build_changed_idx_sql_v2( + ds=ds, + meta_table=self.meta_table, + input_dts=self.input_dts, + transform_keys=self.transform_keys, + offset_table=ds.offset_table, + transformation_id=self.get_name(), + filters_idx=filters_idx, + order_by=order_by, + order=order, + run_config=run_config, + additional_columns=additional_columns, # Передаем дополнительные колонки + ) + else: + keys, sql = build_changed_idx_sql_v1( + ds=ds, + meta_table=self.meta_table, + input_dts=self.input_dts, + transform_keys=self.transform_keys, + filters_idx=filters_idx, + order_by=order_by, + order=order, + run_config=run_config, + additional_columns=additional_columns, # Передаем дополнительные колонки + ) + + query_build_time = time.time() - start_time + + # Логирование времени построения запроса + logger.debug(f"[{self.get_name()}] Query build time ({method}): {query_build_time:.3f}s") + + return keys, sql + @classmethod def compute_transform_schema( cls, @@ -200,11 +318,8 @@ def get_changed_idx_count( run_config: Optional[RunConfig] = None, ) -> int: run_config = self._apply_filters_to_run_config(run_config) - _, sql = build_changed_idx_sql( + _, sql = self._build_changed_idx_sql( ds=ds, - meta_table=self.meta_table, - input_dts=self.input_dts, - transform_keys=self.transform_keys, run_config=run_config, ) @@ -241,11 +356,8 @@ def get_full_process_ids( run_config=run_config, ) - join_keys, u1 = build_changed_idx_sql( + join_keys, u1 = self._build_changed_idx_sql( ds=ds, - meta_table=self.meta_table, - input_dts=self.input_dts, - transform_keys=self.transform_keys, run_config=run_config, order_by=self.order_by, order=self.order, # type: ignore # pylance is stupid @@ -259,14 +371,29 @@ def get_full_process_ids( extra_filters = {} def alter_res_df(): - with ds.meta_dbconn.con.begin() as con: - for df in pd.read_sql_query(u1, con=con, chunksize=chunk_size): - df = df[self.transform_keys] + # Определяем метод для логирования + method = self._get_optimization_method_name(run_config) + + with tracer.start_as_current_span(f"execute_changed_idx_sql_{method}"): + start_time = time.time() + + with ds.meta_dbconn.con.begin() as con: + for df in pd.read_sql_query(u1, con=con, chunksize=chunk_size): + # Используем join_keys (которые включают transform_keys + additional_columns) + # Фильтруем только колонки, которые есть в df + available_keys = [k for k in join_keys if k in df.columns] + df = df[available_keys] + + for k, v in extra_filters.items(): + df[k] = v - for k, v in extra_filters.items(): - df[k] = v + yield cast(IndexDF, df) - yield cast(IndexDF, df) + query_exec_time = time.time() - start_time + logger.debug( + f"[{self.get_name()}] Query execution time ({method}): {query_exec_time:.3f}s, " + f"rows: {idx_count}" + ) return math.ceil(idx_count / chunk_size), alter_res_df() @@ -287,20 +414,30 @@ def get_change_list_process_ids( # TODO пересмотреть эту логику, выглядит избыточной # (возможно, достаточно посчитать один раз для всех # input таблиц) - _, sql = build_changed_idx_sql( + _, sql = self._build_changed_idx_sql( ds=ds, - meta_table=self.meta_table, - input_dts=self.input_dts, - transform_keys=self.transform_keys, filters_idx=idx, run_config=run_config, ) - with ds.meta_dbconn.con.begin() as con: - table_changes_df = pd.read_sql_query( - sql, - con=con, + + # Определяем метод для логирования + method = self._get_optimization_method_name(run_config) + + with tracer.start_as_current_span(f"execute_changed_idx_sql_change_list_{method}"): + start_time = time.time() + + with ds.meta_dbconn.con.begin() as con: + table_changes_df = pd.read_sql_query( + sql, + con=con, + ) + table_changes_df = table_changes_df[self.transform_keys] + + query_exec_time = time.time() - start_time + logger.debug( + f"[{self.get_name()}] Change list query execution time ({method}): " + f"{query_exec_time:.3f}s, rows: {len(table_changes_df)}" ) - table_changes_df = table_changes_df[self.transform_keys] changes.append(table_changes_df) else: @@ -317,6 +454,76 @@ def gen(): return chunk_count, gen() + def _get_max_update_ts_for_batch( + self, + ds: DataStore, + compute_input: "ComputeInput", + processed_idx: IndexDF, + ) -> Optional[float]: + """ + Получить максимальный update_ts из входной таблицы для УСПЕШНО обработанного батча. + + Важно: используем processed_idx который содержит только успешно обработанные записи + из output_dfs (result.index), а не весь батч idx. + + Для JoinSpec таблиц (с join_keys) используем join_keys для фильтрации вместо primary_keys. + Пример: profiles с join_keys={'user_id': 'id'} фильтруется по profiles.id IN (processed_idx.user_id) + """ + from datapipe.sql_util import sql_apply_idx_filter_to_table + + if len(processed_idx) == 0: + return None + + input_dt = compute_input.dt + tbl = input_dt.meta_table.sql_table + + # Построить запрос с фильтром по processed_idx (только успешно обработанные) + # Берем максимум из update_ts и delete_ts (для корректного учета удалений) + # Используем CASE WHEN вместо greatest() для совместимости с SQLite + max_of_both = case( + (tbl.c.delete_ts.isnot(None) & (tbl.c.delete_ts > tbl.c.update_ts), tbl.c.delete_ts), else_=tbl.c.update_ts + ) + max_ts_expr = func.max(max_of_both) + sql = select(max_ts_expr) + + # Для JoinSpec таблиц используем join_keys вместо primary_keys + # join_keys: Dict[idx_col -> dt_col], например {'user_id': 'id'} + # Значит: фильтруем dt.id по значениям из processed_idx.user_id + if compute_input.join_keys: + # Для каждого join key создаём фильтр + # Используем только те join keys, у которых idx_col есть в processed_idx + idx_keys = list(processed_idx.columns) + + # Создаём маппинг dt_col -> idx_col для фильтрации + # Пример: {'id': 'user_id'} - фильтровать dt.id по processed_idx.user_id + filter_mapping = {} + for idx_col, dt_col in compute_input.join_keys.items(): + if idx_col in idx_keys: + filter_mapping[dt_col] = idx_col + + if filter_mapping: + # Применяем фильтр используя dt columns как ключи + # Но берем значения из соответствующих idx columns + filter_keys = list(filter_mapping.keys()) + + # Создаём переименованный processed_idx для sql_apply_idx_filter_to_table + # Например: user_id -> id чтобы фильтровать по dt.id + renamed_idx = processed_idx.rename(columns={v: k for k, v in filter_mapping.items()}) + sql = sql_apply_idx_filter_to_table(sql, tbl, filter_keys, renamed_idx) + else: + # Для обычных таблиц используем primary_keys + idx_keys = list(processed_idx.columns) + filter_keys = [k for k in input_dt.primary_keys if k in idx_keys] + + # Если нет общих ключей, не можем отфильтровать - берем максимум по всей таблице + if len(filter_keys) > 0: + sql = sql_apply_idx_filter_to_table(sql, tbl, filter_keys, processed_idx) + + with ds.meta_dbconn.con.begin() as con: + result = con.execute(sql).scalar() + + return result + def store_batch_result( self, ds: DataStore, @@ -360,6 +567,29 @@ def store_batch_result( self.meta_table.mark_rows_processed_success(idx, process_ts=process_ts, run_config=run_config) + # Вычисляем и возвращаем offset'ы через ChangeList + # Это позволяет работать с RayExecutor и устраняет race conditions + if output_dfs is not None: + # Получаем индекс успешно обработанных записей из output_dfs + # Используем первый output для извлечения processed_idx + if isinstance(output_dfs, (list, tuple)): + first_output = output_dfs[0] + else: + first_output = output_dfs + + # Извлекаем индекс из DataFrame результата + if not first_output.empty: + processed_idx = data_to_index(first_output, self.transform_keys) + + for inp in self.input_dts: + # Найти максимальный update_ts из УСПЕШНО обработанного батча + max_update_ts = self._get_max_update_ts_for_batch(ds, inp, processed_idx) + + if max_update_ts is not None: + offset_key = (self.get_name(), inp.dt.name) + # Добавляем offset в ChangeList для возврата из функции + changes.offsets[offset_key] = max_update_ts + return changes def store_batch_err( @@ -405,7 +635,55 @@ def get_batch_input_dfs( idx: IndexDF, run_config: Optional[RunConfig] = None, ) -> List[DataDF]: - return [inp.dt.get_data(idx) for inp in self.input_dts] + """ + Получить входные данные для батча с поддержкой filtered join. + + Если у ComputeInput указаны join_keys, читаем только связанные записи + для оптимизации производительности. + """ + result = [] + + for inp in self.input_dts: + if inp.join_keys: + # FILTERED JOIN: Читаем только связанные записи + # Извлекаем уникальные значения foreign keys из idx + filtered_idx_data = {} + all_keys_present = True + + for idx_col, dt_col in inp.join_keys.items(): + if idx_col in idx.columns: + # Получаем уникальные значения и создаем маппинг + unique_values = idx[idx_col].unique() + filtered_idx_data[dt_col] = unique_values + else: + # Если хотя бы одного ключа нет - используем fallback + all_keys_present = False + break + + if all_keys_present and filtered_idx_data: + # Создаем filtered_idx для чтения только нужных записей + filtered_idx = IndexDF(pd.DataFrame(filtered_idx_data)) + + logger.debug( + f"[{self.get_name()}] Filtered join for {inp.dt.name}: " + f"reading {len(filtered_idx)} records instead of full table" + ) + + data = inp.dt.get_data(filtered_idx) + else: + # Fallback: если не все ключи присутствуют, читаем по idx + logger.debug( + f"[{self.get_name()}] Filtered join fallback for {inp.dt.name}: " + f"join_keys={inp.join_keys} not found in idx columns {list(idx.columns)}" + ) + data = inp.dt.get_data(idx) + else: + # Обычное чтение по idx + data = inp.dt.get_data(idx) + + result.append(data) + + return result def process_batch_dfs( self, @@ -478,7 +756,8 @@ def run_full( if idx_count is not None and idx_count == 0: return - executor.run_process_batch( + # Получаем результат с offset'ами из executor'а + changes = executor.run_process_batch( name=self.name, ds=ds, idx_count=idx_count, @@ -488,6 +767,25 @@ def run_full( executor_config=self.executor_config, ) + # КРИТИЧНО: Коммит offset'ов ТОЛЬКО после успешного завершения ВСЕГО run_full + # Offset'ы передаются через ChangeList.offsets, что работает с любым executor'ом + # Это обеспечивает атомарность: offset - маркер полностью обработанного прогона + # Предотвращает проблемы: + # 1. "Полудвиги" окна при падении mid-batch + # 2. Порчу offset'ов при вызове run_changelist (он вообще не должен трогать offset'ы) + # 3. Потерю offset'ов в RayExecutor (т.к. они теперь возвращаются через результат) + 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: + # Таблица 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)" + ) + ds.event_logger.log_step_full_complete(self.name) def run_changelist( @@ -552,6 +850,7 @@ class DatatableBatchTransform(PipelineStep): transform_keys: Optional[List[str]] = None kwargs: Optional[Dict] = None labels: Optional[Labels] = None + use_offset_optimization: bool = False def build_compute(self, ds: DataStore, catalog: Catalog) -> List[ComputeStep]: input_dts = [catalog.get_datatable(ds, name) for name in self.inputs] @@ -568,6 +867,7 @@ def build_compute(self, ds: DataStore, catalog: Catalog) -> List[ComputeStep]: transform_keys=self.transform_keys, chunk_size=self.chunk_size, labels=self.labels, + use_offset_optimization=self.use_offset_optimization, ) ] @@ -584,6 +884,7 @@ def __init__( transform_keys: Optional[List[str]] = None, chunk_size: int = 1000, labels: Optional[Labels] = None, + use_offset_optimization: bool = False, ) -> None: super().__init__( ds=ds, @@ -593,6 +894,7 @@ def __init__( transform_keys=transform_keys, chunk_size=chunk_size, labels=labels, + use_offset_optimization=use_offset_optimization, ) self.func = func @@ -626,18 +928,20 @@ class BatchTransform(PipelineStep): filters: Optional[Union[LabelDict, Callable[[], LabelDict]]] = None order_by: Optional[List[str]] = None order: Literal["asc", "desc"] = "asc" + use_offset_optimization: bool = False def pipeline_input_to_compute_input(self, ds: DataStore, catalog: Catalog, input: PipelineInput) -> ComputeInput: if isinstance(input, Required): return ComputeInput( dt=catalog.get_datatable(ds, input.table), join_type="inner", + join_keys=input.join_keys, # Pass join_keys for filtered join ) elif isinstance(input, JoinSpec): - # This should not happen, but just in case return ComputeInput( dt=catalog.get_datatable(ds, input.table), join_type="full", + join_keys=input.join_keys, # Pass join_keys for filtered join ) else: return ComputeInput(dt=catalog.get_datatable(ds, input), join_type="full") @@ -661,6 +965,7 @@ def build_compute(self, ds: DataStore, catalog: Catalog) -> List[ComputeStep]: filters=self.filters, order_by=self.order_by, order=self.order, + use_offset_optimization=self.use_offset_optimization, ) ] @@ -681,6 +986,7 @@ def __init__( filters: Optional[Union[LabelDict, Callable[[], LabelDict]]] = None, order_by: Optional[List[str]] = None, order: Literal["asc", "desc"] = "asc", + use_offset_optimization: bool = False, ) -> None: super().__init__( ds=ds, @@ -694,6 +1000,7 @@ def __init__( filters=filters, order_by=order_by, order=order, + use_offset_optimization=use_offset_optimization, ) self.func = func diff --git a/datapipe/types.py b/datapipe/types.py index f0af4044..f83cbe99 100644 --- a/datapipe/types.py +++ b/datapipe/types.py @@ -9,6 +9,7 @@ Dict, List, NewType, + Optional, Set, Tuple, Type, @@ -63,6 +64,9 @@ @dataclass class JoinSpec: table: TableOrName + # Filtered join optimization: mapping from idx columns to table columns + # Example: {"user_id": "id"} means filter table by table.id IN (idx.user_id) + join_keys: Optional[Dict[str, str]] = None @dataclass @@ -76,6 +80,8 @@ class Required(JoinSpec): @dataclass class ChangeList: changes: Dict[str, IndexDF] = field(default_factory=lambda: cast(Dict[str, IndexDF], {})) + # Offset'ы для оптимизации: (step_name, input_table_name) -> max_update_ts + offsets: Dict[Tuple[str, str], float] = field(default_factory=dict) def append(self, table_name: str, idx: IndexDF) -> None: if table_name in self.changes: @@ -93,6 +99,10 @@ def extend(self, other: ChangeList): for key in other.changes.keys(): self.append(key, other.changes[key]) + # Объединяем offset'ы: берем максимум для каждого ключа + for offset_key, offset_value in other.offsets.items(): + self.offsets[offset_key] = max(self.offsets.get(offset_key, 0), offset_value) + def empty(self): return len(self.changes.keys()) == 0 diff --git a/docs/source/SUMMARY.md b/docs/source/SUMMARY.md index 04c961c7..9cbb2803 100644 --- a/docs/source/SUMMARY.md +++ b/docs/source/SUMMARY.md @@ -7,6 +7,7 @@ - [Table](./concepts-table.md) - [Transform](./concepts-transform.md) - [How merging works](./how-merging-works.md) +- [Offset Optimization](./offset-optimization.md) # Command Line Interface diff --git a/docs/source/offset-optimization-commit-strategy.md b/docs/source/offset-optimization-commit-strategy.md new file mode 100644 index 00000000..61e716e8 --- /dev/null +++ b/docs/source/offset-optimization-commit-strategy.md @@ -0,0 +1,175 @@ +# Стратегия фиксации офсетов + +**Расположение в коде:** `datapipe/step/batch_transform.py:740-789` + +[← Назад к обзору](./offset-optimization.md) + +--- + +## Текущая стратегия: Атомарная фиксация в конце run_full + +**Принцип:** Офсеты фиксируются **только после успешной обработки всех батчей**. + +```python +def run_full(self, ds, run_config=None, executor=None): + idx = self.get_changed_idx(ds, run_config) + changes = executor.run_process_batch(...) + + # Фиксация офсетов ТОЛЬКО в конце + if changes.offsets: + ds.offset_table.update_offsets_bulk(changes.offsets) + + return changes +``` + +### Гарантии + +✅ **Нет потери данных** — при сбое в середине обработки офсет не изменяется, данные переобработаются + +✅ **Изоляция от run_changelist** — `run_changelist` не трогает офсеты, только `run_full` обновляет их + +✅ **Корректное восстановление** — после перезапуска обработка продолжается с последнего гарантированно завершенного run_full + +### Проблема + +**Сценарий:** +``` +Обработка 1000 батчей: + - Батчи 1-998: ✅ Успешно обработаны + - Батч 999: ❌ CRASH! + - Батч 1000: ❌ Не запустился + +Офсет НЕ зафиксирован (crash до конца run_full) + +После перезапуска: + - Все 1000 батчей переобрабатываются с начала + - Потеря прогресса обработки 998 батчей +``` + +**Вывод:** При обработке больших таблиц риск потери значительного прогресса. + +--- + +## Планы развития: Побатчевый коммит офсетов + +**Цель:** Фиксировать офсет после каждого успешно обработанного батча. + +### Планируемая реализация + +```python +def store_batch_result(...): + # Обработка и сохранение данных + ... + + # Вычислить офсеты для текущего батча + batch_offsets = self._get_max_update_ts_for_batch(ds, idx, new_idx) + + # НОВОЕ: Фиксировать офсет сразу после успешной обработки батча + if batch_offsets and use_per_batch_offset_commit: + ds.offset_table.update_offsets_bulk(batch_offsets) + logger.info(f"Updated offsets after batch: {batch_offsets}") + + return ChangeList( + changes={self.name: new_idx}, + offsets=batch_offsets + ) +``` + +### Преимущества + +✅ **Минимизация потерь** — при сбое на батче N продолжение с батча N, а не с начала + +✅ **Прогресс сохраняется** — каждый успешный батч фиксируется + +✅ **Большие таблицы** — возможность обработки очень больших таблиц без риска полной переобработки + +### Соображения безопасности + +Для корректной работы побатчевого коммита необходимо гарантировать: + +1. **Идемпотентность функции трансформации** + - Переобработка батча должна давать тот же результат + - Функция не должна зависеть от порядка вызовов + +2. **Корректный порядок обработки** + - Батчи должны обрабатываться в порядке возрастания update_ts + - Это гарантирует, что офсет движется монотонно вперед + +3. **Атомарность на уровне батча** + - Офсет фиксируется только после успешного сохранения **всех** результатов батча + - Частичная обработка батча не приводит к обновлению офсета + +### Статус + +⏳ **В планах на следующие версии** + +Побатчевый коммит требует тщательного тестирования для гарантии безопасности данных. + +--- + +## ChangeList как транспорт офсетов + +```python +@dataclass +class ChangeList: + changes: Dict[str, IndexDF] + offsets: Dict[Tuple[str, str], float] # (transformation_id, table_name) -> max_update_ts +``` + +**Поток данных:** + +1. **store_batch_result()** — вычисляет офсеты для батча, накапливает в ChangeList +2. **executor.run_process_batch()** — собирает ChangeList от всех батчей, объединяет офсеты (MAX) +3. **run_full()** — получает итоговый ChangeList, фиксирует офсеты через update_offsets_bulk() + +### Объединение офсетов + +При обработке множества батчей офсеты объединяются с помощью MAX: + +```python +def extend(self, other: ChangeList) -> ChangeList: + new_offsets = {**self.offsets} + for key, offset in other.offsets.items(): + if key in new_offsets: + new_offsets[key] = max(new_offsets[key], offset) # MAX! + else: + new_offsets[key] = offset + return ChangeList(changes=..., offsets=new_offsets) +``` + +**Почему MAX?** + +- Гарантирует монотонное движение офсета вперед +- Офсет отражает максимальную обработанную update_ts +- Безопасно при параллельной обработке батчей + +--- + +## Изоляция от run_changelist + +**Правило:** `run_changelist` **НИКОГДА** не изменяет глобальные офсеты. + +**Почему?** + +```python +# Сценарий с проблемой +run_full() # Офсет = 1702345600 + +# Пользователь запускает инкрементальную обработку конкретных записей +run_changelist(idx=custom_idx) # Обработка записей вне очереди + +# Если run_changelist изменит офсет → Офсет = 1702345800 + +# Следующий run_full() пропустит данные между 1702345600 и 1702345800 +# ПОТЕРЯ ДАННЫХ ❌ +``` + +**Решение:** + +- Только `run_full` изменяет офсеты +- `run_changelist` работает с явно указанным индексом, не влияя на офсеты +- Гарантия: все данные обрабатываются через `run_full` + +--- + +[← Назад к обзору](./offset-optimization.md) diff --git a/docs/source/offset-optimization-filtered-join.md b/docs/source/offset-optimization-filtered-join.md new file mode 100644 index 00000000..d8a1d77b --- /dev/null +++ b/docs/source/offset-optimization-filtered-join.md @@ -0,0 +1,217 @@ +# Filtered Join — оптимизация чтения референсных таблиц + +**Расположение в коде:** `datapipe/step/batch_transform.py:632-686` + +[← Назад к обзору](./offset-optimization.md) + +--- + +## Проблема + +Чтение всей референсной таблицы неэффективно: + +``` +Обрабатываем: 100 постов от 10 пользователей +Референсная таблица profiles: 10,000,000 записей + +Без оптимизации: читаем все 10 млн профилей +Нужны только: 10 профилей + +Проблема: 99.9999% данных читается зря! +``` + +--- + +## Решение: Filtered Join + +**Filtered join** — читать из референсной таблицы только те записи, которые действительно нужны для обработки текущего батча. + +### Как это работает + +**Шаг 1: Извлечь уникальные значения внешних ключей из индекса** + +```python +# Индекс измененных постов +idx = IndexDF(pd.DataFrame({ + "post_id": [1, 2, 3, ..., 100], + "user_id": [100, 101, 100, 102, ...] # ~10 уникальных user_id +})) + +# Извлекаем уникальные user_id +unique_user_ids = idx["user_id"].unique() +# → [100, 101, 102, 103, 104, 105, 106, 107, 108, 109] +``` + +**Шаг 2: Создать фильтрованный индекс** + +```python +# Маппинг: posts.user_id → profiles.id +filtered_idx = IndexDF(pd.DataFrame({ + "id": [100, 101, 102, 103, 104, 105, 106, 107, 108, 109] +})) +``` + +**Шаг 3: Прочитать только нужные записи** + +```python +# Читаем ТОЛЬКО 10 профилей вместо 10 млн! +profiles_data = profiles_table.get_data(filtered_idx) +``` + +--- + +## Реализация + +```python +if inp.join_keys: + # Извлечь уникальные user_id из idx + filtered_idx_data = {} + for idx_col, dt_col in inp.join_keys.items(): + if idx_col in idx.columns: + filtered_idx_data[dt_col] = idx[idx_col].unique() + + # Создать фильтрованный индекс и прочитать + filtered_idx = IndexDF(pd.DataFrame(filtered_idx_data)) + data = inp.dt.get_data(filtered_idx) # Только нужные записи! +``` + +--- + +## Производительность + +### Без filtered join +- **Читается:** 10,000,000 записей +- **Память:** ~1 GB (при 100 байт на запись) +- **Время:** ~10 секунд + +### С filtered join +- **Читается:** 10 записей +- **Память:** ~1 KB +- **Время:** ~10 миллисекунд + +**Ускорение: 1000x** 🚀 + +--- + +## Конфигурация + +```python +step = BatchTransformStep( + name="process_posts", + input_dts=[ + # Основная таблица + ComputeInput( + dt=posts_table, + join_type="full" + ), + + # Референсная таблица с filtered join + ComputeInput( + dt=profiles_table, + join_type="inner", + join_keys={"user_id": "id"} # Включает filtered join! + ), + ], + transform_keys=["post_id"], + use_offset_optimization=True +) +``` + +**Важно:** Filtered join автоматически включается при наличии `join_keys`. + +--- + +## Пример + +### Входные данные + +```python +# Индекс постов для обработки (100 записей) +idx = IndexDF(pd.DataFrame({ + "post_id": [1, 2, 3, ..., 100], + "user_id": [100, 101, 100, 102, 100, 103, ...] +})) + +# join_keys конфигурация +join_keys = {"user_id": "id"} # posts.user_id = profiles.id +``` + +### Процесс фильтрации + +```python +# 1. Извлечение уникальных user_id +unique_user_ids = idx["user_id"].unique() +# [100, 101, 102, 103, 104, 105, 106, 107, 108, 109] # 10 значений + +# 2. Маппинг на колонку profiles.id +filtered_idx_data = {"id": unique_user_ids} + +# 3. Создание фильтрованного индекса +filtered_idx = IndexDF(pd.DataFrame(filtered_idx_data)) + +# 4. Чтение только нужных профилей +profiles_data = profiles_table.get_data(filtered_idx) +# → 10 записей вместо 10,000,000 +``` + +### Результат + +```python +# Без filtered join +profiles_data = profiles_table.get_data() # 10 млн записей +# Использование памяти: 1 GB + +# С filtered join +profiles_data = profiles_table.get_data(filtered_idx) # 10 записей +# Использование памяти: 1 KB + +# Экономия памяти: 99.9999% +``` + +--- + +## Граничные случаи + +### Случай 1: join_keys колонки нет в индексе + +```python +join_keys = {"author_id": "id"} + +# Но в idx только post_id, нет author_id +idx = IndexDF(pd.DataFrame({"post_id": [1, 2, 3]})) + +# Fallback: читается вся таблица +data = inp.dt.get_data() +``` + +### Случай 2: Множественные join_keys + +```python +join_keys = {"category_id": "id", "subcategory_id": "sub_id"} + +# Извлекаются оба ключа +filtered_idx_data = { + "id": idx["category_id"].unique(), + "sub_id": idx["subcategory_id"].unique(), +} + +filtered_idx = IndexDF(pd.DataFrame(filtered_idx_data)) +``` + +--- + +## Когда использовать + +**✅ Рекомендуется:** +- Большие референсные таблицы (> 100k записей) +- Малый процент необходимых данных (< 10%) +- Батчи с ограниченным набором ключей + +**⚠️ Не критично:** +- Малые референсные таблицы (< 10k записей) +- Нужно большинство данных (> 50%) +- Стоимость чтения низкая + +--- + +[← Назад к обзору](./offset-optimization.md) diff --git a/docs/source/offset-optimization-initialization.md b/docs/source/offset-optimization-initialization.md new file mode 100644 index 00000000..1ee28564 --- /dev/null +++ b/docs/source/offset-optimization-initialization.md @@ -0,0 +1,279 @@ +# Инициализация офсетов + +**Расположение в коде:** `datapipe/meta/sql_meta.py:1398-1462` + +[← Назад к обзору](./offset-optimization.md) + +--- + +## Проблема + +Как включить offset-оптимизацию на существующей трансформации без потери данных? + +**Сценарий:** + +``` +Трансформация "process_posts" работает 6 месяцев: + - Обработано: 1,000,000 записей + - Метаданных: 1,000,000 записей в transform_meta + - Входная таблица: 1,200,000 постов (включая новые) + +Включаем offset-оптимизацию: + - Офсет не установлен (None) + - Без инициализации: обработает ВСЕ 1,200,000 постов повторно! +``` + +**Вопрос:** Как установить начальный офсет так, чтобы не переобработать уже обработанные 1,000,000 записей? + +--- + +## Решение: initialize_offsets_from_transform_meta() + +### Алгоритм + +1. **Найти MIN(process_ts)** из успешно обработанных записей трансформации + - Это самая ранняя временная метка обработки + - Гарантирует, что все записи до этого момента были обработаны + +2. **Для каждой входной таблицы найти MAX(update_ts)** где `update_ts <= min_process_ts` + - Это максимальная временная метка, которая гарантированно была обработана + +3. **Установить найденное значение как начальный офсет** + +### SQL запросы + +**Шаг 1: Найти минимальную process_ts** + +```sql +SELECT MIN(process_ts) as min_process_ts +FROM transform_meta +WHERE transformation_id = :transformation_id + AND is_success = True + AND process_ts IS NOT NULL +``` + +**Шаг 2: Для каждой входной таблицы найти MAX(update_ts)** + +```sql +SELECT MAX(update_ts) as initial_offset +FROM input_table +WHERE update_ts <= :min_process_ts +``` + +**Шаг 3: Записать офсеты** + +```sql +INSERT INTO transform_input_offset (transformation_id, input_table_name, update_ts_offset) +VALUES (:transformation_id, :table_name, :initial_offset) +ON CONFLICT (transformation_id, input_table_name) +DO UPDATE SET update_ts_offset = GREATEST( + transform_input_offset.update_ts_offset, + EXCLUDED.update_ts_offset +) +``` + +--- + +## Использование + +```python +# Инициализация офсетов для существующей трансформации +initial_offsets = ds.offset_table.initialize_offsets_from_transform_meta( + transformation_id="process_posts", + input_tables=[ + ("posts", posts_sa_table), + ("profiles", profiles_sa_table), + ] +) + +# Результат +# {'posts': 1702345600.123, 'profiles': 1702345500.456} +``` + +**После инициализации:** + +```python +# Теперь можно безопасно включить оптимизацию +step.use_offset_optimization = True +step.run_full(ds) + +# Обработает только новые записи (200,000 вместо 1,200,000) +``` + +--- + +## Гарантии безопасности + +### 1. Нет потери данных + +``` +min_process_ts = 1702345600 + +Логика: + - Все записи с process_ts >= 1702345600 были обработаны + - Следовательно, все записи с update_ts <= 1702345600 гарантированно обработаны + - Офсет = MAX(update_ts WHERE update_ts <= 1702345600) безопасен +``` + +### 2. Нет дублирования + +``` +Начальный офсет = 1702345600 + +Следующий run_full(): + WHERE update_ts >= 1702345600 -- Офсет + AND update_ts > process_ts -- Дополнительная проверка + +Результат: + - Записи до 1702345600 пропускаются (уже обработаны) + - Записи после 1702345600 обрабатываются + - Записи с update_ts = 1702345600 проверяются по process_ts +``` + +### 3. Корректная обработка границ + +**Случай: записи на границе офсета** + +```python +# Офсет инициализирован: 1702345600 + +# Записи в входной таблице +records = [ + {"id": 1, "update_ts": 1702345599}, # До офсета + {"id": 2, "update_ts": 1702345600}, # На границе + {"id": 3, "update_ts": 1702345601}, # После офсета +] + +# Метаданные +meta = [ + {"id": 1, "process_ts": 1702345650, "is_success": True}, + {"id": 2, "process_ts": 1702345650, "is_success": True}, +] + +# Следующий run_full найдет: +WHERE update_ts >= 1702345600 # id=2, id=3 + AND (process_ts IS NULL OR update_ts > process_ts) + +# id=2: update_ts (1702345600) < process_ts (1702345650) → пропускается +# id=3: process_ts IS NULL → обрабатывается + +# Корректно! Только новые записи обрабатываются +``` + +--- + +## Граничные случаи + +### Случай 1: Нет обработанных записей + +```python +min_process_ts = None # Нет записей в transform_meta + +# Результат: офсеты не устанавливаются +initial_offsets = {} + +# Первый run_full обработает все записи (корректное поведение для новой трансформации) +``` + +### Случай 2: Входная таблица пустая + +```python +max_update_ts = None # Нет записей в input_table + +# Результат: офсет не устанавливается для этой таблицы +initial_offsets = {} + +# Фильтр не применяется (корректное поведение для пустой таблицы) +``` + +### Случай 3: Все записи уже обработаны + +```python +min_process_ts = 1702345600 +max_update_ts = 1702345000 # Все записи до min_process_ts + +# Офсет = 1702345000 + +# Следующий run_full: +# WHERE update_ts >= 1702345000 AND update_ts > process_ts + +# Результат: обработает только новые записи после последней обработки +``` + +--- + +## Пример полного сценария + +### Исходная ситуация + +``` +Трансформация "process_posts" без offset-оптимизации: + - transform_meta: 1,000,000 записей + - MIN(process_ts) = 1702000000 (6 месяцев назад) + - MAX(process_ts) = 1702345600 (вчера) + +Входная таблица posts: + - 1,200,000 записей + - MIN(update_ts) = 1701900000 + - MAX(update_ts) = 1702345700 (сегодня) +``` + +### Шаг 1: Инициализация + +```python +initial_offsets = ds.offset_table.initialize_offsets_from_transform_meta( + transformation_id="process_posts", + input_tables=[("posts", posts_sa_table)] +) +``` + +**SQL выполняется:** + +```sql +-- Найти MIN(process_ts) +SELECT MIN(process_ts) FROM transform_meta +WHERE transformation_id = 'process_posts' AND is_success = True +-- Результат: 1702000000 + +-- Найти MAX(update_ts) для posts +SELECT MAX(update_ts) FROM posts +WHERE update_ts <= 1702000000 +-- Результат: 1701999999 + +-- Записать офсет +INSERT INTO transform_input_offset VALUES ('process_posts', 'posts', 1701999999) +``` + +### Шаг 2: Первый run_full с оптимизацией + +```python +step.use_offset_optimization = True +step.run_full(ds) +``` + +**SQL запрос на изменения:** + +```sql +SELECT post_id, update_ts +FROM posts +WHERE update_ts >= 1701999999 -- Офсет +-- Найдено: все записи от 1701999999 до 1702345700 + +LEFT JOIN transform_meta meta ON post_id = meta.post_id +WHERE update_ts IS NULL + OR meta.process_ts IS NULL + OR (meta.is_success = True AND posts.update_ts > meta.process_ts) +-- Отфильтровано: только записи после 1702345600 (последняя обработка) + +-- Результат: ~200,000 новых записей вместо 1,200,000 +``` + +### Выгода + +- **Без инициализации:** обработка 1,200,000 записей (~2 часа) +- **С инициализацией:** обработка 200,000 записей (~20 минут) +- **Экономия:** 83% времени + +--- + +[← Назад к обзору](./offset-optimization.md) diff --git a/docs/source/offset-optimization-monitoring.md b/docs/source/offset-optimization-monitoring.md new file mode 100644 index 00000000..2823b603 --- /dev/null +++ b/docs/source/offset-optimization-monitoring.md @@ -0,0 +1,287 @@ +# Метрики и мониторинг + +[← Назад к обзору](./offset-optimization.md) + +--- + +## Метод get_statistics() + +Таблица офсетов предоставляет метод для получения статистики, которую можно использовать для мониторинга через Prometheus. + +```python +stats = ds.offset_table.get_statistics() + +# Результат: +# [ +# { +# 'transformation_id': 'process_posts', +# 'input_table_name': 'posts', +# 'update_ts_offset': 1702345600.0, +# 'offset_age_seconds': 120.5 +# }, +# { +# 'transformation_id': 'process_posts', +# 'input_table_name': 'profiles', +# 'update_ts_offset': 1702345580.0, +# 'offset_age_seconds': 140.3 +# }, +# ] +``` + +**Для конкретной трансформации:** + +```python +stats = ds.offset_table.get_statistics(transformation_id="process_posts") +``` + +--- + +## Ключевая метрика: offset_age_seconds + +**Назначение:** Показывает, как давно были обработаны последние данные для входной таблицы. + +**Формула:** +``` +offset_age_seconds = текущее_время - update_ts_offset +``` + +### Интерпретация + +| offset_age_seconds | Статус | Действие | +|-------------------|--------|----------| +| ~60 | ✅ Норма | Регулярные запуски каждую минуту | +| 3600 | ⚠️ Внимание | Обработка отстает на 1 час | +| 86400 | ❌ Критично | Трансформация не запускалась сутки | + +### Типичные паттерны + +**1. Стабильный offset_age — нормальная работа** +``` +offset_age_seconds ≈ const (например, ~60 секунд при запуске каждую минуту) + +График: горизонтальная линия +Интерпретация: трансформация работает стабильно +``` + +**2. Линейно растущий offset_age — обработка отстает** +``` +offset_age_seconds растет со временем + +График: восходящая линия +Причина: входные данные поступают быстрее, чем обрабатываются +Действие: увеличить частоту запусков или оптимизировать обработку +``` + +**3. Резкий рост offset_age — трансформация не запускалась** +``` +offset_age_seconds скачок с 60 до 7200 + +График: вертикальный скачок +Причина: трансформация не запускалась последние 2 часа +Действие: проверить scheduler, логи, доступность системы +``` + +--- + +## Prometheus алерты + +### Базовый алерт: высокий возраст офсета + +```yaml +- alert: DatapipeOffsetAgeTooHigh + expr: datapipe_offset_age_seconds > 3600 + for: 5m + labels: + severity: warning + annotations: + summary: "Offset age too high for {{ $labels.transformation_id }}/{{ $labels.input_table_name }}" + description: "Offset age is {{ $value }}s, indicating processing lag or stalled transformation." +``` + +### Алерт: офсет продолжает расти + +```yaml +- alert: DatapipeOffsetAgeGrowing + expr: rate(datapipe_offset_age_seconds[10m]) > 0 + for: 30m + labels: + severity: warning + annotations: + summary: "Offset age growing for {{ $labels.transformation_id }}/{{ $labels.input_table_name }}" + description: "Processing is falling behind incoming data. Consider increasing execution frequency." +``` + +--- + +## Диагностика проблем + +### Проблема 1: offset_age растет + offset_value не меняется + +**Симптомы:** +- offset_age_seconds растет линейно +- offset_value (update_ts_offset) не изменяется + +**Возможные причины:** +- Трансформация не запускается (проблема в scheduler) +- Входные данные не изменяются +- Офсет застрял из-за ошибки + +**Диагностика:** +1. Проверить логи scheduler: есть ли запуски? +2. Проверить логи последнего run_full: успешно ли завершился? +3. Проверить входные таблицы: есть ли новые данные с `update_ts > offset`? + +--- + +### Проблема 2: offset_age растет + offset_value растет медленно + +**Симптомы:** +- offset_age_seconds медленно растет +- offset_value растет, но отстает от текущего времени + +**Возможные причины:** +- Частота запусков недостаточна +- Обработка батчей слишком медленная +- Слишком большой размер батча + +**Диагностика:** +1. Сравнить offset_value с MAX(update_ts) во входных таблицах +2. Измерить время выполнения run_full +3. Проверить размер батчей и количество обрабатываемых записей +4. Профилировать функцию трансформации + +--- + +### Проблема 3: run_full успешен + offset_value не меняется + +**Симптомы:** +- run_full выполняется успешно (логи OK) +- offset_value не изменяется +- offset_age_seconds продолжает расти + +**Возможные причины:** +- Ошибка при фиксации офсета +- run_full не находит измененных записей +- Используется run_changelist вместо run_full + +**Диагностика:** +1. Проверить логи: есть ли сообщение "Updated offsets for {name}: {offsets}"? +2. Проверить, вызывается ли update_offsets_bulk() +3. Проверить, есть ли записи с `update_ts > offset` во входных таблицах +4. Проверить индекс на update_ts: используется ли он? + +--- + +## Экспорт в Prometheus + +```python +from prometheus_client import Gauge + +# Определение метрик +offset_age_gauge = Gauge( + 'datapipe_offset_age_seconds', + 'Age of offset in seconds (time since last processed update_ts)', + ['transformation_id', 'input_table_name'] +) + +offset_value_gauge = Gauge( + 'datapipe_offset_value', + 'Current offset value (timestamp)', + ['transformation_id', 'input_table_name'] +) + +def export_offset_metrics(ds): + """Экспорт метрик офсетов в Prometheus.""" + stats = ds.offset_table.get_statistics() + + for stat in stats: + labels = { + 'transformation_id': stat['transformation_id'], + 'input_table_name': stat['input_table_name'] + } + + # Экспорт возраста офсета + offset_age_gauge.labels(**labels).set(stat['offset_age_seconds']) + + # Экспорт значения офсета + offset_value_gauge.labels(**labels).set(stat['update_ts_offset']) +``` + +--- + +## Визуализация в Grafana + +### Панель 1: Offset Age Timeline + +Временной график offset_age_seconds для отслеживания трендов: + +```json +{ + "title": "Offset Age by Transformation", + "targets": [{ + "expr": "datapipe_offset_age_seconds", + "legendFormat": "{{ transformation_id }}/{{ input_table_name }}" + }], + "yAxes": [{ + "label": "Age (seconds)" + }] +} +``` + +### Панель 2: Processing Lag + +Отставание обработки от поступления данных: + +```promql +# Время отставания в часах +datapipe_offset_age_seconds / 3600 + +# Скорость роста отставания +rate(datapipe_offset_age_seconds[10m]) +``` + +### Панель 3: Offset Progress + +График прогресса обработки: + +```promql +# Значение офсета +datapipe_offset_value + +# Скорость изменения офсета (записей в секунду) +rate(datapipe_offset_value[5m]) +``` + +--- + +## Рекомендации + +### Настройка алертов + +**Порог 1: Предупреждение** +```yaml +expr: datapipe_offset_age_seconds > 3600 # 1 час +severity: warning +``` + +**Порог 2: Критично** +```yaml +expr: datapipe_offset_age_seconds > 86400 # 1 день +severity: critical +``` + +### Регулярный мониторинг + +**Ежедневно проверять:** +- Максимальный offset_age по всем трансформациям +- Трансформации с растущим offset_age +- Трансформации без обновлений офсетов + +**Еженедельно анализировать:** +- Тренды offset_age за неделю +- Корреляция с нагрузкой системы +- Эффективность offset-оптимизации (сравнение времени выполнения) + +--- + +[← Назад к обзору](./offset-optimization.md) diff --git a/docs/source/offset-optimization-reverse-join.md b/docs/source/offset-optimization-reverse-join.md new file mode 100644 index 00000000..d4f3f387 --- /dev/null +++ b/docs/source/offset-optimization-reverse-join.md @@ -0,0 +1,138 @@ +# Reverse Join для референсных таблиц + +**Расположение в коде:** `datapipe/meta/sql_meta.py:917-983` + +[← Назад к обзору](./offset-optimization.md) + +--- + +## Проблема + +```python +# Основная таблица +posts = [{"post_id": 1, "user_id": 100, "content": "Hello"}] + +# Референсная таблица +profiles = [{"id": 100, "name": "Alice"}] + +# При изменении profiles.name для Alice +# Нужно переобработать все посты пользователя Alice +# Но таблица posts НЕ изменялась! +``` + +**Вопрос:** Как найти все записи в основной таблице, которые зависят от измененных записей в референсной таблице? + +--- + +## Решение: Reverse Join + +**Reverse join** = JOIN от **референсной** таблицы к **основной** таблице. + +### SQL паттерн + +```sql +-- Изменения в референсной таблице +profiles_changes AS ( + SELECT id, update_ts + FROM profiles + WHERE update_ts >= :profiles_offset +) + +-- REVERSE JOIN к основной таблице +SELECT DISTINCT + posts.post_id, -- transform_keys основной таблицы + profiles_changes.update_ts -- update_ts из референса +FROM profiles_changes +JOIN posts ON posts.user_id = profiles_changes.id -- ОБРАТНОЕ направление +``` + +**Обратите внимание:** +- JOIN идет **от** profiles_changes **к** posts (обратное направление) +- В SELECT берутся **transform_keys из основной таблицы** (post_id) +- update_ts используется **из референсной таблицы** (момент изменения профиля) + +--- + +## Конфигурация + +```python +ComputeInput( + dt=profiles_table, + join_type="inner", + join_keys={"user_id": "id"} # posts.user_id = profiles.id +) +``` + +**Формат join_keys:** `{колонка_в_основной_таблице: колонка_в_референсной_таблице}` + +### Примеры + +```python +# Простой случай - один ключ +join_keys={"user_id": "id"} +# posts.user_id = profiles.id + +# Составной ключ +join_keys={"category_id": "id", "subcategory_id": "sub_id"} +# posts.category_id = categories.id AND posts.subcategory_id = categories.sub_id +``` + +--- + +## Как это работает + +1. **Находим изменения в референсной таблице:** + ```sql + SELECT id FROM profiles WHERE update_ts >= :offset + ``` + +2. **Обратный JOIN к основной таблице:** + ```sql + JOIN posts ON posts.user_id = profiles.id + ``` + +3. **Извлекаем transform_keys:** + ```sql + SELECT posts.post_id + ``` + +4. **Результат:** Все посты, которые нужно переобработать из-за изменений в профилях + +--- + +## Полный пример + +```python +# Конфигурация трансформации +step = BatchTransformStep( + name="enrich_posts_with_profiles", + input_dts=[ + # Основная таблица + ComputeInput( + dt=posts_table, + join_type="full" + ), + # Референсная таблица с reverse join + ComputeInput( + dt=profiles_table, + join_type="inner", + join_keys={"user_id": "id"} + ), + ], + transform_keys=["post_id"], + use_offset_optimization=True +) +``` + +**Что произойдет при изменении профиля:** + +1. Профиль Alice (id=100) обновлен → update_ts = 1702345700 +2. Offset для profiles = 1702345600 +3. Запрос находит profiles.id = 100 (update_ts >= offset) +4. Reverse join находит все posts где user_id = 100 +5. Эти посты помечаются для переобработки +6. Трансформация обогащает посты новыми данными профиля + +--- + +[← Назад к обзору](./offset-optimization.md) diff --git a/docs/source/offset-optimization-sql-queries.md b/docs/source/offset-optimization-sql-queries.md new file mode 100644 index 00000000..de9a7a2e --- /dev/null +++ b/docs/source/offset-optimization-sql-queries.md @@ -0,0 +1,102 @@ +# Оптимизированные SQL-запросы (v1 vs v2) + +**Расположение в коде:** `datapipe/meta/sql_meta.py:720-1215` + +[← Назад к обзору](./offset-optimization.md) + +--- + +## Алгоритм v1: FULL OUTER JOIN + +**Концепция:** Объединить входные таблицы с метаданными через FULL OUTER JOIN. + +```sql +SELECT transform_keys, input.update_ts +FROM input_table input +FULL OUTER JOIN transform_meta meta ON transform_keys +WHERE + meta.process_ts IS NULL + OR (meta.is_success = True AND input.update_ts > meta.process_ts) + OR meta.is_success != True +``` + +**Производительность:** O(N) где N — размер transform_meta. **Деградирует** с ростом метаданных. + +--- + +## Алгоритм v2: Offset-based + +**Концепция:** Ранняя фильтрация по офсетам + UNION вместо JOIN. + +```sql +WITH +-- Получить офсеты +offsets AS ( + SELECT input_table_name, update_ts_offset + FROM transform_input_offset + WHERE transformation_id = :transformation_id +), + +-- Изменения в каждой входной таблице (ранняя фильтрация!) +input_changes AS ( + SELECT transform_keys, update_ts + FROM input_table + WHERE update_ts >= (SELECT update_ts_offset FROM offsets ...) +), + +-- Записи с ошибками (всегда включены) +error_records AS ( + SELECT transform_keys, NULL as update_ts + FROM transform_meta + WHERE is_success != True +), + +-- UNION всех источников +all_changes AS ( + SELECT * FROM input_changes + UNION ALL + SELECT * FROM error_records +) + +-- Фильтр для исключения уже обработанных +SELECT DISTINCT all_changes.* +FROM all_changes +LEFT JOIN transform_meta meta ON transform_keys +WHERE + all_changes.update_ts IS NULL -- Ошибки + OR meta.process_ts IS NULL -- Новые + OR (meta.is_success = True AND all_changes.update_ts > meta.process_ts) +ORDER BY update_ts, transform_keys +``` + +**Ключевые особенности:** +1. **Ранняя фильтрация:** `WHERE update_ts >= offset` применяется до JOIN с метаданными +2. **Использование индекса:** Фильтр по update_ts использует индекс +3. **UNION вместо JOIN:** Дешевле для больших данных +4. **Проверка process_ts:** Предотвращает зацикливание при использовании `>=` + +**Производительность:** O(M) где M — записи с `update_ts >= offset`. **Константная** производительность. + +--- + +## Критически важно: >= а не > + +```python +# ✅ Правильно +WHERE update_ts >= offset + +# ❌ Неправильно - потеря данных! +WHERE update_ts > offset # Записи с update_ts == offset потеряны +``` + +**Почему >= ?** + +При использовании `>` записи, у которых `update_ts` точно равен `offset`, будут пропущены. Это может произойти, когда: +- Несколько записей имеют одинаковую временную метку +- Офсет зафиксирован на границе батча + +Использование `>=` гарантирует, что все записи будут обработаны, а дополнительная проверка `update_ts > process_ts` предотвращает дублирование. + +--- + +[← Назад к обзору](./offset-optimization.md) diff --git a/docs/source/offset-optimization-storage.md b/docs/source/offset-optimization-storage.md new file mode 100644 index 00000000..e95fa589 --- /dev/null +++ b/docs/source/offset-optimization-storage.md @@ -0,0 +1,64 @@ +# Хранение и управление офсетами + +**Расположение в коде:** `datapipe/meta/sql_meta.py:1218-1396` + +[← Назад к обзору](./offset-optimization.md) + +--- + +## Класс TransformInputOffsetTable + +Таблица для хранения максимальных обработанных временных меток (офсетов) для каждой комбинации трансформации и входной таблицы. + +**Схема:** +```sql +Table: transform_input_offset +Primary Key: (transformation_id, input_table_name) +Columns: + - transformation_id: VARCHAR + - input_table_name: VARCHAR + - update_ts_offset: FLOAT +``` + +--- + +## Основные методы API + +### get_offsets_for_transformation() + +Получить все офсеты для трансформации **одним запросом** (оптимизировано): + +```python +offsets = ds.offset_table.get_offsets_for_transformation("process_posts") +# {'posts': 1702345678.123, 'profiles': 1702345600.456} +``` + +### update_offsets_bulk() + +Атомарное обновление множества офсетов в одной транзакции: + +```python +offsets = { + ("process_posts", "posts"): 1702345678.123, + ("process_posts", "profiles"): 1702345600.456, +} +ds.offset_table.update_offsets_bulk(offsets) +``` + +**Критическая деталь:** Используется `GREATEST(existing, new)` — офсет **никогда не уменьшается**, что предотвращает потерю данных при race conditions. + +### reset_offset() + +Сброс офсета для повторной обработки: + +```python +# Сбросить офсет для одной таблицы +ds.offset_table.reset_offset("process_posts", "posts") + +# Сбросить все офсеты трансформации +ds.offset_table.reset_offset("process_posts") +``` + +--- + +[← Назад к обзору](./offset-optimization.md) diff --git a/docs/source/offset-optimization.md b/docs/source/offset-optimization.md new file mode 100644 index 00000000..15d2b171 --- /dev/null +++ b/docs/source/offset-optimization.md @@ -0,0 +1,238 @@ +# Оптимизация на основе офсетов (Offset Optimization) + +## Введение + +Оптимизация на основе офсетов (offset optimization) — это функция, которая значительно повышает производительность инкрементальной обработки данных путем отслеживания временной метки последней обработки (offset) для каждой входной таблицы трансформации. Это позволяет Datapipe пропускать уже обработанные записи без полного сканирования таблицы метаданных трансформации. + +## Краткая концепция + +### Традиционный подход (v1) + +Без оптимизации Datapipe использует FULL OUTER JOIN между входными таблицами и таблицей метаданных трансформации для поиска измененных записей. Этот подход корректен, но требует полного сканирования метаданных на каждом запуске, что замедляет обработку по мере роста объема обработанных данных. + +**Сложность:** O(N), где N — размер таблицы метаданных + +### Оптимизированный подход (v2) + +С включенной оптимизацией Datapipe: +1. **Отслеживает офсеты** — для каждой входной таблицы трансформации хранится максимальная обработанная временная метка `update_ts` +2. **Фильтрует данные рано** — применяет фильтр `WHERE update_ts >= offset` на уровне входных таблиц, используя индекс +3. **Минимизирует сканирование** — обрабатывает только записи, измененные с момента последнего запуска +4. **Атомарно фиксирует офсеты** — обновляет офсеты только после успешного завершения всего run_full + +**Сложность:** O(M), где M — количество записей, измененных с последнего запуска + +## Основные возможности (Features) + +### 1. Хранение и управление офсетами + +**TransformInputOffsetTable** — таблица для хранения офсетов с API для: +- Получения офсетов для трансформации +- Атомарного обновления одного или нескольких офсетов +- Инициализации офсетов из существующих метаданных +- Сброса офсетов для полной переобработки + +[Подробнее →](./offset-optimization-storage.md) + +### 2. Оптимизированные SQL-запросы (v1 vs v2) + +Два алгоритма построения запросов на поиск измененных записей: +- **v1 (FULL OUTER JOIN)** — традиционный подход без офсетов +- **v2 (Offset-based)** — оптимизированный подход с ранней фильтрацией по офсетам + +[Подробнее →](./offset-optimization-sql-queries.md) + +### 3. Reverse Join для референсных таблиц + +При изменении записей в референсной таблице (например, справочник пользователей) Datapipe автоматически находит все зависимые записи в основной таблице через **reverse join** с использованием `join_keys`. + +**Пример:** При обновлении `profiles.name` находятся все `posts` этого пользователя через `posts.profile_id = profiles.id`. + +[Подробнее →](./offset-optimization-reverse-join.md) + +### 4. Filtered Join — оптимизация чтения референсных таблиц + +Вместо чтения всей референсной таблицы, Datapipe извлекает только те записи, которые реально нужны для обработки текущего батча. Это достигается через фильтрацию по уникальным значениям внешних ключей из индекса измененных записей. + +**Пример:** Если обрабатываются посты 10 пользователей, из таблицы `profiles` читаются только профили этих 10 пользователей, а не все миллионы. + +[Подробнее →](./offset-optimization-filtered-join.md) + +### 5. Стратегия фиксации офсетов + +Офсеты фиксируются **атомарно в конце run_full** после успешной обработки всех батчей. Это гарантирует: +- **Отсутствие потери данных** — при сбое в середине обработки офсет не изменяется, данные переобработаются +- **Изоляцию от run_changelist** — инкрементальные запуски не меняют глобальные офсеты +- **Корректное восстановление** — после перезапуска обработка продолжается с последнего гарантированно завершенного run_full + +**Планы развития:** В планах вернуть коммит офсетов после каждого батча, чтобы избежать ситуации, когда при сбое во время обработки большой таблицы приходится начинать с самого начала. При побатчевом коммите офсет будет сохраняться после каждого успешно обработанного батча, что позволит продолжить обработку с последнего завершенного батча вместо полной переобработки. + +[Подробнее →](./offset-optimization-commit-strategy.md) + +### 6. Инициализация офсетов + +Функция **initialize_offsets_from_transform_meta()** позволяет включить оптимизацию на существующих трансформациях без потери данных: +1. Находит MIN(process_ts) из успешно обработанных записей +2. Для каждой входной таблицы находит MAX(update_ts) где update_ts <= min_process_ts +3. Устанавливает это значение как начальный офсет + +[Подробнее →](./offset-optimization-initialization.md) + +## Как включить оптимизацию + +### В определении трансформации + +```python +from datapipe.step.batch_transform import BatchTransformStep +from datapipe.compute import ComputeInput + +step = BatchTransformStep( + ds=ds, + name="process_posts", + func=transform_function, + input_dts=[ + # Основная таблица (без join_keys) + ComputeInput(dt=posts_table, join_type="full"), + + # Референсная таблица с reverse join + ComputeInput( + dt=profiles_table, + join_type="inner", + join_keys={"user_id": "id"} # posts.user_id = profiles.id + ), + ], + output_dts=[output_table], + transform_keys=["post_id"], + + # Включение offset-оптимизации + use_offset_optimization=True +) +``` + +### В runtime через RunConfig + +```python +from datapipe.run_config import RunConfig + +run_config = RunConfig( + labels={"use_offset_optimization": True} +) + +step.run_full(ds, run_config=run_config) +``` + +## Когда использовать + +### Рекомендуется использовать + +- ✅ **Большие таблицы метаданных** — трансформации с большим количеством обработанных записей (> 100k) +- ✅ **Инкрементальные обновления** — малая доля данных изменяется на каждом запуске (< 10%) +- ✅ **Индекс на update_ts** — входные таблицы имеют индекс на поле update_ts +- ✅ **Референсные таблицы** — используются справочники с join_keys +- ✅ **Production окружение** — стабильные пайплайны с регулярными запусками + +### Не рекомендуется использовать + +- ❌ **Полная переобработка** — если обрабатываются все данные на каждом запуске +- ❌ **Малый объем метаданных** — если таблица метаданных небольшая (< 10k записей) +- ❌ **Высокая доля изменений** — если большинство записей обновляется на каждом запуске (> 50%) +- ❌ **Разработка/отладка** — на этапе разработки удобнее использовать v1 для простоты + +## Ограничения + +- **Reverse join и filtered join с join_keys** — при использовании `join_keys` data table должна находиться в той же SQL базе данных, что и meta table. Для non-SQL источников (MongoDB, файлы) используется fallback без JOIN — выбираются только колонки из meta table, filtered join не применяется. + +- **run_changelist использует офсеты, но не обновляет их** — `run_changelist` применяет v2 алгоритм с фильтрацией по офсетам (если оптимизация включена), но не обновляет офсеты после выполнения. Это означает, что регулярные запуски `run_full` необходимы для продвижения офсетов и обработки новых данных. Если запускать только `run_changelist`, новые записи, появившиеся после последнего `run_full`, могут не попасть в обработку. + +- ⚠️ **ВАЖНО: run_changelist может пропустить явно переданные записи** — если оптимизация включена, записи из changelist с `update_ts < offset` будут отфильтрованы и не попадут в обработку, даже если явно переданы в changelist. Это конфликт между философией "обработать всё что передали" и фильтрацией по offset. **Рекомендация:** для сценариев с частым использованием `run_changelist` рассмотрите отключение оптимизации через `RunConfig(labels={"use_offset_optimization": False})` при вызове `run_changelist`. + +## Архитектура и поток данных + +``` +┌─────────────────────────────────────────────────────────────┐ +│ DataStore │ +│ ┌──────────────────────────────────────────────────────┐ │ +│ │ offset_table: TransformInputOffsetTable │ │ +│ │ - get_offset(transformation_id, table_name) │ │ +│ │ - update_offsets_bulk(offsets) │ │ +│ └──────────────────────────────────────────────────────┘ │ +└─────────────────────────────────────────────────────────────┘ + ↑ + │ используется в run_full() + │ +┌─────────────────────────────────────────────────────────────┐ +│ BaseBatchTransformStep │ +│ │ +│ 1. build_changed_idx_sql() → v2 с офсетами │ +│ 2. get_batch_input_dfs() → filtered join │ +│ 3. store_batch_result() → накопление офсетов в ChangeList │ +│ 4. run_full() → фиксация офсетов через │ +│ offset_table.update_offsets_bulk() │ +└─────────────────────────────────────────────────────────────┘ +``` + +## Критически важные детали реализации + +1. **Инклюзивное неравенство** — используется `>=` а не `>` для избежания потери граничных записей +2. **Проверка process_ts** — даже с офсетом проверяется `process_ts > update_ts` для предотвращения зацикливания +3. **Reverse join для референсов** — автоматическое построение обратного join при наличии join_keys +4. **Атомарная фиксация офсетов** — офсеты фиксируются только после успешного run_full +5. **Записи с ошибками всегда включены** — записи с `is_success != True` включаются вне зависимости от офсета +6. **NULL update_ts для ошибок** — записи с ошибками маркируются NULL update_ts для отличия от измененных записей +7. **Сортировка по update_ts** — хронологический порядок критичен для корректности офсетов + +## Дополнительные материалы + +**Детальное описание фич:** +- [Хранение и управление офсетами](./offset-optimization-storage.md) +- [Оптимизированные SQL-запросы](./offset-optimization-sql-queries.md) +- [Reverse Join](./offset-optimization-reverse-join.md) +- [Filtered Join](./offset-optimization-filtered-join.md) +- [Стратегия фиксации офсетов](./offset-optimization-commit-strategy.md) +- [Инициализация офсетов](./offset-optimization-initialization.md) +- [Метрики и мониторинг](./offset-optimization-monitoring.md) + +## Тестирование + +Offset-оптимизация покрыта комплексным набором тестов: +- `tests/test_offset_table.py` — операции с таблицей офсетов +- `tests/test_offset_auto_update.py` — автоматическое обновление офсетов +- `tests/test_offset_joinspec.py` — reverse join с join_keys +- `tests/test_multi_table_filtered_join.py` — filtered join +- `tests/test_batch_transform_with_offset_optimization.py` — интеграционные тесты +- `tests/offset_edge_cases/` — граничные случаи и edge cases (13+ тестов) + +Нагрузочное тестирование: https://github.com/epoch8/datapipe-perf + +## Мониторинг и метрики + +Offset-оптимизация предоставляет метрики для мониторинга через Prometheus. Таблица офсетов поддерживает метод `get_statistics()`, который возвращает статистику по офсетам: + +```python +# Получить статистику по всем офсетам +stats = ds.offset_table.get_statistics() + +# Получить статистику для конкретной трансформации +stats = ds.offset_table.get_statistics(transformation_id="process_posts") +``` + +**Доступные метрики:** +- `transformation_id` — ID трансформации +- `input_table_name` — имя входной таблицы +- `update_ts_offset` — текущее значение офсета (timestamp) +- `offset_age_seconds` — возраст офсета в секундах (сколько времени прошло с момента last processed update_ts) + +**Использование для мониторинга:** + +Метрика `offset_age_seconds` особенно полезна для мониторинга: +- **Растущий offset_age** — индикатор того, что обработка отстает от поступления данных +- **Стабильный offset_age** — нормальная работа с регулярными запусками +- **Большой offset_age** — возможная проблема (трансформация долго не запускалась или упала) + +Эти метрики можно экспортировать в Prometheus для визуализации в Grafana и настройки алертов. + +## История и текущий статус + +Offset-оптимизация прошла первое тестирование в production окружении и показала значительное улучшение производительности для трансформаций с большими объемами обработанных данных. Функция стабильна и готова для использования в production. + +**Ветка разработки:** `Looky-7769/offsets` diff --git a/tests/conftest.py b/tests/conftest.py index 0f545247..75cb9a72 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,6 +16,26 @@ from datapipe.store.database import DBConn +def pytest_collection_modifyitems(config, items): + """ + Исключить performance тесты из обычного test run. + + Performance тесты должны запускаться отдельно: pytest tests/performance/ + """ + # Если явно указана директория performance в аргументах, не фильтруем + args_str = " ".join(config.invocation_params.args) + if "performance" in args_str: + return + + # Исключить все тесты из tests/performance/ если не указано явно + remaining = [] + for item in items: + if "performance" not in str(item.fspath): + remaining.append(item) + + items[:] = remaining + + @pytest.fixture def tmp_dir(): with tempfile.TemporaryDirectory() as d: @@ -51,7 +71,12 @@ def assert_df_equal(a: pd.DataFrame, b: pd.DataFrame) -> bool: @pytest.fixture def dbconn(): if os.environ.get("TEST_DB_ENV") == "sqlite": - DBCONNSTR = "sqlite+pysqlite3:///:memory:" + # Try pysqlite3 first (CI), fallback to pysqlite (local dev) + try: + import pysqlite3 # noqa: F401 + DBCONNSTR = "sqlite+pysqlite3:///:memory:" + except ImportError: + DBCONNSTR = "sqlite+pysqlite:///:memory:" DB_TEST_SCHEMA = None else: pg_host = os.getenv("POSTGRES_HOST", "localhost") diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml index 83f63bcd..27464d12 100644 --- a/tests/docker-compose.yml +++ b/tests/docker-compose.yml @@ -25,8 +25,8 @@ services: environment: discovery.type: single-node ES_JAVA_OPTS: -Xms4g -Xmx4g - xpack.security.enabled: false - xpack.security.http.ssl.enabled: false + xpack.security.enabled: "false" + xpack.security.http.ssl.enabled: "false" node.name: node cluster.name: cluster - http.cors.enabled: true + http.cors.enabled: "true" diff --git a/tests/offset_edge_cases/__init__.py b/tests/offset_edge_cases/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/offset_edge_cases/test_offset_commit_at_run_full_end.py b/tests/offset_edge_cases/test_offset_commit_at_run_full_end.py new file mode 100644 index 00000000..05390d2e --- /dev/null +++ b/tests/offset_edge_cases/test_offset_commit_at_run_full_end.py @@ -0,0 +1,678 @@ +""" +Тесты для нового механизма коммита offset'ов в конце run_full. + +Покрываемые сценарии: +1. Одна трансформация - offset коммитится в конце run_full +2. Цепочка трансформаций - offset'ы коммитятся независимо для каждой трансформации +3. Многотабличная трансформация - offset'ы обновляются для всех входных таблиц +4. Частичная обработка (сбой mid-batch) - offset НЕ обновляется +5. Пустой результат - offset НЕ обновляется +6. run_changelist - offset НЕ обновляется +""" +import time + +import pandas as pd +import pytest +from sqlalchemy import Column, Integer, String + +from datapipe.compute import ComputeInput +from datapipe.datatable import DataStore +from datapipe.step.batch_transform import BatchTransformStep +from datapipe.store.database import DBConn, TableStoreDB +from datapipe.types import ChangeList + + +def test_single_transform_offset_committed_at_end(dbconn: DBConn): + """ + Тест: Одна трансформация, offset коммитится в конце run_full. + + Сценарий: + 1. Создать 12 записей с разными update_ts + 2. Запустить run_full (batch_size=5) - 3 батча + 3. Проверить что offset = MAX(update_ts) из всех обработанных записей + 4. Повторный запуск - выборка пустая + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Входная таблица + input_store = TableStoreDB( + dbconn, + "single_input", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("single_input", input_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "single_output", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("single_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="single_transform", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, # 3 батча по 5, 5, 2 + ) + + # Создаем 12 записей с разными update_ts + t1 = time.time() + df1 = pd.DataFrame({"id": [f"rec_{i:02d}" for i in range(5)], "value": list(range(5))}) + input_dt.store_chunk(df1, now=t1) + + time.sleep(0.01) + t2 = time.time() + df2 = pd.DataFrame({"id": [f"rec_{i:02d}" for i in range(5, 10)], "value": list(range(5, 10))}) + input_dt.store_chunk(df2, now=t2) + + time.sleep(0.01) + t3 = time.time() + df3 = pd.DataFrame({"id": [f"rec_{i:02d}" for i in range(10, 12)], "value": list(range(10, 12))}) + input_dt.store_chunk(df3, now=t3) + + # Запускаем run_full + step.run_full(ds) + + # Проверяем что все записи обработаны + output_data = output_dt.get_data() + assert len(output_data) == 12, "Все 12 записей должны быть обработаны" + + # Проверяем offset - должен быть MAX(update_ts) = t3 + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert "single_input" in offsets + offset_value = offsets["single_input"] + assert offset_value >= t3, f"Offset должен быть >= {t3}, получено {offset_value}" + + # Повторный запуск - выборка должна быть пустой + idx_count, idx_gen = step.get_full_process_ids(ds=ds, run_config=None) + assert idx_count == 0, "После полной обработки выборка должна быть пустой" + + +def test_transform_chain_independent_offsets(dbconn: DBConn): + """ + Тест: Цепочка трансформаций, offset'ы коммитятся независимо. + + Сценарий: + 1. Создать цепочку: input -> transform1 -> intermediate -> transform2 -> output + 2. Запустить transform1 полностью + 3. Проверить что offset для input обновлен + 4. Запустить transform2 полностью + 5. Проверить что offset для intermediate обновлен + 6. Offset'ы независимы + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Входная таблица + input_store = TableStoreDB( + dbconn, + "chain_input", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("chain_input", input_store) + + # Промежуточная таблица + intermediate_store = TableStoreDB( + dbconn, + "chain_intermediate", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + intermediate_dt = ds.create_table("chain_intermediate", intermediate_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "chain_output", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("chain_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + # Первая трансформация + step1 = BatchTransformStep( + ds=ds, + name="chain_transform1", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[intermediate_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Вторая трансформация + step2 = BatchTransformStep( + ds=ds, + name="chain_transform2", + func=copy_func, + input_dts=[ComputeInput(dt=intermediate_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Создаем данные в input + t1 = time.time() + df1 = pd.DataFrame({"id": [f"rec_{i:02d}" for i in range(7)], "value": list(range(7))}) + input_dt.store_chunk(df1, now=t1) + + # Запускаем первую трансформацию + step1.run_full(ds) + + # Проверяем offset для input + offsets1 = ds.offset_table.get_offsets_for_transformation(step1.get_name()) + assert "chain_input" in offsets1 + offset1_value = offsets1["chain_input"] + assert offset1_value >= t1 + + # Проверяем что intermediate заполнена + intermediate_data = intermediate_dt.get_data() + assert len(intermediate_data) == 7 + + # Запускаем вторую трансформацию + step2.run_full(ds) + + # Проверяем offset для intermediate + offsets2 = ds.offset_table.get_offsets_for_transformation(step2.get_name()) + assert "chain_intermediate" in offsets2 + offset2_value = offsets2["chain_intermediate"] + + # Получаем update_ts intermediate таблицы + intermediate_meta = intermediate_dt.meta_table.get_metadata() + max_intermediate_ts = intermediate_meta["update_ts"].max() + assert offset2_value >= max_intermediate_ts + + # Проверяем что output заполнена + output_data = output_dt.get_data() + assert len(output_data) == 7 + + # Offset'ы независимы - у каждой трансформации свой набор offset'ов + assert offsets1.keys() != offsets2.keys(), "Offset'ы должны быть независимы для разных трансформаций" + + +def test_multi_table_transform_all_offsets_updated(dbconn: DBConn): + """ + Тест: Многотабличная трансформация, offset'ы для всех входных таблиц. + + Сценарий: + 1. Создать трансформацию с 2 входными таблицами + 2. Добавить данные в обе таблицы + 3. Запустить run_full + 4. Проверить что offset'ы обновлены для ОБЕИХ входных таблиц + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Первая входная таблица + input1_store = TableStoreDB( + dbconn, + "multi_input1", + [ + Column("id", String, primary_key=True), + Column("value1", Integer), + ], + create_table=True, + ) + input1_dt = ds.create_table("multi_input1", input1_store) + + # Вторая входная таблица + input2_store = TableStoreDB( + dbconn, + "multi_input2", + [ + Column("id", String, primary_key=True), + Column("value2", Integer), + ], + create_table=True, + ) + input2_dt = ds.create_table("multi_input2", input2_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "multi_output", + [ + Column("id", String, primary_key=True), + Column("value1", Integer), + Column("value2", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("multi_output", output_store) + + def merge_func(df1, df2): + # Объединяем по id + result = df1.merge(df2, on="id", how="outer") + return result[["id", "value1", "value2"]].fillna(0).astype({"value1": int, "value2": int}) + + step = BatchTransformStep( + ds=ds, + name="multi_transform", + func=merge_func, + input_dts=[ + ComputeInput(dt=input1_dt, join_type="full"), + ComputeInput(dt=input2_dt, join_type="full"), + ], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Создаем данные в первой таблице + t1 = time.time() + df1 = pd.DataFrame({"id": [f"rec_{i:02d}" for i in range(6)], "value1": list(range(6))}) + input1_dt.store_chunk(df1, now=t1) + + time.sleep(0.01) + + # Создаем данные во второй таблице + t2 = time.time() + df2 = pd.DataFrame({"id": [f"rec_{i:02d}" for i in range(3, 9)], "value2": list(range(3, 9))}) + input2_dt.store_chunk(df2, now=t2) + + # Запускаем трансформацию + step.run_full(ds) + + # Проверяем что данные обработаны + output_data = output_dt.get_data() + assert len(output_data) > 0, "Должны быть обработанные записи" + + # Проверяем offset'ы для ОБЕИХ входных таблиц + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert "multi_input1" in offsets, "Offset для первой входной таблицы должен быть установлен" + assert "multi_input2" in offsets, "Offset для второй входной таблицы должен быть установлен" + + offset1_value = offsets["multi_input1"] + offset2_value = offsets["multi_input2"] + + assert offset1_value >= t1, f"Offset для input1 должен быть >= {t1}" + assert offset2_value >= t2, f"Offset для input2 должен быть >= {t2}" + + +def test_partial_processing_offset_not_updated(dbconn: DBConn): + """ + Тест: Частичная обработка (сбой mid-batch), offset НЕ обновляется. + + Сценарий: + 1. Создать 12 записей + 2. Симулировать обработку только первого батча (без run_full) + 3. Проверить что offset НЕ обновился + 4. Запустить полный run_full + 5. Проверить что offset обновился до максимума + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Входная таблица + input_store = TableStoreDB( + dbconn, + "partial_input", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("partial_input", input_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "partial_output", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("partial_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="partial_transform", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Создаем записи + t1 = time.time() + df1 = pd.DataFrame({"id": [f"rec_{i:02d}" for i in range(12)], "value": list(range(12))}) + input_dt.store_chunk(df1, now=t1) + + # Получаем первый батч и обрабатываем его напрямую (симуляция частичной обработки) + idx_count, idx_gen = step.get_full_process_ids(ds=ds, run_config=None) + assert idx_count == 3, "Должно быть 3 батча" + + # Обрабатываем только первый батч (без вызова run_full) + first_batch = next(idx_gen) + step.process_batch(ds, first_batch, run_config=None) + + # Проверяем что offset НЕ обновился (нет вызова run_full) + offsets_after_partial = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert len(offsets_after_partial) == 0, "Offset не должен обновиться при частичной обработке" + + # Теперь запускаем полный run_full + step.run_full(ds) + + # Проверяем что offset обновился + offsets_after_full = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert "partial_input" in offsets_after_full + offset_value = offsets_after_full["partial_input"] + assert offset_value >= t1 + + # Проверяем что все записи обработаны (не дублируются) + output_data = output_dt.get_data() + assert len(output_data) == 12, "Все 12 записей должны быть обработаны ровно один раз" + + +def test_empty_result_offset_not_updated(dbconn: DBConn): + """ + Тест: Пустой результат трансформации, offset НЕ обновляется. + + Сценарий: + 1. Создать данные + 2. Трансформация возвращает пустой DataFrame + 3. Запустить run_full + 4. Проверить что offset НЕ обновился + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Входная таблица + input_store = TableStoreDB( + dbconn, + "empty_input", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("empty_input", input_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "empty_output", + [ + Column("id", String, primary_key=True), + Column("result", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("empty_output", output_store) + + # Трансформация всегда возвращает пустой DataFrame + def empty_func(df): + return pd.DataFrame(columns=["id", "result"]) + + step = BatchTransformStep( + ds=ds, + name="empty_transform", + func=empty_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Создаем данные + now = time.time() + df = pd.DataFrame({"id": [f"rec_{i:02d}" for i in range(7)], "value": list(range(7))}) + input_dt.store_chunk(df, now=now) + + # Запускаем run_full + step.run_full(ds) + + # Проверяем что output пустой + output_data = output_dt.get_data() + assert len(output_data) == 0, "Output должен быть пустым" + + # Проверяем что offset НЕ обновился + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert len(offsets) == 0, "Offset не должен обновиться при пустом результате" + + +def test_run_changelist_does_not_update_offset(dbconn: DBConn): + """ + Тест: run_changelist НЕ обновляет offset. + + Сценарий: + 1. Создать 3 записи (rec_1, rec_2, rec_3) с update_ts = T1, T2, T3 + 2. Запустить run_full - offset = T3 + 3. Создать новые записи (rec_4, rec_5) с update_ts = T4, T5 + 4. Запустить run_changelist только для rec_5 + 5. Проверить что offset остался = T3 (НЕ перепрыгнул на T5) + 6. Запустить run_full + 7. Проверить что rec_4 обработана (не пропущена) + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Входная таблица + input_store = TableStoreDB( + dbconn, + "changelist_input", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("changelist_input", input_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "changelist_output", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("changelist_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="changelist_transform", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Создаем начальные записи + t1 = time.time() + df1 = pd.DataFrame({"id": ["rec_1"], "value": [1]}) + input_dt.store_chunk(df1, now=t1) + + time.sleep(0.01) + t2 = time.time() + df2 = pd.DataFrame({"id": ["rec_2"], "value": [2]}) + input_dt.store_chunk(df2, now=t2) + + time.sleep(0.01) + t3 = time.time() + df3 = pd.DataFrame({"id": ["rec_3"], "value": [3]}) + input_dt.store_chunk(df3, now=t3) + + # Запускаем run_full + step.run_full(ds) + + # Проверяем offset + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_after_full = offsets[input_dt.name] + assert offset_after_full >= t3 + + # Создаем новые записи + time.sleep(0.01) + t4 = time.time() + df4 = pd.DataFrame({"id": ["rec_4"], "value": [4]}) + input_dt.store_chunk(df4, now=t4) + + time.sleep(0.01) + t5 = time.time() + df5 = pd.DataFrame({"id": ["rec_5"], "value": [5]}) + input_dt.store_chunk(df5, now=t5) + + # Запускаем run_changelist ТОЛЬКО для rec_5 + change_list = ChangeList() + change_list.append(input_dt.name, df5[["id"]]) + step.run_changelist(ds, change_list) + + # Проверяем что rec_5 обработана + output_data_after_changelist = output_dt.get_data() + assert "rec_5" in output_data_after_changelist["id"].values + + # КРИТИЧНО: Проверяем что offset НЕ изменился (остался T3) + offsets_after_changelist = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_after_changelist = offsets_after_changelist[input_dt.name] + assert offset_after_changelist == offset_after_full, ( + f"Offset не должен измениться после run_changelist. " + f"Было: {offset_after_full}, стало: {offset_after_changelist}" + ) + + # Запускаем run_full - rec_4 должна быть обработана + step.run_full(ds) + + # Проверяем что rec_4 обработана + output_data_final = output_dt.get_data() + assert len(output_data_final) == 5, "Все 5 записей должны быть обработаны" + assert "rec_4" in output_data_final["id"].values, "rec_4 не должна быть пропущена" + + # Проверяем что offset обновился до максимума обработанных записей + # После run_full обрабатываются rec_4 (t4) и rec_5 (t5 - уже была обработана в changelist, но обновилась) + # Offset должен быть >= t4 (так как rec_4 точно обработана в этом run_full) + offsets_final = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_final = offsets_final[input_dt.name] + assert offset_final >= t4, "Offset должен обновиться после run_full (минимум до t4)" + # Offset может быть >= t5 если rec_5 тоже переобработалась + assert offset_final >= offset_after_full, "Offset должен увеличиться по сравнению с предыдущим значением" + + +def test_multiple_batches_offset_is_max(dbconn: DBConn): + """ + Тест: При обработке нескольких батчей offset = MAX(update_ts) из всех батчей. + + Сценарий: + 1. Создать 3 группы записей с разными update_ts: T1, T2, T3 + 2. Размер батча = 5, будет 3 батча + 3. Запустить run_full + 4. Проверить что offset = MAX(T1, T2, T3) = T3 + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Входная таблица + input_store = TableStoreDB( + dbconn, + "maxbatch_input", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("maxbatch_input", input_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "maxbatch_output", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("maxbatch_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="maxbatch_transform", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Создаем записи с разными update_ts + # Группа 1: T1 + t1 = time.time() + df1 = pd.DataFrame({"id": [f"rec_t1_{i:02d}" for i in range(4)], "value": list(range(4))}) + input_dt.store_chunk(df1, now=t1) + + time.sleep(0.01) + + # Группа 2: T2 + t2 = time.time() + df2 = pd.DataFrame({"id": [f"rec_t2_{i:02d}" for i in range(5)], "value": list(range(5))}) + input_dt.store_chunk(df2, now=t2) + + time.sleep(0.01) + + # Группа 3: T3 + t3 = time.time() + df3 = pd.DataFrame({"id": [f"rec_t3_{i:02d}" for i in range(3)], "value": list(range(3))}) + input_dt.store_chunk(df3, now=t3) + + # Запускаем run_full + step.run_full(ds) + + # Проверяем что все записи обработаны + output_data = output_dt.get_data() + assert len(output_data) == 12, "Все 12 записей должны быть обработаны" + + # Проверяем offset - должен быть MAX(T1, T2, T3) = T3 + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_value = offsets[input_dt.name] + assert offset_value >= t3, f"Offset должен быть >= MAX(T1,T2,T3) = {t3}, получено {offset_value}" + + # Проверяем что offset не меньше минимального + assert offset_value >= t1, f"Offset должен быть >= T1 = {t1}" diff --git a/tests/offset_edge_cases/test_offset_custom_ordering.py b/tests/offset_edge_cases/test_offset_custom_ordering.py new file mode 100644 index 00000000..74f65a91 --- /dev/null +++ b/tests/offset_edge_cases/test_offset_custom_ordering.py @@ -0,0 +1,143 @@ +""" +Тесты для кастомной сортировки (order_by) с offset-оптимизацией. + +Проблема: +При передаче кастомного order_by новый порядок по update_ts пропадает: +сортируется только по переданным колонкам. Это создает риск нехронологического +чтения данных, что может привести к пропуску записей при использовании offset. + +Решение: +Всегда префиксовать order_by колонкой update_ts ASC NULLS LAST перед +пользовательским order_by для обеспечения хронологического порядка обработки. + +Код исправления в sql_meta.py (строки 1194-1204): +- При наличии кастомного order_by всегда добавляется сортировка по update_ts перед пользовательской +""" +import time + +import pandas as pd +import pytest +from sqlalchemy import Column, Integer, String + +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_custom_order_by_preserves_update_ts_ordering(dbconn: DBConn): + """ + Проблема: При передаче кастомного order_by теряется сортировка по update_ts. + + Сценарий: + 1. Создать записи с разными update_ts в обратном порядке по id + - rec_3 создается раньше (меньший update_ts = T1) + - rec_2 создается позже (update_ts = T2) + - rec_1 создается последним (больший update_ts = T3) + 2. Передать order_by=['id'] (сортировка по id) + 3. Проверить что записи обрабатываются в порядке update_ts (T1, T2, T3), + а не по id (rec_1, rec_2, rec_3) + + Ожидание: order_by должен всегда префиксоваться update_ts ASC NULLS LAST + """ + # SQLite имеет проблемы с NULLS LAST в сложных запросах с CTE + # Пропускаем для SQLite, так как в production используется PostgreSQL + if "sqlite" in dbconn.connstr.lower() or "pysqlite" in dbconn.connstr.lower(): + pytest.skip("SQLite does not handle NULLS LAST in complex ORDER BY with CTE correctly") + + ds = DataStore(dbconn, create_meta_table=True) + + # Создаем входную таблицу + input_store = TableStoreDB( + dbconn, + "order_test_input", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("order_test_input", input_store) + + # Создаем выходную таблицу для отслеживания порядка + output_store = TableStoreDB( + dbconn, + "order_test_output", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + Column("processing_order", Integer), # Порядок обработки + ], + create_table=True, + ) + output_dt = ds.create_table("order_test_output", output_store) + + # Счетчик для отслеживания порядка обработки + processing_counter = {"count": 0} + + def tracking_func(df): + result = df[["id", "value"]].copy() + # Добавляем порядок обработки + result["processing_order"] = list(range(processing_counter["count"], processing_counter["count"] + len(df))) + processing_counter["count"] += len(df) + return result + + step = BatchTransformStep( + ds=ds, + name="order_test", + func=tracking_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=10, + order_by=["id"], # Кастомный порядок - сортировка по id + order="asc", + ) + + # Создаем записи с update_ts в обратном порядке по id + # id = "rec_3" создается раньше (меньший update_ts = T1) + # id = "rec_1" создается позже (больший update_ts = T3) + t1 = time.time() + df1 = pd.DataFrame({"id": ["rec_3"], "value": [3]}) + input_dt.store_chunk(df1, now=t1) + + time.sleep(0.1) # Увеличена задержка для надежности в CI + t2 = time.time() + df2 = pd.DataFrame({"id": ["rec_2"], "value": [2]}) + input_dt.store_chunk(df2, now=t2) + + time.sleep(0.1) # Увеличена задержка для надежности в CI + t3 = time.time() + df3 = pd.DataFrame({"id": ["rec_1"], "value": [1]}) + input_dt.store_chunk(df3, now=t3) + + # Проверяем update_ts + meta = input_dt.meta_table.get_metadata(pd.DataFrame({"id": ["rec_1", "rec_2", "rec_3"]})) + meta_sorted = meta.sort_values("id") + assert t3 > t2 > t1 + assert meta_sorted[meta_sorted["id"] == "rec_1"].iloc[0]["update_ts"] == t3 + assert meta_sorted[meta_sorted["id"] == "rec_2"].iloc[0]["update_ts"] == t2 + assert meta_sorted[meta_sorted["id"] == "rec_3"].iloc[0]["update_ts"] == t1 + + # Обработать + step.run_full(ds) + + # Проверить порядок обработки + output_data = output_dt.get_data() + assert len(output_data) == 3 + + # Сортируем по порядку обработки + output_sorted = output_data.sort_values("processing_order") + + # Ожидаем что записи обработаны в порядке update_ts: rec_3 (t1), rec_2 (t2), rec_1 (t3) + # Несмотря на order_by=['id'], сортировка по update_ts должна быть приоритетнее + expected_order = ["rec_3", "rec_2", "rec_1"] + actual_order = output_sorted["id"].tolist() + + assert actual_order == expected_order, ( + f"Записи должны обрабатываться в порядке update_ts, а не id. " + f"Ожидалось: {expected_order}, получено: {actual_order}. " + f"Это означает что кастомный order_by=['id'] игнорирует сортировку по update_ts, " + f"что может привести к пропуску данных при использовании offset." + ) diff --git a/tests/offset_edge_cases/test_offset_delete_ts_filtering.py b/tests/offset_edge_cases/test_offset_delete_ts_filtering.py new file mode 100644 index 00000000..ff1e59a1 --- /dev/null +++ b/tests/offset_edge_cases/test_offset_delete_ts_filtering.py @@ -0,0 +1,129 @@ +""" +Тесты для обработки удалений (delete_ts) в offset-оптимизации. + +Проблема: +При фильтрации повторной обработки (update_ts > process_ts) не учитывается delete_ts. +Если запись была обработана, а затем удалена (delete_ts установлен, update_ts не изменился), +то условие update_ts > process_ts будет ложным и удаление не попадет в выборку. + +Решение: +Использовать coalesce(update_ts, delete_ts) > process_ts или greatest(update_ts, delete_ts) > process_ts +для учета обоих типов изменений. + +Код исправления в sql_meta.py: +- В фильтре changed records добавлена проверка delete_ts >= offset +- В фильтре повторной обработки используется OR для учета удалений +""" +import time + +import pandas as pd +from sqlalchemy import Column, Integer, String + +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_delete_ts_not_processed_after_delete(dbconn: DBConn): + """ + Проблема: delete_ts не учитывается при фильтрации повторной обработки. + + Сценарий: + 1. Создать запись с update_ts = T1 + 2. Обработать (process_ts = T_proc) + 3. Удалить запись (delete_ts = T_del > T_proc), update_ts не меняется + 4. Повторно запустить - удаление должно попасть в выборку + + Ожидание: Фильтр должен проверять coalesce(update_ts, delete_ts) > process_ts + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Создаем входную таблицу + input_store = TableStoreDB( + dbconn, + "delete_test_input", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("delete_test_input", input_store) + + # Создаем выходную таблицу + output_store = TableStoreDB( + dbconn, + "delete_test_output", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("delete_test_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="delete_test_copy", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=10, + ) + + # 1. Создать запись + t1 = time.time() + df = pd.DataFrame({"id": ["rec_1"], "value": [100]}) + input_dt.store_chunk(df, now=t1) + + # Проверяем метаданные + meta = input_dt.meta_table.get_metadata(df[["id"]]) + assert len(meta) == 1 + assert meta.iloc[0]["update_ts"] == t1 + assert pd.isna(meta.iloc[0]["delete_ts"]) + + # 2. Обработать запись + time.sleep(0.01) + step.run_full(ds) + + # Проверяем что запись обработана + output_data = output_dt.get_data() + assert len(output_data) == 1 + assert output_data.iloc[0]["id"] == "rec_1" + + # Записываем t_proc для дальнейшей проверки + time.sleep(0.01) + t_proc = time.time() + + # 3. Удалить запись (delete_ts устанавливается, update_ts НЕ меняется) + time.sleep(0.01) + t_del = time.time() + input_dt.delete_by_idx(df[["id"]], now=t_del) + + # Проверяем метаданные - delete_ts установлен + meta_after_delete = input_dt.meta_table.get_metadata(df[["id"]], include_deleted=True) + assert len(meta_after_delete) == 1 + assert pd.notna(meta_after_delete.iloc[0]["delete_ts"]), "delete_ts должен быть установлен" + delete_ts_value = meta_after_delete.iloc[0]["delete_ts"] + update_ts_after_delete = meta_after_delete.iloc[0]["update_ts"] + + # Главная проблема: если delete_ts > process_ts, но update_ts <= process_ts, + # то фильтр (update_ts > process_ts) не пропустит эту запись + # Нужно использовать coalesce(update_ts, delete_ts) > process_ts + + # 4. Повторно запустить - удаление должно попасть в выборку + time.sleep(0.01) + + # Проверяем что батчи для обработки есть (есть удаленная запись) + idx_count, idx_gen = step.get_full_process_ids(ds=ds, run_config=None) + assert idx_count > 0, ( + f"БАГ: После удаления записи выборка пустая (idx_count={idx_count}). " + f"Это происходит потому что фильтр проверяет только update_ts > process_ts, " + f"но не учитывает delete_ts. update_ts={update_ts_after_delete}, delete_ts={delete_ts_value}, t_proc={t_proc}" + ) diff --git a/tests/offset_edge_cases/test_offset_first_run_bug.py b/tests/offset_edge_cases/test_offset_first_run_bug.py new file mode 100644 index 00000000..17fe8700 --- /dev/null +++ b/tests/offset_edge_cases/test_offset_first_run_bug.py @@ -0,0 +1,366 @@ +""" +Тест воспроизводит РЕАЛЬНЫЙ баг из production. + +ПРОБЛЕМА: +При первом запуске трансформации с offset optimization, +если записи обрабатываются в порядке ORDER BY (id, hashtag) а не по update_ts, +и в батч попадают записи с РАЗНЫМИ update_ts (созданные в разное время), +то offset устанавливается на MAX(update_ts) из батча. + +Все записи с id ПОСЛЕ последней обработанной, но с update_ts < offset, +будут ПРОПУЩЕНЫ при следующих запусках! + +РЕАЛЬНЫЙ СЦЕНАРИЙ ИЗ PRODUCTION (hashtag_issue.md): +- 16:21 - Пост b927ca71 создан, хештеги извлечены +- 20:29 - Пост e26f9c4b создан +- copy_to_online ПЕРВЫЙ РАЗ обрабатывает: + - Батч содержит (в порядке id): b927ca71(16:21), e26f9c4b(20:29), ... + - offset = MAX(16:21, 20:29) = 20:29 +- Следующий запуск: WHERE update_ts > 20:29 + - Пропускаются ВСЕ записи с id > e26f9c4b и update_ts < 20:29! + - Результат: 60% данных потеряно + +Этот тест воспроизводит эту проблему. +""" +import time + +import pandas as pd +import pytest +from sqlalchemy import Column, Integer, String + +from datapipe.compute import ComputeInput +from datapipe.datatable import DataStore +from datapipe.step.batch_transform import BatchTransformStep +from datapipe.store.database import DBConn, TableStoreDB + + +@pytest.mark.xfail(reason="Test uses run_idx() which no longer commits offsets (offsets only commit at end of run_full)") +def test_first_run_with_mixed_update_ts_and_order_by_id(dbconn: DBConn): + """ + Воспроизводит ТОЧНЫЙ сценарий production бага. + + Симуляция накопления данных за несколько часов, + затем первый запуск copy_to_online который обрабатывает + записи в порядке (id, hashtag), НЕ по update_ts. + + Результат: данные с id после "границы батча" но с старым update_ts ТЕРЯЮТСЯ. + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "first_run_input", + [ + Column("id", String, primary_key=True), + Column("hashtag", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("first_run_input", input_store) + + output_store = TableStoreDB( + dbconn, + "first_run_output", + [ + Column("id", String, primary_key=True), + Column("hashtag", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("first_run_output", output_store) + + def copy_func(df): + return df[["id", "hashtag", "value"]] + + step = BatchTransformStep( + ds=ds, + name="first_run_copy", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id", "hashtag"], + use_offset_optimization=True, + chunk_size=10, # Маленький размер для демонстрации + ) + + # ========== Симулируем накопление данных за несколько часов ========== + base_time = time.time() + + # Имитируем посты которые приходили в течение 4 часов + # 16:21 - Пост b927ca71 (UUID начинается с 'b') + t_16_21 = base_time + 1 + input_dt.store_chunk( + pd.DataFrame({ + "id": ["b927ca71-0001", "b927ca71-0001"], + "hashtag": ["322", "anime"], + "value": [1, 2], + }), + now=t_16_21 + ) + + time.sleep(0.001) + + # 17:00 - Еще посты с разными id, но старыми timestamps + t_17_00 = base_time + 2 + input_dt.store_chunk( + pd.DataFrame({ + "id": ["a111aaaa-0002", "c222cccc-0003", "d333dddd-0004"], + "hashtag": ["test1", "test2", "test3"], + "value": [3, 4, 5], + }), + now=t_17_00 + ) + + time.sleep(0.001) + + # 18:00 - Больше постов + t_18_00 = base_time + 3 + input_dt.store_chunk( + pd.DataFrame({ + "id": ["e444eeee-0005", "f555ffff-0006", "g666gggg-0007"], + "hashtag": ["hash1", "hash2", "hash3"], + "value": [6, 7, 8], + }), + now=t_18_00 + ) + + time.sleep(0.001) + + # 20:29 - Новый пост e26f9c4b (UUID начинается с 'e') + t_20_29 = base_time + 4 + input_dt.store_chunk( + pd.DataFrame({ + "id": ["e26f9c4b-0008"], + "hashtag": ["looky"], + "value": [9], + }), + now=t_20_29 + ) + + time.sleep(0.001) + + # 20:30 - Еще несколько постов ПОСЛЕ e26f9c4b, но с РАЗНЫМИ timestamps + # Эти посты критичны - у них id > e26f9c4b, но update_ts может быть старым! + t_20_30 = base_time + 5 + input_dt.store_chunk( + pd.DataFrame({ + "id": ["h777hhhh-0009", "i888iiii-0010", "j999jjjj-0011"], + "hashtag": ["new1", "new2", "new3"], + "value": [10, 11, 12], + }), + now=t_20_30 # Новое время! + ) + + time.sleep(0.001) + + # Критично: добавляем записи с id МЕЖДУ уже созданными, но со СТАРЫМ timestamp + # Симулируем ситуацию где записи приходят не в порядке id + t_19_00 = base_time + 2.5 # Старый timestamp (между 17:00 и 18:00) + input_dt.store_chunk( + pd.DataFrame({ + "id": ["f111ffff-late1", "f222ffff-late2"], # id в середине диапазона + "hashtag": ["late1", "late2"], + "value": [98, 99], + }), + now=t_19_00 # СТАРЫЙ timestamp! + ) + + time.sleep(0.001) + + # Добавляем еще записей для полного второго батча (чтобы было 20+ записей → 2 батча) + t_20_31 = base_time + 5.1 + input_dt.store_chunk( + pd.DataFrame({ + "id": ["k111kkkk-0012", "l222llll-0013", "m333mmmm-0014", + "n444nnnn-0015", "o555oooo-0016", "p666pppp-0017"], + "hashtag": ["extra1", "extra2", "extra3", "extra4", "extra5", "extra6"], + "value": [12, 13, 14, 15, 16, 17], + }), + now=t_20_31 + ) + + # ========== Проверяем накопленные данные ========== + all_meta = input_dt.meta_table.get_metadata() + print(f"\nВсего записей накоплено: {len(all_meta)}") + print("Распределение по update_ts:") + for idx, row in all_meta.sort_values("id").iterrows(): + print(f" id={row['id'][:15]:15} hashtag={row['hashtag']:10} update_ts={row['update_ts']:.2f}") + + # ========== ПЕРВЫЙ ЗАПУСК copy_to_online ========== + # Имитируем прерывание: обработаем только первые chunk_size=10 записей + (idx_count, idx_gen) = step.get_full_process_ids(ds=ds, run_config=None) + print(f"\nБатчей доступно для обработки: {idx_count}") + + # Обрабатываем ТОЛЬКО первый батч (как если бы джоба прервалась) + first_batch_idx = next(idx_gen) + idx_gen.close() # Закрываем генератор + print(f"Обрабатываем первый батч, размер: {len(first_batch_idx)}") + step.run_idx(ds=ds, idx=first_batch_idx, run_config=None) + + # Получаем offset после первого запуска + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_after_first = offsets["first_run_input"] + + # Проверяем что обработано + output_after_first = output_dt.get_data() + processed_ids = set(output_after_first["id"].tolist()) + + print(f"\n=== ПОСЛЕ ПЕРВОГО ЗАПУСКА ===") + print(f"Обработано записей: {len(output_after_first)}") + print(f"offset установлен на: {offset_after_first:.2f}") + print(f"Обработанные id: {sorted(processed_ids)}") + + # ========== КРИТИЧНО: Какие записи НЕ обработаны? ========== + all_input_ids = set(all_meta["id"].tolist()) + unprocessed_ids = all_input_ids - processed_ids + + if unprocessed_ids: + print(f"\n=== НЕОБРАБОТАННЫЕ ЗАПИСИ ===") + unprocessed_meta = all_meta[all_meta["id"].isin(unprocessed_ids)] + for idx, row in unprocessed_meta.sort_values("id").iterrows(): + below_offset = row["update_ts"] < offset_after_first + status = "БУДЕТ ПОТЕРЯНА!" if below_offset else "будет обработана" + print( + f" id={row['id'][:15]:15} update_ts={row['update_ts']:.2f} " + f"< offset={offset_after_first:.2f} ? {below_offset} → {status}" + ) + + # ========== ВТОРОЙ ЗАПУСК ========== + # Получаем оставшиеся батчи (с учетом offset) + (idx_count_second, idx_gen_second) = step.get_full_process_ids(ds=ds, run_config=None) + print(f"\n=== ВТОРОЙ ЗАПУСК ===") + print(f"Батчей доступно для обработки: {idx_count_second}") + + if idx_count_second > 0: + # Обрабатываем оставшиеся батчи + for idx in idx_gen_second: + print(f"Обрабатываем батч, размер: {len(idx)}") + step.run_idx(ds=ds, idx=idx, run_config=None) + idx_gen_second.close() + + # ========== КРИТИЧНАЯ ПРОВЕРКА: ВСЕ записи должны быть обработаны ========== + final_output = output_dt.get_data() + final_processed_ids = set(final_output["id"].tolist()) + + # Список потерянных записей + lost_records = all_input_ids - final_processed_ids + + if lost_records: + lost_meta = all_meta[all_meta["id"].isin(lost_records)] + print(f"\n=== 🚨 ПОТЕРЯННЫЕ ЗАПИСИ (БАГ!) ===") + for idx, row in lost_meta.sort_values("id").iterrows(): + print( + f" id={row['id'][:15]:15} hashtag={row['hashtag']:10} " + f"update_ts={row['update_ts']:.2f} < offset={offset_after_first:.2f}" + ) + + pytest.fail( + f"КРИТИЧЕСКИЙ БАГ ВОСПРОИЗВЕДЕН: {len(lost_records)} записей ПОТЕРЯНЫ!\n" + f"Ожидалось: {len(all_input_ids)} записей\n" + f"Получено: {len(final_output)} записей\n" + f"Потеряно: {len(lost_records)} записей\n" + f"Потерянные id: {sorted(lost_records)}\n\n" + f"Это ТОЧНО воспроизводит production баг где было потеряно 60% данных!" + ) + + print(f"\n=== ФИНАЛЬНЫЙ РЕЗУЛЬТАТ ===") + print(f"Всего записей в input: {len(all_input_ids)}") + print(f"Обработано в output: {len(final_output)}") + print(f"✅ Все записи обработаны корректно!") + + +def test_first_run_invariant_all_records_below_offset_must_be_processed(dbconn: DBConn): + """ + Проверяет инвариант для первого запуска: + После первого запуска ВСЕ записи с update_ts <= offset должны быть обработаны. + + Это более общий тест который проверяет что независимо от: + - Порядка создания записей + - Порядка их id + - Размера батча + + После установки offset НЕ ДОЛЖНО остаться необработанных записей с update_ts < offset. + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "invariant_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("invariant_input", input_store) + + output_store = TableStoreDB( + dbconn, + "invariant_output", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + output_dt = ds.create_table("invariant_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="invariant_copy", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Создаем записи с разными update_ts и разными id (не в порядке времени) + base_time = time.time() + + records = [ + ("z999", base_time + 1), # Поздний id, ранний timestamp + ("a111", base_time + 5), # Ранний id, поздний timestamp + ("m555", base_time + 2), # Средний id, средний timestamp + ("b222", base_time + 3), + ("y888", base_time + 1.5), + ("c333", base_time + 4), + ("x777", base_time + 2.5), + ] + + for record_id, timestamp in records: + input_dt.store_chunk( + pd.DataFrame({"id": [record_id], "value": [int(timestamp)]}), + now=timestamp + ) + time.sleep(0.001) + + # Первый запуск + step.run_full(ds) + + # Получаем offset + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + current_offset = offsets["invariant_input"] + + # ИНВАРИАНТ: ВСЕ записи с update_ts <= current_offset должны быть обработаны + all_meta = input_dt.meta_table.get_metadata() + output_data = output_dt.get_data() + processed_ids = set(output_data["id"].tolist()) + + violations = [] + for idx, row in all_meta.iterrows(): + if row["update_ts"] <= current_offset: + if row["id"] not in processed_ids: + violations.append(row) + + if violations: + print(f"\n🚨 НАРУШЕНИЕ ИНВАРИАНТА!") + print(f"offset = {current_offset}") + print(f"Необработанные записи с update_ts <= offset:") + for row in violations: + print(f" id={row['id']} update_ts={row['update_ts']}") + + pytest.fail( + f"НАРУШЕНИЕ ИНВАРИАНТА: {len(violations)} записей с update_ts <= offset НЕ обработаны!\n" + f"Это означает потерю данных." + ) diff --git a/tests/offset_edge_cases/test_offset_invariants.py b/tests/offset_edge_cases/test_offset_invariants.py new file mode 100644 index 00000000..cb403d69 --- /dev/null +++ b/tests/offset_edge_cases/test_offset_invariants.py @@ -0,0 +1,403 @@ +""" +Тесты на инварианты offset optimization. + +ГЛАВНЫЙ ИНВАРИАНТ: +После каждой итерации трансформации, ВСЕ записи с update_ts <= offset должны быть обработаны. +НЕ ДОЛЖНО быть необработанных записей с update_ts < offset. + +Проверяем два режима: +1. Синхронное выполнение (1 под) - последовательная обработка +2. Асинхронное выполнение (несколько подов) - параллельная обработка + +РЕАЛЬНЫЙ СЦЕНАРИЙ ИЗ PRODUCTION: +- Записи создаются с разными update_ts +- Батч обрабатывает их в порядке (id, hashtag), НЕ по update_ts +- offset = MAX(update_ts) из батча +- МЕЖДУ итерациями могут появляться новые записи +- Эти записи НЕ ДОЛЖНЫ пропускаться +""" +import time +from concurrent.futures import ThreadPoolExecutor, as_completed + +import pandas as pd +import pytest +from sqlalchemy import Column, Integer, String + +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_invariant_synchronous(dbconn: DBConn): + """ + Тест инварианта в синхронном режиме. + + Проверяем что после каждой итерации: + - offset обновляется корректно + - НЕТ необработанных записей с update_ts < offset + - update_ts новых записей всегда >= предыдущего offset + + Это базовый тест - если он падает, система работает некорректно. + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "invariant_input", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("invariant_input", input_store) + + output_store = TableStoreDB( + dbconn, + "invariant_output", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("invariant_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="invariant_copy", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, # Маленький размер батча для тестирования + ) + + previous_offset = 0.0 + + # Выполняем 5 итераций создания данных и обработки + # NOTE: Уменьшено с 10 до 5 итераций для ускорения тестов (каждая итерация делает run_full()) + for iteration in range(5): + time.sleep(0.01) # Небольшая задержка между итерациями + + # Создаем новые записи + current_time = time.time() + new_records = pd.DataFrame({ + "id": [f"iter{iteration}_rec{i}" for i in range(3)], + "value": [iteration * 10 + i for i in range(3)], + }) + + input_dt.store_chunk(new_records, now=current_time) + + # Проверяем метаданные созданных записей + meta_df = input_dt.meta_table.get_metadata(new_records[["id"]]) + for idx, row in meta_df.iterrows(): + record_update_ts = row["update_ts"] + + # ИНВАРИАНТ 1: update_ts новых записей должен быть >= предыдущего offset + assert record_update_ts >= previous_offset, ( + f"НАРУШЕНИЕ ИНВАРИАНТА (итерация {iteration}): " + f"Новая запись {row['id']} имеет update_ts={record_update_ts} < " + f"предыдущий offset={previous_offset}!" + ) + + # Запускаем трансформацию + step.run_full(ds) + + # Получаем текущий offset + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + current_offset = offsets.get("invariant_input", 0.0) + + # ИНВАРИАНТ 2: offset должен расти монотонно (или оставаться прежним) + assert current_offset >= previous_offset, ( + f"НАРУШЕНИЕ ИНВАРИАНТА (итерация {iteration}): " + f"offset уменьшился! previous={previous_offset}, current={current_offset}" + ) + + # ИНВАРИАНТ 3: ВСЕ записи с update_ts <= current_offset должны быть обработаны + # Проверяем: нет ли необработанных записей в input с update_ts < current_offset + all_input_meta = input_dt.meta_table.get_metadata() + all_output_ids = set(output_dt.get_data()["id"].tolist()) if len(output_dt.get_data()) > 0 else set() + + for idx, row in all_input_meta.iterrows(): + if row["update_ts"] <= current_offset: + # Эта запись должна быть обработана! + assert row["id"] in all_output_ids, ( + f"НАРУШЕНИЕ ИНВАРИАНТА (итерация {iteration}): " + f"Запись {row['id']} с update_ts={row['update_ts']} <= offset={current_offset} " + f"НЕ обработана! Это означает потерю данных." + ) + + print(f"Итерация {iteration}: offset={current_offset}, обработано {len(all_output_ids)} записей") + previous_offset = current_offset + + # Финальная проверка: все записи обработаны + total_records = 5 * 3 # 5 итераций по 3 записи + final_output = output_dt.get_data() + assert len(final_output) == total_records, ( + f"Ожидалось {total_records} записей в output, получено {len(final_output)}" + ) + + +def test_offset_invariant_concurrent(dbconn: DBConn): + """ + Тест инварианта в асинхронном режиме (несколько подов параллельно). + + Сценарий: + - Несколько потоков одновременно: + 1. Создают записи (с текущим time.time()) + 2. Запускают трансформацию + - Проверяем что offset корректно обрабатывает race conditions + + ПОТЕНЦИАЛЬНАЯ ПРОБЛЕМА: + - Поток 1: создал запись с update_ts=T1, запустил обработку + - Поток 2 (параллельно): создал запись с update_ts=T2 (T2 > T1) + - Поток 2: обработал раньше, установил offset=T2 + - Поток 1: обработал позже, но его запись с update_ts=T1 < offset=T2 + - Результат: может быть проблема с видимостью данных + """ + # SQLite имеет проблемы с concurrent connections, пропускаем для SQLite + if "sqlite" in dbconn.connstr.lower() or "pysqlite" in dbconn.connstr.lower(): + pytest.skip("SQLite does not support concurrent writes reliably") + + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "concurrent_input", + [ + Column("id", String, primary_key=True), + Column("thread_id", Integer), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("concurrent_input", input_store) + + output_store = TableStoreDB( + dbconn, + "concurrent_output", + [ + Column("id", String, primary_key=True), + Column("thread_id", Integer), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("concurrent_output", output_store) + + def copy_func(df): + return df[["id", "thread_id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="concurrent_copy", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=10, + ) + + def worker_thread(thread_id, iterations): + """ + Рабочий поток: создает записи и запускает трансформацию. + Симулирует работу отдельного пода. + """ + created_ids = [] + + for i in range(iterations): + # Небольшая случайная задержка чтобы потоки не синхронизировались + time.sleep(0.001 * (thread_id % 3)) + + # Создаем запись с ТЕКУЩИМ временем + record_id = f"thread{thread_id}_iter{i}" + record = pd.DataFrame({ + "id": [record_id], + "thread_id": [thread_id], + "value": [thread_id * 1000 + i], + }) + + # Записываем с текущим now (как в реальной системе) + input_dt.store_chunk(record, now=time.time()) + created_ids.append(record_id) + + # Запускаем трансформацию (каждый под обрабатывает данные) + step.run_full(ds) + + return created_ids + + # Запускаем несколько потоков параллельно (симулируем несколько подов) + # NOTE: Этот тест медленный (каждая итерация делает run_full()). + # Баланс: 2 threads × 6 iterations = 12 records → 2 batches при chunk_size=10 + # Это достаточно для тестирования concurrent behavior без чрезмерной медленности. + num_threads = 2 + iterations_per_thread = 6 + + with ThreadPoolExecutor(max_workers=num_threads) as executor: + futures = [ + executor.submit(worker_thread, thread_id, iterations_per_thread) + for thread_id in range(num_threads) + ] + + all_created_ids = [] + for future in as_completed(futures): + all_created_ids.extend(future.result()) + + # Финальный запуск для обработки оставшихся данных + step.run_full(ds) + + # Получаем финальный offset + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + final_offset = offsets.get("concurrent_input", 0.0) + + # КРИТИЧНАЯ ПРОВЕРКА: ВСЕ записи должны быть обработаны + final_output = output_dt.get_data() + processed_ids = set(final_output["id"].tolist()) + + expected_count = num_threads * iterations_per_thread + assert len(final_output) == expected_count, ( + f"Ожидалось {expected_count} записей, получено {len(final_output)}" + ) + + # Проверяем что КАЖДАЯ созданная запись обработана + for record_id in all_created_ids: + assert record_id in processed_ids, ( + f"КРИТИЧЕСКИЙ БАГ (CONCURRENT): Запись {record_id} была создана но НЕ обработана!" + ) + + # ИНВАРИАНТ: Нет необработанных записей с update_ts <= final_offset + all_input_meta = input_dt.meta_table.get_metadata() + for idx, row in all_input_meta.iterrows(): + if row["update_ts"] <= final_offset: + assert row["id"] in processed_ids, ( + f"НАРУШЕНИЕ ИНВАРИАНТА (CONCURRENT): " + f"Запись {row['id']} с update_ts={row['update_ts']} <= final_offset={final_offset} " + f"НЕ обработана при параллельном выполнении!" + ) + + print(f"Concurrent test: {expected_count} записей обработано, final_offset={final_offset}") + + +@pytest.mark.xfail( + reason="Known limitation: records with update_ts < offset are lost (Hypothesis 4). " + "Warning added in get_changes_for_store_chunk() to detect this edge case." +) +def test_offset_with_delayed_records(dbconn: DBConn): + """ + Тест проверяет сценарий когда запись создается "между итерациями". + + Сценарий: + 1. Обрабатываем батч 1, offset устанавливается на MAX(update_ts) = T3 + 2. МЕЖДУ итерациями создается запись с update_ts = T2 (T2 < T3, но запись НОВАЯ) + 3. Следующая итерация должна обработать эту запись + + ВОПРОС: Может ли это произойти на проде? + - Если время создания = time.time(), то update_ts всегда растет + - НО: если несколько серверов с разными часами... + - ИЛИ: если система обрабатывает данные из очереди с задержкой... + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "delayed_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("delayed_input", input_store) + + output_store = TableStoreDB( + dbconn, + "delayed_output", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + output_dt = ds.create_table("delayed_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="delayed_copy", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # Создаем записи с монотонно растущими timestamps + base_time = time.time() + + # Первый батч: записи с T1, T2, T3 + input_dt.store_chunk( + pd.DataFrame({"id": ["rec_T1"], "value": [1]}), + now=base_time + 1 + ) + input_dt.store_chunk( + pd.DataFrame({"id": ["rec_T2"], "value": [2]}), + now=base_time + 2 + ) + input_dt.store_chunk( + pd.DataFrame({"id": ["rec_T3"], "value": [3]}), + now=base_time + 3 + ) + + # Первая итерация + step.run_full(ds) + + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_after_first = offsets["delayed_input"] + + # offset должен быть >= T3 + assert offset_after_first >= base_time + 3 + + # Теперь создаем запись с update_ts МЕЖДУ T2 и T3 + # Вопрос: как это может произойти если now=time.time()? + # Ответ: НЕ МОЖЕТ в нормальной ситуации! + # + # Но проверим что ЕСЛИ это произойдет, система обработает корректно + time.sleep(0.01) + + # Создаем "запоздалую" запись (симулируем запись от другого сервера с отстающими часами) + delayed_time = base_time + 2.5 # Между T2 и T3 + input_dt.store_chunk( + pd.DataFrame({"id": ["rec_delayed"], "value": [999]}), + now=delayed_time + ) + + # Проверяем метаданные + delayed_meta = input_dt.meta_table.get_metadata(pd.DataFrame({"id": ["rec_delayed"]})) + delayed_update_ts = delayed_meta.iloc[0]["update_ts"] + + # КРИТИЧНО: update_ts < offset! + if delayed_update_ts < offset_after_first: + # Запись НЕ ВИДНА для следующей итерации + changed_count = step.get_changed_idx_count(ds) + + # Эта проверка покажет есть ли баг + print(f"ПРЕДУПРЕЖДЕНИЕ: Запись с update_ts={delayed_update_ts} < offset={offset_after_first}") + print(f"changed_count={changed_count} (ожидается 0 - запись не видна)") + + # Запускаем обработку + step.run_full(ds) + + # Проверяем результат + final_output = output_dt.get_data() + output_ids = set(final_output["id"].tolist()) + + # КРИТИЧНАЯ ПРОВЕРКА + if "rec_delayed" not in output_ids: + pytest.fail( + f"КРИТИЧЕСКИЙ БАГ: Запись 'rec_delayed' с update_ts={delayed_update_ts} < " + f"offset={offset_after_first} НЕ обработана! " + f"Это означает что система теряет данные при создании записей 'из прошлого'." + ) diff --git a/tests/offset_edge_cases/test_offset_production_bug.py b/tests/offset_edge_cases/test_offset_production_bug.py new file mode 100644 index 00000000..77b8d395 --- /dev/null +++ b/tests/offset_edge_cases/test_offset_production_bug.py @@ -0,0 +1,778 @@ +""" +Тесты для воспроизведения production бага с offset optimization. + +Production issue: +- После деплоя offset optimization было потеряно 60% данных (48,915 из 82,000 записей) +- Корневая причина: update_ts vs process_ts расхождение при чтении таблиц +- Промежуточные трансформации обновляют process_ts но не update_ts +- Offset фильтрация использует update_ts, пропуская "устаревшие" записи + +Эти тесты ДОЛЖНЫ ПАДАТЬ на текущей реализации (демонстрируя баг). +После фикса они должны проходить. +""" +import time + +import pandas as pd +import pytest +from sqlalchemy import Column, Integer, String + +from datapipe.compute import ComputeInput +from datapipe.datatable import DataStore +from datapipe.step.batch_transform import BatchTransformStep, DatatableBatchTransformStep +from datapipe.store.database import DBConn, TableStoreDB + + +@pytest.mark.xfail(reason="Reproduces production bug: offset skips records after intermediate transformation") +def test_offset_skips_records_with_intermediate_transformation(dbconn: DBConn): + """ + Воспроизводит основной баг из production. + + Сценарий: + 1. Создается запись в таблице (update_ts=T1, process_ts=T1) + 2. Промежуточная трансформация (DatatableBatchTransform) ЧИТАЕТ запись + - Обновляет process_ts=T2 + - НЕ обновляет update_ts (остается T1) + 3. Создается новая запись (update_ts=T3) + 4. Финальная трансформация с offset обрабатывает батч + - Батч содержит: старую запись (update_ts=T1) и новую (update_ts=T3) + - После обработки: offset = MAX(T1, T3) = T3 + 5. Следующий запуск: WHERE update_ts > T3 + - Старые записи с update_ts=T1 пропускаются! + + Реальный пример из production: + - 16:21 - Пост создается (update_ts=16:21, process_ts=16:21) + - 20:04 - hashtag_statistics_aggregation читает (update_ts=16:21, process_ts=20:04) + - 20:29 - Новый пост (update_ts=20:29) + - copy_to_online устанавливает offset=20:29 + - Следующий запуск пропускает все с update_ts < 20:29 + """ + ds = DataStore(dbconn, create_meta_table=True) + + # ========== Создаем таблицу данных (как post_hashtag_lower) ========== + data_store = TableStoreDB( + dbconn, + "data_table", + [ + Column("id", String, primary_key=True), + Column("hashtag", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + data_dt = ds.create_table("data_table", data_store) + + # ========== Создаем таблицу агрегации (как hashtag_lower) ========== + stats_store = TableStoreDB( + dbconn, + "stats_table", + [ + Column("hashtag", String, primary_key=True), + Column("count", Integer), + ], + create_table=True, + ) + stats_dt = ds.create_table("stats_table", stats_store) + + # ========== Создаем таблицу назначения (как post_hashtag_lower_online) ========== + target_store = TableStoreDB( + dbconn, + "target_table", + [ + Column("id", String, primary_key=True), + Column("hashtag", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + target_dt = ds.create_table("target_table", target_store) + + # ========== Промежуточная трансформация (читает, но не изменяет данные) ========== + def stats_aggregation(ds, idx, input_dts, run_config=None, kwargs=None): + """Агрегирует статистику - аналог hashtag_statistics_aggregation""" + df = input_dts[0].get_data(idx) + result = df.groupby("hashtag").size().reset_index(name="count") + return [(result, None)] + + intermediate_step = DatatableBatchTransformStep( + ds=ds, + name="stats_aggregation", + func=stats_aggregation, + input_dts=[ComputeInput(dt=data_dt, join_type="full")], + output_dts=[stats_dt], + transform_keys=["hashtag"], + use_offset_optimization=True, + ) + + # ========== Финальная трансформация (копирует в target) ========== + def copy_transform(df): + """Копирует данные - аналог copy_to_online""" + return df[["id", "hashtag", "value"]] + + final_step = BatchTransformStep( + ds=ds, + name="copy_to_target", + func=copy_transform, + input_dts=[ComputeInput(dt=data_dt, join_type="full")], + output_dts=[target_dt], + transform_keys=["id", "hashtag"], + use_offset_optimization=True, + ) + + # ========== Шаг 1: Создаем первую запись ========== + t1 = time.time() + data_dt.store_chunk( + pd.DataFrame({ + "id": ["post1"], + "hashtag": ["test"], + "value": [100], + }), + now=t1 + ) + + # Проверяем метаданные + meta1 = data_dt.meta_table.get_metadata() + assert len(meta1) == 1 + initial_update_ts = meta1.iloc[0]["update_ts"] + initial_process_ts = meta1.iloc[0]["process_ts"] + assert initial_update_ts == initial_process_ts # Изначально равны + + # ========== Шаг 2: Промежуточная трансформация ЧИТАЕТ данные ========== + time.sleep(0.01) + intermediate_step.run_full(ds) + + # Проверяем что process_ts обновился, но update_ts НЕТ + meta2 = data_dt.meta_table.get_metadata() + after_read_update_ts = meta2.iloc[0]["update_ts"] + after_read_process_ts = meta2.iloc[0]["process_ts"] + + # КРИТИЧНАЯ ПРОВЕРКА: update_ts НЕ изменился, process_ts изменился + assert after_read_update_ts == initial_update_ts, "update_ts не должен измениться при чтении" + assert after_read_process_ts > initial_process_ts, "process_ts должен обновиться после чтения" + + # ========== Шаг 3: Создаем новую запись (с более поздним update_ts) ========== + time.sleep(0.01) + t3 = time.time() + data_dt.store_chunk( + pd.DataFrame({ + "id": ["post2"], + "hashtag": ["demo"], + "value": [200], + }), + now=t3 + ) + + # ========== Шаг 4: Финальная трансформация обрабатывает ОБЕ записи ========== + final_step.run_full(ds) + + # Проверяем что обе записи скопировались + target_data_1 = target_dt.get_data() + assert len(target_data_1) == 2, "Обе записи должны быть скопированы" + assert set(target_data_1["id"].tolist()) == {"post1", "post2"} + + # ========== Шаг 5: Добавляем еще одну запись ========== + time.sleep(0.01) + t5 = time.time() + data_dt.store_chunk( + pd.DataFrame({ + "id": ["post3"], + "hashtag": ["new"], + "value": [300], + }), + now=t5 + ) + + # ========== Шаг 6: Финальная трансформация запускается снова ========== + # Из-за offset optimization должна обработать ТОЛЬКО post3 + # НО: если в батч попадет post1 (у которой старый update_ts), + # offset обновится на MAX(update_ts), и post1 будет потеряна при следующем запуске + + changed_count = final_step.get_changed_idx_count(ds) + # Ожидаем что видна только 1 новая запись (post3) + assert changed_count == 1, f"Должна быть видна только новая запись, получено: {changed_count}" + + final_step.run_full(ds) + + # Проверяем что все 3 записи в target + target_data_2 = target_dt.get_data() + assert len(target_data_2) == 3, f"Все 3 записи должны быть в target, получено: {len(target_data_2)}" + assert set(target_data_2["id"].tolist()) == {"post1", "post2", "post3"} + + # ========== Шаг 7: Симулируем что промежуточная трансформация снова читает старые данные ========== + # Это может произойти если она запускается без offset или перезапускается + time.sleep(0.01) + + # Обновляем process_ts для post1 снова (симулируем повторное чтение) + meta_post1 = data_dt.meta_table.get_metadata(pd.DataFrame({"id": ["post1"], "hashtag": ["test"]})) + current_update_ts_post1 = meta_post1.iloc[0]["update_ts"] + + # После повторного чтения process_ts должен обновиться + intermediate_step.run_full(ds) + + meta_post1_after = data_dt.meta_table.get_metadata(pd.DataFrame({"id": ["post1"], "hashtag": ["test"]})) + new_process_ts_post1 = meta_post1_after.iloc[0]["process_ts"] + new_update_ts_post1 = meta_post1_after.iloc[0]["update_ts"] + + # update_ts по-прежнему старый! + assert new_update_ts_post1 == current_update_ts_post1 + + # ========== Шаг 8: Добавляем еще записи ========== + time.sleep(0.01) + t8 = time.time() + data_dt.store_chunk( + pd.DataFrame({ + "id": ["post4"], + "hashtag": ["another"], + "value": [400], + }), + now=t8 + ) + + # ========== Шаг 9: Финальная трансформация ========== + # ПРОБЛЕМА: Если в changed records попадут post1 (update_ts старый) и post4 (update_ts новый) + # То offset = MAX(старый, новый) = новый + # И при следующем запуске post1 будет пропущена + + final_step.run_full(ds) + + # Финальная проверка: ВСЕ записи должны быть в target + final_target_data = target_dt.get_data() + assert len(final_target_data) == 4, ( + f"БАГ ВОСПРОИЗВЕДЕН: Ожидалось 4 записи в target, получено {len(final_target_data)}. " + f"Пропущенные записи с старым update_ts!" + ) + assert set(final_target_data["id"].tolist()) == {"post1", "post2", "post3", "post4"} + + +@pytest.mark.xfail(reason="Reproduces production bug: ORDER BY non-temporal causes data loss") +def test_offset_with_non_temporal_ordering(dbconn: DBConn): + """ + Воспроизводит Сценарий 2 из production: ORDER BY (id, hashtag) вместо update_ts. + + Сценарий: + 1. В таблице есть записи с разными update_ts и id + 2. Запрос использует ORDER BY id, hashtag LIMIT N + 3. Обработаны записи в порядке id (не по времени!) + 4. offset = MAX(update_ts) из обработанных записей + 5. Добавляется новая запись с id в середине диапазона + 6. Следующий запрос: WHERE update_ts > offset + - Запись отфильтрована по update_ts + - Но она также "пропущена" в сортировке по id + + Реальный пример из production: + Данные: [ + {id: "a1", hashtag: "test", update_ts: 19:00}, + {id: "b2", hashtag: "demo", update_ts: 18:00}, + {id: "c3", hashtag: "foo", update_ts: 17:00}, + {id: "z9", hashtag: "bar", update_ts: 20:00} + ] + Запрос 1: ORDER BY id, hashtag LIMIT 3 + → a1(19:00), b2(18:00), c3(17:00) + → offset = 19:00 + Новая запись: {id: "d4", hashtag: "new", update_ts: 18:30} + Запрос 2: WHERE update_ts > 19:00 + → d4 пропущена (18:30 < 19:00)! + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "input_ordering", + [ + Column("id", String, primary_key=True), + Column("hashtag", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("input_ordering", input_store) + + output_store = TableStoreDB( + dbconn, + "output_ordering", + [ + Column("id", String, primary_key=True), + Column("hashtag", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("output_ordering", output_store) + + def copy_func(df): + return df[["id", "hashtag", "value"]] + + step = BatchTransformStep( + ds=ds, + name="copy_ordering", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id", "hashtag"], + use_offset_optimization=True, + chunk_size=3, # Ограничиваем размер батча чтобы обработать только первые 3 + ) + + # ========== Шаг 1: Создаем записи с разными timestamp в несортированном порядке ========== + base_time = time.time() + + # a1 - update_ts будет средним + input_dt.store_chunk( + pd.DataFrame({"id": ["a1"], "hashtag": ["test"], "value": [1]}), + now=base_time + 2 # T=2 (средний) + ) + + time.sleep(0.01) + + # b2 - update_ts будет ранним + input_dt.store_chunk( + pd.DataFrame({"id": ["b2"], "hashtag": ["demo"], "value": [2]}), + now=base_time + 1 # T=1 (ранний) + ) + + time.sleep(0.01) + + # c3 - update_ts будет самым ранним + input_dt.store_chunk( + pd.DataFrame({"id": ["c3"], "hashtag": ["foo"], "value": [3]}), + now=base_time + 0 # T=0 (самый ранний) + ) + + time.sleep(0.01) + + # z9 - update_ts будет самым поздним + input_dt.store_chunk( + pd.DataFrame({"id": ["z9"], "hashtag": ["bar"], "value": [4]}), + now=base_time + 3 # T=3 (самый поздний) + ) + + # ========== Шаг 2: Первый запуск - обрабатываем только первые 3 (по ORDER BY id) ========== + # Из-за ORDER BY id, hashtag получим: a1, b2, c3 + # update_ts этих записей: T=2, T=1, T=0 + # offset = MAX(2, 1, 0) = 2 + + step.run_full(ds) + + # Проверяем что обработались первые 3 записи (по id) + output_data_1 = output_dt.get_data() + output_ids_1 = sorted(output_data_1["id"].tolist()) + assert len(output_data_1) == 3 + assert output_ids_1 == ["a1", "b2", "c3"], f"Ожидались a1,b2,c3, получено: {output_ids_1}" + + # Проверяем offset + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_value = offsets["input_ordering"] + + # offset должен быть максимальным из обработанных (T=2) + assert offset_value >= base_time + 2 + + # ========== Шаг 3: Добавляем запись d4 с промежуточным timestamp ========== + time.sleep(0.01) + # update_ts между T=1 и T=2, но меньше offset=T=2 + input_dt.store_chunk( + pd.DataFrame({"id": ["d4"], "hashtag": ["new"], "value": [5]}), + now=base_time + 1.5 # T=1.5 (между b2 и a1) + ) + + # ========== Шаг 4: Проверяем видимость новых записей ========== + # Должны быть видны: z9 (T=3 > offset=2) и d4 (T=1.5 < offset=2 - НЕ ВИДНА!) + changed_count = step.get_changed_idx_count(ds) + + # На самом деле будет видна только z9, так как d4 отфильтрована по offset + # Это и есть баг! + + # ========== Шаг 5: Второй запуск ========== + step.run_full(ds) + + # ========== КРИТИЧНАЯ ПРОВЕРКА: Все записи должны быть в output ========== + final_output = output_dt.get_data() + final_ids = sorted(final_output["id"].tolist()) + + # Эта проверка должна падать + assert len(final_output) == 5, ( + f"БАГ ВОСПРОИЗВЕДЕН: Ожидалось 5 записей, получено {len(final_output)}. " + f"Запись d4 с промежуточным update_ts пропущена!" + ) + assert final_ids == ["a1", "b2", "c3", "d4", "z9"], ( + f"Ожидались все записи включая d4, получено: {final_ids}" + ) + + +@pytest.mark.xfail(reason="Reproduces production bug: process_ts vs update_ts divergence") +def test_process_ts_vs_update_ts_divergence(dbconn: DBConn): + """ + Проверяет семантику update_ts vs process_ts. + + Проблема: + - update_ts обновляется только когда ДАННЫЕ изменяются + - process_ts обновляется когда запись ОБРАБАТЫВАЕТСЯ (читается) + - Offset фильтрация использует update_ts + - Если трансформация только читает (не изменяет), process_ts != update_ts + - Следующая трансформация с offset может пропустить "обработанные но не измененные" записи + + Этот тест проверяет базовый инвариант: + - После создания: update_ts == process_ts ✓ + - После чтения: update_ts < process_ts (расхождение!) + - Offset использует update_ts → данные теряются + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Создаем таблицу + table_store = TableStoreDB( + dbconn, + "timestamps_table", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + table_dt = ds.create_table("timestamps_table", table_store) + + # Создаем выходную таблицу для чтения (не изменяет данные) + read_output_store = TableStoreDB( + dbconn, + "read_output", + [Column("id", String, primary_key=True), Column("count", Integer)], + create_table=True, + ) + read_output_dt = ds.create_table("read_output", read_output_store) + + # Трансформация которая ЧИТАЕТ но НЕ ИЗМЕНЯЕТ входные данные + def read_only_transform(ds, idx, input_dts, run_config=None, kwargs=None): + """Просто считает записи - не изменяет источник""" + df = input_dts[0].get_data(idx) + result = pd.DataFrame({"id": ["summary"], "count": [len(df)]}) + return [(result, None)] + + read_step = DatatableBatchTransformStep( + ds=ds, + name="read_only", + func=read_only_transform, + input_dts=[ComputeInput(dt=table_dt, join_type="full")], + output_dts=[read_output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # Создаем выходную таблицу для копирования + copy_output_store = TableStoreDB( + dbconn, + "copy_output", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + copy_output_dt = ds.create_table("copy_output", copy_output_store) + + # Трансформация которая копирует данные + def copy_transform(df): + return df[["id", "value"]] + + copy_step = BatchTransformStep( + ds=ds, + name="copy_step", + func=copy_transform, + input_dts=[ComputeInput(dt=table_dt, join_type="full")], + output_dts=[copy_output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # ========== Шаг 1: Создаем записи ========== + t1 = time.time() + table_dt.store_chunk( + pd.DataFrame({"id": ["rec1", "rec2"], "value": [10, 20]}), + now=t1 + ) + + # Проверяем что после создания update_ts == process_ts + meta_after_create = table_dt.meta_table.get_metadata() + for idx, row in meta_after_create.iterrows(): + assert row["update_ts"] == row["process_ts"], ( + f"После создания update_ts должен равняться process_ts" + ) + create_update_ts = row["update_ts"] + create_process_ts = row["process_ts"] + + # ========== Шаг 2: Трансформация ЧИТАЕТ данные ========== + time.sleep(0.01) + read_step.run_full(ds) + + # Проверяем что process_ts обновился, но update_ts НЕТ + meta_after_read = table_dt.meta_table.get_metadata() + for idx, row in meta_after_read.iterrows(): + assert row["update_ts"] == create_update_ts, ( + f"После чтения update_ts НЕ должен измениться" + ) + assert row["process_ts"] > create_process_ts, ( + f"После чтения process_ts должен обновиться" + ) + # КРИТИЧНО: Теперь update_ts < process_ts (расхождение!) + assert row["update_ts"] < row["process_ts"], ( + "РАСХОЖДЕНИЕ: update_ts < process_ts после чтения данных" + ) + + # ========== Шаг 3: Добавляем новую запись с более поздним timestamp ========== + time.sleep(0.01) + t3 = time.time() + table_dt.store_chunk( + pd.DataFrame({"id": ["rec3"], "value": [30]}), + now=t3 + ) + + # ========== Шаг 4: Копируем все записи ========== + copy_step.run_full(ds) + + # Проверяем что все 3 записи скопировались + copy_output_1 = copy_output_dt.get_data() + assert len(copy_output_1) == 3 + assert set(copy_output_1["id"].tolist()) == {"rec1", "rec2", "rec3"} + + # Проверяем offset - будет MAX(update_ts) = t3 + offsets = ds.offset_table.get_offsets_for_transformation(copy_step.get_name()) + offset_after_copy = offsets["timestamps_table"] + assert offset_after_copy >= t3 + + # ========== Шаг 5: Читаем старые записи снова ========== + time.sleep(0.01) + read_step.run_full(ds) + + # process_ts для rec1, rec2 снова обновился + meta_after_read2 = table_dt.meta_table.get_metadata() + rec1_meta = meta_after_read2[meta_after_read2["id"] == "rec1"].iloc[0] + + # update_ts все еще старый (t1), но process_ts свежий + assert rec1_meta["update_ts"] == create_update_ts + assert rec1_meta["process_ts"] > create_process_ts + + # ========== Шаг 6: Добавляем еще одну запись ========== + time.sleep(0.01) + t6 = time.time() + table_dt.store_chunk( + pd.DataFrame({"id": ["rec4"], "value": [40]}), + now=t6 + ) + + # ========== Шаг 7: Копируем снова ========== + # ПРОБЛЕМА: Если offset основан на update_ts, + # а rec1, rec2 имеют старый update_ts (t1 < offset), + # они могут быть пропущены при определенных условиях + + copy_step.run_full(ds) + + # ========== КРИТИЧНАЯ ПРОВЕРКА: Все 4 записи должны быть скопированы ========== + final_copy_output = copy_output_dt.get_data() + assert len(final_copy_output) == 4, ( + f"БАГ: Ожидалось 4 записи, получено {len(final_copy_output)}. " + f"Записи с update_ts < process_ts могут быть пропущены!" + ) + assert set(final_copy_output["id"].tolist()) == {"rec1", "rec2", "rec3", "rec4"} + + +@pytest.mark.xfail(reason="Reproduces production bug: full chain like in production") +def test_copy_to_online_with_stats_aggregation_chain(dbconn: DBConn): + """ + Интеграционный тест: полная цепочка как в production. + + Цепочка: + 1. post_hashtag_scraping_lower создает записи в post_hashtag_lower + 2. hashtag_statistics_aggregation ЧИТАЕТ post_hashtag_lower (агрегирует статистику) + - Обновляет process_ts для прочитанных записей + - НЕ обновляет update_ts (данные не изменились) + 3. copy_to_online копирует из post_hashtag_lower в post_hashtag_lower_online + - Использует offset на основе update_ts + - Пропускает записи с "устаревшим" update_ts + + Реальный сценарий из production: + - 16:21 - Пост создан, хештеги извлечены + - 20:04 - hashtag_statistics_aggregation прочитал (process_ts=20:04, update_ts=16:21) + - 20:29 - Новый пост создан (update_ts=20:29) + - copy_to_online обработал обе записи, offset=20:29 + - Следующий запуск copy_to_online пропустил все записи с update_ts<20:29 + - Результат: 60% данных потеряно + """ + ds = DataStore(dbconn, create_meta_table=True) + + # ========== Таблица: post_hashtag_lower (offline) ========== + post_hashtag_store = TableStoreDB( + dbconn, + "post_hashtag_lower", + [ + Column("id", String, primary_key=True), + Column("hashtag", String, primary_key=True), + Column("created_at", Integer), + ], + create_table=True, + ) + post_hashtag_dt = ds.create_table("post_hashtag_lower", post_hashtag_store) + + # ========== Таблица: hashtag_lower (статистика) ========== + hashtag_stats_store = TableStoreDB( + dbconn, + "hashtag_lower", + [ + Column("hashtag", String, primary_key=True), + Column("post_count", Integer), + Column("last_seen", Integer), + ], + create_table=True, + ) + hashtag_stats_dt = ds.create_table("hashtag_lower", hashtag_stats_store) + + # ========== Таблица: post_hashtag_lower_online ========== + post_hashtag_online_store = TableStoreDB( + dbconn, + "post_hashtag_lower_online", + [ + Column("id", String, primary_key=True), + Column("hashtag", String, primary_key=True), + Column("created_at", Integer), + ], + create_table=True, + ) + post_hashtag_online_dt = ds.create_table("post_hashtag_lower_online", post_hashtag_online_store) + + # ========== Трансформация 1: hashtag_statistics_aggregation ========== + def aggregate_hashtag_stats(ds, idx, input_dts, run_config=None, kwargs=None): + """Агрегирует статистику по хештегам""" + df = input_dts[0].get_data(idx) + stats = df.groupby("hashtag").agg({ + "id": "count", + "created_at": "max" + }).reset_index() + stats.columns = ["hashtag", "post_count", "last_seen"] + return [(stats, None)] + + stats_step = DatatableBatchTransformStep( + ds=ds, + name="hashtag_statistics_aggregation", + func=aggregate_hashtag_stats, + input_dts=[ComputeInput(dt=post_hashtag_dt, join_type="full")], + output_dts=[hashtag_stats_dt], + transform_keys=["hashtag"], + use_offset_optimization=True, + ) + + # ========== Трансформация 2: copy_to_online ========== + def copy_to_online(df): + """Копирует данные в online БД""" + return df[["id", "hashtag", "created_at"]] + + copy_step = BatchTransformStep( + ds=ds, + name="copy_to_online", + func=copy_to_online, + input_dts=[ComputeInput(dt=post_hashtag_dt, join_type="full")], + output_dts=[post_hashtag_online_dt], + transform_keys=["id", "hashtag"], + use_offset_optimization=True, + ) + + # ========== Симуляция production сценария ========== + + # Время 16:21 - Пост b927ca71 создан с 5 хештегами + t_16_21 = time.time() + post_hashtag_dt.store_chunk( + pd.DataFrame({ + "id": ["b927ca71"] * 5, + "hashtag": ["322", "anime", "looky", "test77", "ошош"], + "created_at": [int(t_16_21)] * 5, + }), + now=t_16_21 + ) + + # Проверяем метаданные после создания + meta_after_create = post_hashtag_dt.meta_table.get_metadata() + assert len(meta_after_create) == 5 + initial_update_ts = meta_after_create.iloc[0]["update_ts"] + + # Время 20:04 - hashtag_statistics_aggregation обрабатывает + time.sleep(0.01) + stats_step.run_full(ds) + + # Проверяем что статистика создана + stats_data = hashtag_stats_dt.get_data() + assert len(stats_data) == 5 + + # КРИТИЧНО: process_ts обновился, update_ts НЕТ + meta_after_stats = post_hashtag_dt.meta_table.get_metadata() + for idx, row in meta_after_stats.iterrows(): + assert row["update_ts"] == initial_update_ts, "update_ts не должен измениться" + assert row["process_ts"] > initial_update_ts, "process_ts должен обновиться" + + # Время 20:29 - Новый пост e26f9c4b создан + time.sleep(0.01) + t_20_29 = time.time() + post_hashtag_dt.store_chunk( + pd.DataFrame({ + "id": ["e26f9c4b"], + "hashtag": ["demo"], + "created_at": [int(t_20_29)], + }), + now=t_20_29 + ) + + # Время 20:30 - copy_to_online запускается ВПЕРВЫЕ + time.sleep(0.01) + copy_step.run_full(ds) + + # Проверяем что ВСЕ 6 хештегов скопированы + online_data_1 = post_hashtag_online_dt.get_data() + assert len(online_data_1) == 6, f"Ожидалось 6 записей, получено {len(online_data_1)}" + + # Проверяем offset + offsets = ds.offset_table.get_offsets_for_transformation(copy_step.get_name()) + offset_after_first_copy = offsets["post_hashtag_lower"] + + # Offset будет MAX(update_ts) из обработанных записей + # Это будет t_20_29 (самый новый update_ts) + assert offset_after_first_copy >= t_20_29 + + # Время 09:00 (следующий день) - Еще посты создаются + time.sleep(0.01) + t_09_00 = time.time() + post_hashtag_dt.store_chunk( + pd.DataFrame({ + "id": ["f79ec772", "f79ec772"], + "hashtag": ["новый", "хештег"], + "created_at": [int(t_09_00)] * 2, + }), + now=t_09_00 + ) + + # hashtag_statistics_aggregation снова обрабатывает (включая старые записи) + time.sleep(0.01) + stats_step.run_full(ds) + + # Проверяем что process_ts для старых записей снова обновился + meta_old_posts = post_hashtag_dt.meta_table.get_metadata( + pd.DataFrame({ + "id": ["b927ca71"] * 5, + "hashtag": ["322", "anime", "looky", "test77", "ошош"], + }) + ) + for idx, row in meta_old_posts.iterrows(): + # update_ts все еще старый! + assert row["update_ts"] == initial_update_ts + # process_ts свежий + assert row["process_ts"] > offset_after_first_copy + + # copy_to_online запускается снова + time.sleep(0.01) + copy_step.run_full(ds) + + # ========== КРИТИЧНАЯ ПРОВЕРКА: Все 8 записей должны быть в online ========== + final_online_data = post_hashtag_online_dt.get_data() + + # БАГ: Старые записи (b927ca71) с update_ts < offset будут пропущены + # если они попадут в батч вместе с новыми записями + assert len(final_online_data) == 8, ( + f"БАГ ВОСПРОИЗВЕДЕН: Ожидалось 8 записей в online, получено {len(final_online_data)}. " + f"Это симулирует потерю 60% данных из production!" + ) + + # Проверяем что все посты присутствуют + online_post_ids = set(final_online_data["id"].unique()) + expected_ids = {"b927ca71", "e26f9c4b", "f79ec772"} + assert online_post_ids == expected_ids, ( + f"Ожидались посты {expected_ids}, получено {online_post_ids}" + ) diff --git a/tests/performance/__init__.py b/tests/performance/__init__.py new file mode 100644 index 00000000..a3bb79a5 --- /dev/null +++ b/tests/performance/__init__.py @@ -0,0 +1,5 @@ +""" +Performance tests for offset optimization. + +Run with: pytest tests/performance/ -v +""" diff --git a/tests/performance/conftest.py b/tests/performance/conftest.py new file mode 100644 index 00000000..882a62b4 --- /dev/null +++ b/tests/performance/conftest.py @@ -0,0 +1,167 @@ +""" +Shared fixtures for performance tests. +""" +import os +import time + +import pandas as pd +import pytest +from sqlalchemy import Column, Integer, String, create_engine, text + +from datapipe.compute import ComputeInput +from datapipe.datatable import DataStore +from datapipe.step.batch_transform import BatchTransformStep +from datapipe.store.database import DBConn, TableStoreDB + + +@pytest.fixture +def dbconn(): + """Database connection fixture for performance tests.""" + if os.environ.get("TEST_DB_ENV") == "sqlite": + DBCONNSTR = "sqlite+pysqlite3:///:memory:" + DB_TEST_SCHEMA = 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" + DB_TEST_SCHEMA = "test" + + if DB_TEST_SCHEMA: + eng = create_engine(DBCONNSTR) + + try: + with eng.begin() as conn: + conn.execute(text(f"DROP SCHEMA {DB_TEST_SCHEMA} CASCADE")) + except Exception: + pass + + with eng.begin() as conn: + conn.execute(text(f"CREATE SCHEMA {DB_TEST_SCHEMA}")) + + dbconn = DBConn(DBCONNSTR, DB_TEST_SCHEMA) + yield dbconn + + +def fast_bulk_insert(dbconn: DBConn, table_name: str, data: pd.DataFrame): + """ + Быстрая вставка данных используя bulk insert (PostgreSQL). + + Это намного быстрее чем обычный store_chunk, т.к. минует метаданные. + """ + # Используем to_sql с method='multi' для быстрой вставки + with dbconn.con.begin() as con: + data.to_sql( + table_name, + con=con, + schema=dbconn.schema, + if_exists='append', + index=False, + method='multi', + chunksize=10000, + ) + + +@pytest.fixture +def fast_data_loader(dbconn: DBConn): + """ + Фикстура для быстрой загрузки больших объемов тестовых данных. + """ + def prepare_large_dataset( + ds: DataStore, + table_name: str, + num_records: int, + primary_key: str = "id" + ) -> "DataTable": # noqa: F821 + """ + Быстрая подготовка большого объема тестовых данных. + + Использует прямую вставку в data таблицу и метатаблицу для скорости. + """ + # Создать таблицу + store = TableStoreDB( + dbconn, + table_name, + [ + Column(primary_key, String, primary_key=True), + Column("value", Integer), + Column("category", String), + ], + create_table=True, + ) + dt = ds.create_table(table_name, store) + + # Генерировать данные партиями для эффективности памяти + batch_size = 50000 + now = time.time() + + print(f"\nPreparing {num_records:,} records for {table_name}...") + start_time = time.time() + + for i in range(0, num_records, batch_size): + chunk_size = min(batch_size, num_records - i) + + # Генерация данных + data = pd.DataFrame({ + primary_key: [f"id_{j:010d}" for j in range(i, i + chunk_size)], + "value": range(i, i + chunk_size), + "category": [f"cat_{j % 100}" for j in range(i, i + chunk_size)], + }) + + # Быстрая вставка в data таблицу + fast_bulk_insert(dbconn, table_name, data) + + # Вставка метаданных (необходимо для работы offset) + meta_data = data[[primary_key]].copy() + meta_data['hash'] = 0 # Упрощенный hash + meta_data['create_ts'] = now + meta_data['update_ts'] = now + meta_data['process_ts'] = None + meta_data['delete_ts'] = None + + fast_bulk_insert(dbconn, f"{table_name}_meta", meta_data) + + if (i + chunk_size) % 100000 == 0: + elapsed = time.time() - start_time + rate = (i + chunk_size) / elapsed + print(f" Inserted {i + chunk_size:,} records ({rate:,.0f} rec/sec)") + + total_time = time.time() - start_time + print(f"Data preparation completed in {total_time:.2f}s ({num_records/total_time:,.0f} rec/sec)") + + return dt + + return prepare_large_dataset + + +@pytest.fixture +def perf_pipeline_factory(dbconn: DBConn): + """ + Фабрика для создания пайплайнов в нагрузочных тестах. + """ + def create_pipeline( + ds: DataStore, + name: str, + source_dt, + target_dt, + use_offset: bool + ) -> BatchTransformStep: + """ + Создать пайплайн для тестирования производительности. + """ + def copy_transform(df): + # Простая трансформация - копирование с небольшим изменением + result = df.copy() + result['value'] = result['value'] * 2 + return result + + return BatchTransformStep( + ds=ds, + name=name, + func=copy_transform, + input_dts=[ComputeInput(dt=source_dt, join_type="full")], + output_dts=[target_dt], + transform_keys=["id"], + use_offset_optimization=use_offset, + ) + + return create_pipeline diff --git a/tests/performance/test_offset_performance.py b/tests/performance/test_offset_performance.py new file mode 100644 index 00000000..3dddee29 --- /dev/null +++ b/tests/performance/test_offset_performance.py @@ -0,0 +1,568 @@ +""" +Нагрузочные тесты для offset-оптимизации. + +Запуск: pytest tests/performance/ -v +Обычные тесты: pytest tests/ -v (исключает tests/performance автоматически) +""" +import signal +import time +from contextlib import contextmanager +from typing import List, Optional, Tuple + +import pandas as pd +from sqlalchemy import Column, Integer, String + +from datapipe.datatable import DataStore +from datapipe.store.database import DBConn, TableStoreDB + +from .conftest import fast_bulk_insert + + +class TimeoutError(Exception): + """Raised when operation exceeds timeout""" + pass + + +@contextmanager +def timeout(seconds: int): + """ + Context manager для ограничения времени выполнения. + + Если операция не завершилась за указанное время, поднимается TimeoutError. + """ + def timeout_handler(signum, frame): + raise TimeoutError(f"Operation timed out after {seconds} seconds") + + # Установить обработчик сигнала + old_handler = signal.signal(signal.SIGALRM, timeout_handler) + signal.alarm(seconds) + + try: + yield + finally: + # Отменить таймер и восстановить старый обработчик + signal.alarm(0) + signal.signal(signal.SIGALRM, old_handler) + + +def test_performance_small_dataset(fast_data_loader, perf_pipeline_factory, dbconn: DBConn): + """ + Тест производительности на маленьком датасете (10K записей). + + Проверяет что оба метода работают корректно и измеряет время. + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Подготовка данных + num_records = 10_000 + source_dt = fast_data_loader(ds, "perf_small_source", num_records) + + # Создать target таблицы + target_v1_store = TableStoreDB( + dbconn, + "perf_small_target_v1", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + Column("category", String), + ], + create_table=True, + ) + target_v1 = ds.create_table("perf_small_target_v1", target_v1_store) + + target_v2_store = TableStoreDB( + dbconn, + "perf_small_target_v2", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + Column("category", String), + ], + create_table=True, + ) + target_v2 = ds.create_table("perf_small_target_v2", target_v2_store) + + # Тест v1 (старый метод) + step_v1 = perf_pipeline_factory( + ds, "perf_small_v1", source_dt, target_v1, use_offset=False + ) + + print("\n=== Testing v1 (FULL OUTER JOIN) ===") + start = time.time() + step_v1.run_full(ds) + v1_time = time.time() - start + print(f"v1 time: {v1_time:.3f}s") + + # Тест v2 (offset метод) + step_v2 = perf_pipeline_factory( + ds, "perf_small_v2", source_dt, target_v2, use_offset=True + ) + + print("\n=== Testing v2 (offset-based) ===") + start = time.time() + step_v2.run_full(ds) + v2_time = time.time() - start + print(f"v2 time: {v2_time:.3f}s") + + # Проверка корректности + v1_data = target_v1.get_data().sort_values("id").reset_index(drop=True) + v2_data = target_v2.get_data().sort_values("id").reset_index(drop=True) + + assert len(v1_data) == num_records + assert len(v2_data) == num_records + pd.testing.assert_frame_equal(v1_data, v2_data) + + print("\n=== Results ===") + print(f"Records processed: {num_records:,}") + print(f"v1 (FULL OUTER JOIN): {v1_time:.3f}s") + print(f"v2 (offset-based): {v2_time:.3f}s") + if v1_time > 0: + print(f"Speedup: {v1_time/v2_time:.2f}x") + + +def test_performance_large_dataset_with_timeout( + fast_data_loader, perf_pipeline_factory, dbconn: DBConn +): + """ + Тест производительности на большом датасете (100K записей) с ограничением времени. + + Если метод не завершается за MAX_TIME секунд, считаем что он превысил лимит. + """ + MAX_TIME = 60 # 60 секунд максимум + + ds = DataStore(dbconn, create_meta_table=True) + + # Подготовка данных + num_records = 100_000 + source_dt = fast_data_loader(ds, "perf_large_source", num_records) + + # Создать target таблицы + target_v1_store = TableStoreDB( + dbconn, + "perf_large_target_v1", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + Column("category", String), + ], + create_table=True, + ) + target_v1 = ds.create_table("perf_large_target_v1", target_v1_store) + + target_v2_store = TableStoreDB( + dbconn, + "perf_large_target_v2", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + Column("category", String), + ], + create_table=True, + ) + target_v2 = ds.create_table("perf_large_target_v2", target_v2_store) + + # Тест v1 с таймаутом + step_v1 = perf_pipeline_factory( + ds, "perf_large_v1", source_dt, target_v1, use_offset=False + ) + + print("\n=== Testing v1 (FULL OUTER JOIN) with timeout ===") + v1_time: Optional[float] = None + v1_timed_out = False + + try: + with timeout(MAX_TIME): + start = time.time() + step_v1.run_full(ds) + v1_time = time.time() - start + print(f"v1 time: {v1_time:.3f}s") + except TimeoutError: + v1_timed_out = True + v1_time = None + print(f"v1 timed out (>{MAX_TIME}s)") + + # Тест v2 (должен быть быстрее) + step_v2 = perf_pipeline_factory( + ds, "perf_large_v2", source_dt, target_v2, use_offset=True + ) + + print("\n=== Testing v2 (offset-based) ===") + start = time.time() + step_v2.run_full(ds) + v2_time = time.time() - start + print(f"v2 time: {v2_time:.3f}s") + + # Проверка что v2 не превысил таймаут + assert v2_time < MAX_TIME, f"v2 method should complete within {MAX_TIME}s, took {v2_time:.3f}s" + + # Проверка корректности результатов v2 + v2_data = target_v2.get_data() + assert len(v2_data) == num_records, f"Expected {num_records} records, got {len(v2_data)}" + + # Вывод результатов + print("\n=== Results ===") + print(f"Records processed: {num_records:,}") + if v1_timed_out: + print(f"v1 (FULL OUTER JOIN): >{MAX_TIME}s (timed out)") + print(f"v2 (offset-based): {v2_time:.3f}s") + print(f"Speedup: >{MAX_TIME/v2_time:.2f}x (minimum)") + else: + print(f"v1 (FULL OUTER JOIN): {v1_time:.3f}s") + print(f"v2 (offset-based): {v2_time:.3f}s") + print(f"Speedup: {v1_time/v2_time:.2f}x") + + +def test_performance_incremental_updates( + fast_data_loader, perf_pipeline_factory, dbconn: DBConn +): + """ + Тест производительности инкрементальных обновлений. + + Создаем большой датасет, обрабатываем его, затем добавляем небольшое количество + новых записей и измеряем время обработки только новых данных. + + Это основной use case для offset-оптимизации. + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Подготовка начальных данных + initial_records = 50_000 + new_records = 1_000 + + source_dt = fast_data_loader(ds, "perf_incr_source", initial_records) + + # Создать target таблицы + target_v1_store = TableStoreDB( + dbconn, + "perf_incr_target_v1", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + Column("category", String), + ], + create_table=True, + ) + target_v1 = ds.create_table("perf_incr_target_v1", target_v1_store) + + target_v2_store = TableStoreDB( + dbconn, + "perf_incr_target_v2", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + Column("category", String), + ], + create_table=True, + ) + target_v2 = ds.create_table("perf_incr_target_v2", target_v2_store) + + # Создать пайплайны + step_v1 = perf_pipeline_factory( + ds, "perf_incr_v1", source_dt, target_v1, use_offset=False + ) + step_v2 = perf_pipeline_factory( + ds, "perf_incr_v2", source_dt, target_v2, use_offset=True + ) + + # Первоначальная обработка (оба метода) + print("\n=== Initial processing ===") + step_v1.run_full(ds) + step_v2.run_full(ds) + print(f"Processed {initial_records:,} initial records") + + # Добавить новые записи + time.sleep(0.01) # Небольшая задержка для разделения timestamp + now = time.time() + + new_data = pd.DataFrame({ + "id": [f"new_id_{j:010d}" for j in range(new_records)], + "value": range(1000000, 1000000 + new_records), + "category": [f"new_cat_{j % 10}" for j in range(new_records)], + }) + + # Вставка новых данных + fast_bulk_insert(dbconn, "perf_incr_source", new_data) + + # Вставка метаданных для новых записей + meta_data = new_data[["id"]].copy() + meta_data['hash'] = 0 + meta_data['create_ts'] = now + meta_data['update_ts'] = now + meta_data['process_ts'] = None + meta_data['delete_ts'] = None + fast_bulk_insert(dbconn, "perf_incr_source_meta", meta_data) + + print(f"\nAdded {new_records:,} new records") + + # Инкрементальная обработка v1 (будет обрабатывать все данные заново) + print("\n=== v1 incremental (re-processes all data) ===") + start = time.time() + step_v1.run_full(ds) + v1_incr_time = time.time() - start + print(f"v1 incremental time: {v1_incr_time:.3f}s") + + # Инкрементальная обработка v2 (обработает только новые записи) + print("\n=== v2 incremental (processes only new data) ===") + + # Проверяем что видны только новые записи + changed_count = step_v2.get_changed_idx_count(ds) + print(f"v2 detected {changed_count} changed records") + assert changed_count == new_records, f"Expected {new_records} changes, got {changed_count}" + + start = time.time() + step_v2.run_full(ds) + v2_incr_time = time.time() - start + print(f"v2 incremental time: {v2_incr_time:.3f}s") + + # Проверка корректности + v1_data = target_v1.get_data().sort_values("id").reset_index(drop=True) + v2_data = target_v2.get_data().sort_values("id").reset_index(drop=True) + + total_records = initial_records + new_records + assert len(v1_data) == total_records + assert len(v2_data) == total_records + pd.testing.assert_frame_equal(v1_data, v2_data) + + # Вывод результатов + print("\n=== Incremental Update Results ===") + print(f"Total records: {total_records:,}") + print(f"New records: {new_records:,} ({new_records/total_records*100:.1f}%)") + print(f"v1 incremental: {v1_incr_time:.3f}s (re-processes all {total_records:,} records)") + print(f"v2 incremental: {v2_incr_time:.3f}s (processes only {new_records:,} new records)") + print(f"Speedup: {v1_incr_time/v2_incr_time:.2f}x") + print(f"v2 efficiency: {v2_incr_time/v1_incr_time*100:.1f}% of v1 time") + + +def test_performance_scalability_analysis( + fast_data_loader, perf_pipeline_factory, dbconn: DBConn +): + """ + Тест масштабируемости с экстраполяцией на большие датасеты. + + Тестирует 3 размера: 100K, 500K, 1M записей. + Использует initialize_offsets_from_transform_meta() для ускорения. + Экстраполирует результаты на 10M и 100M записей. + """ + MAX_TIME = 180 # 3 минуты максимум для v1 + + ds = DataStore(dbconn, create_meta_table=True) + + # Размеры датасетов для тестирования + test_sizes = [ + (100_000, True), # 100K - с v1 initial + (500_000, False), # 500K - без v1 initial + (1_000_000, False), # 1M - без v1 initial + ] + + new_records = 1_000 # Фиксированное количество новых записей для всех тестов + + results: List[Tuple[int, float, float]] = [] # (size, v1_time, v2_time) + + print(f"\n{'='*70}") + print(f"SCALABILITY ANALYSIS: Testing {len(test_sizes)} dataset sizes") + print(f"New records per test: {new_records:,} (constant)") + print(f"{'='*70}") + + for dataset_size, run_v1_initial in test_sizes: + print(f"\n{'='*70}") + print(f"Dataset size: {dataset_size:,} records") + print(f"{'='*70}") + + # Уникальные имена таблиц для каждого размера + size_suffix = f"{dataset_size//1000}k" + source_table = f"perf_scale_{size_suffix}_source" + target_v1_table = f"perf_scale_{size_suffix}_target_v1" + target_v2_table = f"perf_scale_{size_suffix}_target_v2" + + # Подготовка данных + source_dt = fast_data_loader(ds, source_table, dataset_size) + + # Создать target таблицы + target_v1_store = TableStoreDB( + dbconn, + target_v1_table, + [ + Column("id", String, primary_key=True), + Column("value", Integer), + Column("category", String), + ], + create_table=True, + ) + target_v1 = ds.create_table(target_v1_table, target_v1_store) + + target_v2_store = TableStoreDB( + dbconn, + target_v2_table, + [ + Column("id", String, primary_key=True), + Column("value", Integer), + Column("category", String), + ], + create_table=True, + ) + target_v2 = ds.create_table(target_v2_table, target_v2_store) + + # Создать пайплайны + step_v1 = perf_pipeline_factory( + ds, f"perf_scale_v1_{size_suffix}", source_dt, target_v1, use_offset=False + ) + step_v2 = perf_pipeline_factory( + ds, f"perf_scale_v2_{size_suffix}", source_dt, target_v2, use_offset=True + ) + + # ========== Первоначальная обработка ========== + if run_v1_initial: + print("\n--- Initial processing (both methods) ---") + print("v1 initial processing...") + start = time.time() + step_v1.run_full(ds) + v1_initial_time = time.time() - start + print(f"v1 initial: {v1_initial_time:.3f}s") + + print("v2 initial processing...") + start = time.time() + step_v2.run_full(ds) + v2_initial_time = time.time() - start + print(f"v2 initial: {v2_initial_time:.3f}s") + else: + # Быстрая инициализация через initialize_offsets + print("\n--- Fast initialization (v2 only, using init_offsets) ---") + print("v2 processing all data...") + start = time.time() + step_v2.run_full(ds) + v2_initial_time = time.time() - start + print(f"v2 initial: {v2_initial_time:.3f}s") + + # Инициализируем offset для v1 тоже (чтобы simulate что v1 уже работал) + print("Initializing v1 target (simulate previous processing)...") + step_v1.run_full(ds) + print("Done") + + # ========== Добавить новые записи ========== + time.sleep(0.01) + now = time.time() + + new_data = pd.DataFrame({ + "id": [f"new_{size_suffix}_{j:010d}" for j in range(new_records)], + "value": range(2000000, 2000000 + new_records), + "category": [f"new_cat_{j % 10}" for j in range(new_records)], + }) + + print(f"\n--- Adding {new_records:,} new records ({new_records/dataset_size*100:.3f}%) ---") + + # Вставка новых данных + fast_bulk_insert(dbconn, source_table, new_data) + + # Вставка метаданных для новых записей + meta_data = new_data[["id"]].copy() + meta_data['hash'] = 0 + meta_data['create_ts'] = now + meta_data['update_ts'] = now + meta_data['process_ts'] = None + meta_data['delete_ts'] = None + fast_bulk_insert(dbconn, f"{source_table}_meta", meta_data) + + # ========== Инкрементальная обработка v1 с таймаутом ========== + print("\n--- v1 incremental (FULL OUTER JOIN) ---") + v1_incr_time: Optional[float] = None + v1_timed_out = False + + try: + with timeout(MAX_TIME): + start = time.time() + step_v1.run_full(ds) + v1_incr_time = time.time() - start + print(f"v1 incremental: {v1_incr_time:.3f}s") + except TimeoutError: + v1_timed_out = True + v1_incr_time = MAX_TIME # Используем MAX_TIME для экстраполяции + print(f"v1 TIMED OUT (>{MAX_TIME}s)") + + # ========== Инкрементальная обработка v2 ========== + print("\n--- v2 incremental (offset-based) ---") + + # Проверяем что видны только новые записи + changed_count = step_v2.get_changed_idx_count(ds) + print(f"v2 detected {changed_count:,} changed records") + assert changed_count == new_records, f"Expected {new_records} changes, got {changed_count}" + + start = time.time() + step_v2.run_full(ds) + v2_incr_time = time.time() - start + print(f"v2 incremental: {v2_incr_time:.3f}s") + + # Проверка корректности + v1_data = target_v1.get_data() + v2_data = target_v2.get_data() + + total_records = dataset_size + new_records + assert len(v1_data) == total_records, f"v1: expected {total_records}, got {len(v1_data)}" + assert len(v2_data) == total_records, f"v2: expected {total_records}, got {len(v2_data)}" + + # Сохранить результаты + results.append((dataset_size, v1_incr_time, v2_incr_time)) + + # Вывод результатов для текущего размера + print(f"\n--- Results for {dataset_size:,} records ---") + if v1_timed_out: + print(f"v1: >{MAX_TIME:.1f}s (TIMED OUT)") + else: + print(f"v1: {v1_incr_time:.3f}s") + print(f"v2: {v2_incr_time:.3f}s") + if not v1_timed_out: + print(f"Speedup: {v1_incr_time/v2_incr_time:.1f}x") + + # ========== Анализ масштабируемости и экстраполяция ========== + print(f"\n{'='*70}") + print("SCALABILITY ANALYSIS & EXTRAPOLATION") + print(f"{'='*70}") + + print(f"\n{'Dataset Size':<15} {'v1 Time (s)':<15} {'v2 Time (s)':<15} {'Speedup':<10}") + print("-" * 70) + for size, v1_time, v2_time in results: + speedup = v1_time / v2_time if v2_time > 0 else float('inf') + timeout_marker = ">" if v1_time >= MAX_TIME else "" + print(f"{size:>13,} {timeout_marker}{v1_time:>14.3f} {v2_time:>14.3f} {speedup:>9.1f}x") + + # Экстраполяция на большие размеры + print("\n--- Extrapolation to larger datasets ---") + + # v1: линейная экстраполяция (O(N)) + # Находим коэффициент: time = k * size + valid_v1_results = [(s, t) for s, t, _ in results if t < MAX_TIME] + + if len(valid_v1_results) >= 2: + # Линейная регрессия для v1 + sizes = [s for s, _ in valid_v1_results] + times = [t for _, t in valid_v1_results] + k_v1 = sum(s * t for s, t in zip(sizes, times)) / sum(s * s for s in sizes) + + print(f"\nv1 (FULL OUTER JOIN) scaling: ~{k_v1*1e6:.2f} µs per record") + else: + # Если не хватает данных, используем последнее измерение + k_v1 = results[-1][1] / results[-1][0] + print(f"\nv1 (FULL OUTER JOIN) scaling: ~{k_v1*1e6:.2f} µs per record (estimated)") + + # v2: константная сложность (зависит только от new_records) + v2_avg = sum(t for _, _, t in results) / len(results) + print(f"v2 (offset-based) average: {v2_avg:.3f}s (constant, ~{v2_avg/new_records*1e6:.2f} µs per new record)") + + # Прогноз для больших датасетов + extrapolation_sizes = [10_000_000, 100_000_000] # 10M, 100M + + print(f"\n{'Dataset Size':<15} {'v1 Estimated':<20} {'v2 Estimated':<20} {'Est. Speedup':<15}") + print("-" * 70) + for size in extrapolation_sizes: + v1_est = k_v1 * size + v2_est = v2_avg # Константа + speedup_est = v1_est / v2_est + + print(f"{size:>13,} {v1_est:>17.1f}s {v2_est:>17.3f}s {speedup_est:>14.0f}x") + + print(f"\n{'='*70}") + print("CONCLUSION") + print(f"{'='*70}") + print("v1 complexity: O(N) - scales linearly with dataset size") + print(f"v2 complexity: O(M) - depends only on number of changes ({new_records:,})") + print("\nFor incremental updates on large datasets (10M+), v2 is 100-1000x faster!") + print(f"{'='*70}\n") diff --git a/tests/test_batch_transform_with_offset_optimization.py b/tests/test_batch_transform_with_offset_optimization.py new file mode 100644 index 00000000..fbe52a31 --- /dev/null +++ b/tests/test_batch_transform_with_offset_optimization.py @@ -0,0 +1,426 @@ +""" +Интеграционные тесты для проверки работы BatchTransformStep с offset оптимизацией (v2) +""" +import time + +import pandas as pd +from sqlalchemy import Column, Integer, String + +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_batch_transform_with_offset_basic(dbconn: DBConn): + """Базовый тест работы BatchTransformStep с offset оптимизацией""" + ds = DataStore(dbconn, create_meta_table=True) + + # Создаем входную таблицу + input_store = TableStoreDB( + dbconn, + "test_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("test_input", input_store) + + # Создаем выходную таблицу + output_store = TableStoreDB( + dbconn, + "test_output", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("test_output", output_store) + + # Функция трансформации + def transform_func(df): + return df.rename(columns={"value": "result"}) + + # Создаем step с включенной offset оптимизацией + step = BatchTransformStep( + ds=ds, + name="test_transform", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2", "3"], "value": [10, 20, 30]}), now=now + ) + + # Первый запуск - обработать все + step.run_full(ds) + + # Проверяем результат + output_data = output_dt.get_data() + assert len(output_data) == 3 + assert sorted(output_data["id"].tolist()) == ["1", "2", "3"] + assert sorted(output_data["result"].tolist()) == [10, 20, 30] + + # Добавляем новые данные + time.sleep(0.01) + now2 = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["4", "5"], "value": [40, 50]}), now=now2 + ) + + # Второй запуск - должны обработаться только новые записи + step.run_full(ds) + + # Проверяем, что все данные присутствуют + output_data = output_dt.get_data() + assert len(output_data) == 5 + assert sorted(output_data["id"].tolist()) == ["1", "2", "3", "4", "5"] + assert sorted(output_data["result"].tolist()) == [10, 20, 30, 40, 50] + + +def test_batch_transform_offset_filters_old_records(dbconn: DBConn): + """Тест что offset фильтрует старые записи при инкрементальной обработке""" + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "test_input2", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("test_input2", input_store) + + output_store = TableStoreDB( + dbconn, + "test_output2", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("test_output2", output_store) + + call_count = {"count": 0} + + def transform_func(df): + call_count["count"] += 1 + return df.rename(columns={"value": "result"}) + + step = BatchTransformStep( + ds=ds, + name="test_transform2", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # Добавляем старые данные + old_time = time.time() - 100 + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2"], "value": [10, 20]}), now=old_time + ) + + # Первый запуск + step.run_full(ds) + first_call_count = call_count["count"] + + # Добавляем новые данные + time.sleep(0.01) + new_time = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["3", "4"], "value": [30, 40]}), now=new_time + ) + + # Второй запуск - должны обработаться только новые + step.run_full(ds) + + # Проверяем, что функция вызывалась оба раза + assert call_count["count"] > first_call_count + + # Проверяем результат + output_data = output_dt.get_data() + assert len(output_data) == 4 + assert sorted(output_data["id"].tolist()) == ["1", "2", "3", "4"] + + +def test_batch_transform_offset_with_error_retry(dbconn: DBConn): + """Тест что ошибочные записи переобрабатываются""" + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "test_input3", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("test_input3", input_store) + + output_store = TableStoreDB( + dbconn, + "test_output3", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("test_output3", output_store) + + call_data = {"calls": [], "error_count": 0} + + def transform_func(df): + call_data["calls"].append(sorted(df["id"].tolist())) + # Имитируем ошибку на первом запуске для id=2 + if call_data["error_count"] == 0 and "2" in df["id"].tolist(): + call_data["error_count"] += 1 + # Выбрасываем исключение чтобы запись была помечена как error + raise ValueError("Test error for id=2") + return df.rename(columns={"value": "result"}) + + step = BatchTransformStep( + ds=ds, + name="test_transform3", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=1, # Обрабатываем по одной записи чтобы ошибка была только для id=2 + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2", "3"], "value": [10, 20, 30]}), now=now + ) + + # Первый запуск - id=2 не обработается + step.run_full(ds) + + # Проверяем, что обработались только 1 и 3 + output_data = output_dt.get_data() + assert len(output_data) == 2 + assert sorted(output_data["id"].tolist()) == ["1", "3"] + + # Второй запуск - должен переобработать id=2 (т.к. он остался с is_success=False) + step.run_full(ds) + + # Проверяем, что теперь все обработалось + output_data = output_dt.get_data() + assert len(output_data) == 3 + assert sorted(output_data["id"].tolist()) == ["1", "2", "3"] + + +def test_batch_transform_offset_multiple_inputs(dbconn: DBConn): + """Тест работы с несколькими входными таблицами и offset оптимизацией""" + ds = DataStore(dbconn, create_meta_table=True) + + # Две входных таблицы + input1_store = TableStoreDB( + dbconn, + "test_input_a", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input1_dt = ds.create_table("test_input_a", input1_store) + + input2_store = TableStoreDB( + dbconn, + "test_input_b", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input2_dt = ds.create_table("test_input_b", input2_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "test_output_multi", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("test_output_multi", output_store) + + def transform_func(df1, df2): + # Объединяем данные из обеих таблиц + return pd.concat([df1, df2]).rename(columns={"value": "result"}) + + step = BatchTransformStep( + ds=ds, + name="test_transform_multi", + func=transform_func, + input_dts=[ + ComputeInput(dt=input1_dt, join_type="full"), + ComputeInput(dt=input2_dt, join_type="full"), + ], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # Добавляем данные в обе таблицы + now = time.time() + input1_dt.store_chunk(pd.DataFrame({"id": ["1", "2"], "value": [10, 20]}), now=now) + input2_dt.store_chunk(pd.DataFrame({"id": ["3", "4"], "value": [30, 40]}), now=now) + + # Первый запуск + step.run_full(ds) + + output_data = output_dt.get_data() + assert len(output_data) == 4 + assert sorted(output_data["id"].tolist()) == ["1", "2", "3", "4"] + + # Добавляем новые данные только в первую таблицу + time.sleep(0.01) + now2 = time.time() + input1_dt.store_chunk(pd.DataFrame({"id": ["5"], "value": [50]}), now=now2) + + # Второй запуск - должны обработаться только новые из первой таблицы + step.run_full(ds) + + output_data = output_dt.get_data() + assert len(output_data) == 5 + assert sorted(output_data["id"].tolist()) == ["1", "2", "3", "4", "5"] + + +def test_batch_transform_without_offset_vs_with_offset(dbconn: DBConn): + """Сравнение работы без и с offset оптимизацией""" + ds = DataStore(dbconn, create_meta_table=True) + + # Входная таблица (общая) + input_store = TableStoreDB( + dbconn, + "test_input_compare", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("test_input_compare", input_store) + + # Две выходные таблицы - одна без offset, другая с offset + output_no_offset_store = TableStoreDB( + dbconn, + "test_output_no_offset", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_no_offset_dt = ds.create_table("test_output_no_offset", output_no_offset_store) + + output_with_offset_store = TableStoreDB( + dbconn, + "test_output_with_offset", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_with_offset_dt = ds.create_table("test_output_with_offset", output_with_offset_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + # Step без offset + step_no_offset = BatchTransformStep( + ds=ds, + name="test_transform_no_offset", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_no_offset_dt], + transform_keys=["id"], + use_offset_optimization=False, + ) + + # Step с offset + step_with_offset = BatchTransformStep( + ds=ds, + name="test_transform_with_offset", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_with_offset_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2", "3"], "value": [10, 20, 30]}), now=now + ) + + # Первый запуск обоих + step_no_offset.run_full(ds) + step_with_offset.run_full(ds) + + # Результаты должны быть идентичны + data_no_offset = output_no_offset_dt.get_data().sort_values("id").reset_index(drop=True) + data_with_offset = output_with_offset_dt.get_data().sort_values("id").reset_index(drop=True) + pd.testing.assert_frame_equal(data_no_offset, data_with_offset) + + # Добавляем новые данные + time.sleep(0.01) + now2 = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["4", "5"], "value": [40, 50]}), now=now2 + ) + + # Второй запуск + step_no_offset.run_full(ds) + step_with_offset.run_full(ds) + + # Результаты снова должны быть идентичны + data_no_offset = output_no_offset_dt.get_data().sort_values("id").reset_index(drop=True) + data_with_offset = output_with_offset_dt.get_data().sort_values("id").reset_index(drop=True) + pd.testing.assert_frame_equal(data_no_offset, data_with_offset) + + +def test_batch_transform_offset_no_new_data(dbconn: DBConn): + """Тест что при отсутствии новых данных ничего не обрабатывается""" + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "test_input_no_new", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("test_input_no_new", input_store) + + output_store = TableStoreDB( + dbconn, + "test_output_no_new", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("test_output_no_new", output_store) + + call_count = {"count": 0} + + def transform_func(df): + call_count["count"] += 1 + return df.rename(columns={"value": "result"}) + + step = BatchTransformStep( + ds=ds, + name="test_transform_no_new", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2"], "value": [10, 20]}), now=now + ) + + # Первый запуск + step.run_full(ds) + assert call_count["count"] == 1 + + # Второй запуск без новых данных + step.run_full(ds) + # Функция не должна вызваться, если нет новых данных для обработки + # (или вызваться с пустым df) + # В зависимости от реализации может быть 1 или 2 + # Проверяем что результаты не изменились + output_data = output_dt.get_data() + assert len(output_data) == 2 diff --git a/tests/test_build_changed_idx_sql_v2.py b/tests/test_build_changed_idx_sql_v2.py new file mode 100644 index 00000000..e01fd45b --- /dev/null +++ b/tests/test_build_changed_idx_sql_v2.py @@ -0,0 +1,277 @@ +""" +Тесты для build_changed_idx_sql_v2 - новой оптимизированной версии с offset'ами +""" +import time + +import pandas as pd +from sqlalchemy import Column, Integer, String + +from datapipe.compute import ComputeInput, Table +from datapipe.datatable import DataStore +from datapipe.meta.sql_meta import TransformMetaTable, build_changed_idx_sql_v2 +from datapipe.store.database import DBConn, TableStoreDB + + +def test_build_changed_idx_sql_v2_basic(dbconn: DBConn): + """Тест базовой работы build_changed_idx_sql_v2""" + ds = DataStore(dbconn, create_meta_table=True) + + # Создаем входную таблицу + input_table_store = TableStoreDB( + dbconn, + "test_input", + [ + Column("id", String, primary_key=True), + Column("value", String), + ], + create_table=True, + ) + input_dt = ds.create_table("test_input", input_table_store) + + # Создаем TransformMetaTable + transform_meta = TransformMetaTable( + dbconn, + "test_transform_meta", + [Column("id", String, primary_key=True)], + create_table=True, + ) + + # Добавляем данные во входную таблицу + now = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2", "3"], "value": ["a", "b", "c"]}), now=now + ) + + # Устанавливаем offset + ds.offset_table.update_offset("test_transform", "test_input", now - 10) + + # Создаем ComputeInput + compute_input = ComputeInput(dt=input_dt, join_type="full") + + # Строим SQL с использованием v2 + transform_keys, sql = build_changed_idx_sql_v2( + ds=ds, + meta_table=transform_meta, + input_dts=[compute_input], + transform_keys=["id"], + offset_table=ds.offset_table, + transformation_id="test_transform", + ) + + # Проверяем, что SQL компилируется + # После исправления гипотезы 2, update_ts добавляется в all_select_keys для ORDER BY + assert transform_keys == ["id", "update_ts"] + assert sql is not None + + # Выполняем SQL и проверяем результат + with dbconn.con.begin() as con: + result = con.execute(sql).fetchall() + # Должны получить записи, добавленные после offset + assert len(result) == 3 + + +def test_build_changed_idx_sql_v2_with_offset_filters_new_records(dbconn: DBConn): + """Тест что offset фильтрует старые записи""" + ds = DataStore(dbconn, create_meta_table=True) + + input_table_store = TableStoreDB( + dbconn, + "test_input2", + [ + Column("id", String, primary_key=True), + Column("value", String), + ], + create_table=True, + ) + input_dt = ds.create_table("test_input2", input_table_store) + + transform_meta = TransformMetaTable( + dbconn, + "test_transform_meta2", + [Column("id", String, primary_key=True)], + create_table=True, + ) + + # Добавляем старые данные + old_time = time.time() - 100 + input_dt.store_chunk(pd.DataFrame({"id": ["1", "2"], "value": ["a", "b"]}), now=old_time) + + # Устанавливаем offset после старых данных + ds.offset_table.update_offset("test_transform2", "test_input2", old_time + 10) + + # Добавляем новые данные + new_time = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["3", "4"], "value": ["c", "d"]}), now=new_time) + + compute_input = ComputeInput(dt=input_dt, join_type="full") + + transform_keys, sql = build_changed_idx_sql_v2( + ds=ds, + meta_table=transform_meta, + input_dts=[compute_input], + transform_keys=["id"], + offset_table=ds.offset_table, + transformation_id="test_transform2", + ) + + # Выполняем SQL + with dbconn.con.begin() as con: + result = con.execute(sql).fetchall() + ids = sorted([row[1] for row in result]) # row[0] is _datapipe_dummy + # Должны получить только новые записи (3, 4), старые (1, 2) отфильтрованы offset'ом + assert ids == ["3", "4"] + + +def test_build_changed_idx_sql_v2_with_error_records(dbconn: DBConn): + """Тест что ошибочные записи попадают в выборку""" + ds = DataStore(dbconn, create_meta_table=True) + + input_table_store = TableStoreDB( + dbconn, + "test_input3", + [ + Column("id", String, primary_key=True), + Column("value", String), + ], + create_table=True, + ) + input_dt = ds.create_table("test_input3", input_table_store) + + transform_meta = TransformMetaTable( + dbconn, + "test_transform_meta3", + [Column("id", String, primary_key=True)], + create_table=True, + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["1", "2", "3"], "value": ["a", "b", "c"]}), now=now) + + # Устанавливаем offset + ds.offset_table.update_offset("test_transform3", "test_input3", now + 10) + + # Добавляем ошибочную запись в transform_meta + transform_meta.insert_rows(pd.DataFrame({"id": ["error_id"]})) + # Оставляем is_success=False и process_ts=0 (дефолтные значения) + + compute_input = ComputeInput(dt=input_dt, join_type="full") + + transform_keys, sql = build_changed_idx_sql_v2( + ds=ds, + meta_table=transform_meta, + input_dts=[compute_input], + transform_keys=["id"], + offset_table=ds.offset_table, + transformation_id="test_transform3", + ) + + # Выполняем SQL + with dbconn.con.begin() as con: + result = con.execute(sql).fetchall() + ids = sorted([row[1] for row in result]) + # Должны получить только error_id (новые записи отфильтрованы offset'ом) + assert ids == ["error_id"] + + +def test_build_changed_idx_sql_v2_multiple_inputs(dbconn: DBConn): + """Тест с несколькими входными таблицами""" + ds = DataStore(dbconn, create_meta_table=True) + + # Создаем две входных таблицы + input1_store = TableStoreDB( + dbconn, + "test_input_a", + [Column("id", String, primary_key=True), Column("value", String)], + create_table=True, + ) + input1_dt = ds.create_table("test_input_a", input1_store) + + input2_store = TableStoreDB( + dbconn, + "test_input_b", + [Column("id", String, primary_key=True), Column("value", String)], + create_table=True, + ) + input2_dt = ds.create_table("test_input_b", input2_store) + + transform_meta = TransformMetaTable( + dbconn, + "test_transform_multi", + [Column("id", String, primary_key=True)], + create_table=True, + ) + + # Добавляем данные в обе таблицы + now = time.time() + input1_dt.store_chunk(pd.DataFrame({"id": ["1", "2"], "value": ["a", "b"]}), now=now) + input2_dt.store_chunk(pd.DataFrame({"id": ["3", "4"], "value": ["c", "d"]}), now=now) + + # Устанавливаем offset для обеих таблиц + ds.offset_table.update_offset("test_transform_multi", "test_input_a", now - 10) + ds.offset_table.update_offset("test_transform_multi", "test_input_b", now - 10) + + compute_inputs = [ + ComputeInput(dt=input1_dt, join_type="full"), + ComputeInput(dt=input2_dt, join_type="full"), + ] + + transform_keys, sql = build_changed_idx_sql_v2( + ds=ds, + meta_table=transform_meta, + input_dts=compute_inputs, + transform_keys=["id"], + offset_table=ds.offset_table, + transformation_id="test_transform_multi", + ) + + # Выполняем SQL + with dbconn.con.begin() as con: + result = con.execute(sql).fetchall() + ids = sorted([row[1] for row in result]) + # Должны получить UNION записей из обеих таблиц + assert ids == ["1", "2", "3", "4"] + + +def test_build_changed_idx_sql_v2_no_offset_processes_all(dbconn: DBConn): + """Тест что при отсутствии offset обрабатываются все данные""" + ds = DataStore(dbconn, create_meta_table=True) + + input_table_store = TableStoreDB( + dbconn, + "test_input_no_offset", + [Column("id", String, primary_key=True), Column("value", String)], + create_table=True, + ) + input_dt = ds.create_table("test_input_no_offset", input_table_store) + + transform_meta = TransformMetaTable( + dbconn, + "test_transform_no_offset", + [Column("id", String, primary_key=True)], + create_table=True, + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["1", "2", "3"], "value": ["a", "b", "c"]}), now=now) + + # НЕ устанавливаем offset - первый запуск + + compute_input = ComputeInput(dt=input_dt, join_type="full") + + transform_keys, sql = build_changed_idx_sql_v2( + ds=ds, + meta_table=transform_meta, + input_dts=[compute_input], + transform_keys=["id"], + offset_table=ds.offset_table, + transformation_id="test_transform_no_offset", + ) + + # Выполняем SQL + with dbconn.con.begin() as con: + result = con.execute(sql).fetchall() + ids = sorted([row[1] for row in result]) + # При offset=None (дефолт 0.0) должны получить все записи + assert ids == ["1", "2", "3"] diff --git a/tests/test_complex_pipeline_with_offset.py b/tests/test_complex_pipeline_with_offset.py new file mode 100644 index 00000000..48ee248f --- /dev/null +++ b/tests/test_complex_pipeline_with_offset.py @@ -0,0 +1,439 @@ +from typing import cast + +import pandas as pd +import pytest +from sqlalchemy import Column +from sqlalchemy.sql.sqltypes import Integer, String + +from datapipe.compute import Catalog, Pipeline, Table, build_compute, run_steps +from datapipe.datatable import DataStore +from datapipe.step.batch_generate import BatchGenerate +from datapipe.step.batch_transform import BatchTransform +from datapipe.store.database import TableStoreDB +from datapipe.tests.util import assert_datatable_equal, assert_df_equal +from datapipe.types import IndexDF, Required, JoinSpec + +TEST__ITEM = pd.DataFrame( + { + "item_id": [f"item_id{i}" for i in range(10)], + "item__attribute": [f"item__attribute{i}" for i in range(10)], + } +) + +TEST__KEYPOINT = pd.DataFrame( + { + "keypoint_id": list(range(5)), + "keypoint_name": [f"keypoint_name{i}" for i in range(5)], + } +) + +TEST__PIPELINE = pd.DataFrame( + { + "pipeline_id": [f"pipeline_id{i}" for i in range(3)], + "pipeline__attribute": [f"pipeline_attribute{i}" for i in range(3)], + } +) + +TEST__PREDICTION_LEFT = pd.DataFrame( + { + "item_id": [f"item_id{i}" for i in range(10)], + } +) +TEST__PREDICTION_CENTER = pd.DataFrame( + { + "pipeline_id": [f"pipeline_id{i}" for i in range(3)], + } +) +TEST__PREDICTION_RIGHT = pd.DataFrame( + { + "keypoint_name": [f"keypoint_name{i}" for i in range(5)], + "prediction__attribute": [f"prediction__attribute{i}" for i in range(5)], + } +) +TEST__PREDICTION = pd.merge(TEST__PREDICTION_LEFT, TEST__PREDICTION_CENTER, how="cross") +TEST__PREDICTION = pd.merge(TEST__PREDICTION, TEST__PREDICTION_RIGHT, how="cross") +TEST__PREDICTION = TEST__PREDICTION[["item_id", "pipeline_id", "keypoint_name", "prediction__attribute"]] + + +def test_complex_pipeline_with_offset_optimization(dbconn): + ds = DataStore(dbconn, create_meta_table=True) + catalog = Catalog( + { + "item": Table( + store=TableStoreDB( + dbconn, + "item", + [ + Column("item_id", String, primary_key=True), + Column("item__attribute", String), + ], + True, + ) + ), + "pipeline": Table( + store=TableStoreDB( + dbconn, + "pipeline", + [ + Column("pipeline_id", String, primary_key=True), + Column("pipeline__attribute", String), + ], + True, + ) + ), + "prediction": Table( + store=TableStoreDB( + dbconn, + "prediction", + [ + Column("item_id", String, primary_key=True), + Column("pipeline_id", String, primary_key=True), + Column("keypoint_name", String, primary_key=True), + Column("prediction__attribute", String), + ], + True, + ) + ), + "keypoint": Table( + store=TableStoreDB( + dbconn, + "keypoint", + [ + Column("keypoint_id", Integer, primary_key=True), + Column("keypoint_name", String, primary_key=True), + ], + True, + ) + ), + "output": Table( + store=TableStoreDB( + dbconn, + "output", + [ + Column("item_id", String, primary_key=True), + Column("pipeline_id", String, primary_key=True), + Column("attirbute", String), + ], + True, + ) + ), + } + ) + + def complex_function(df__item, df__pipeline, df__prediction, df__keypoint, idx: IndexDF): + assert idx[idx[["item_id", "pipeline_id"]].duplicated()].empty + assert len(df__keypoint) == len(TEST__KEYPOINT) + df__output = pd.merge(df__item, df__prediction, on=["item_id"]) + df__output = pd.merge(df__output, df__pipeline, on=["pipeline_id"]) + df__output = pd.merge(df__output, df__keypoint, on=["keypoint_name"]) + df__output = df__output[["item_id", "pipeline_id"]].drop_duplicates() + df__output["attirbute"] = "attribute" + return df__output + + pipeline = Pipeline( + [ + BatchTransform( + func=complex_function, + inputs=[ + Required("item", join_keys={"item_id": "item_id"}), # inner join с reverse join + Required("pipeline", join_keys={"pipeline_id": "pipeline_id"}), # inner join с reverse join + "prediction", # full join, основная таблица (содержит все transform_keys) + Required("keypoint"), # inner join, референсная таблица без reverse join + ], + outputs=["output"], + transform_keys=["item_id", "pipeline_id"], + chunk_size=50, + use_offset_optimization=True, + ), + ] + ) + steps = build_compute(ds, catalog, pipeline) + ds.get_table("item").store_chunk(TEST__ITEM) + ds.get_table("pipeline").store_chunk(TEST__PIPELINE) + ds.get_table("prediction").store_chunk(TEST__PREDICTION) + ds.get_table("keypoint").store_chunk(TEST__KEYPOINT) + TEST_RESULT = complex_function( + TEST__ITEM, + TEST__PIPELINE, + TEST__PREDICTION, + TEST__KEYPOINT, + idx=cast(IndexDF, pd.DataFrame(columns=["item_id", "pipeline_id"])), + ) + run_steps(ds, steps) + assert_datatable_equal(ds.get_table("output"), TEST_RESULT) + + +TEST__FROZEN_DATASET = pd.DataFrame( + { + "frozen_dataset_id": [f"frozen_dataset_id{i}" for i in range(2)], + } +) + +TEST__TRAIN_CONFIG = pd.DataFrame( + { + "train_config_id": [f"train_config_id{i}" for i in range(2)], + "train_config__params": [f"train_config__params{i}" for i in range(2)], + } +) + + +def test_complex_train_pipeline_with_offset_optimization(dbconn): + ds = DataStore(dbconn, create_meta_table=True) + catalog = Catalog( + { + "frozen_dataset": Table( + store=TableStoreDB( + dbconn, + "frozen_dataset", + [ + Column("frozen_dataset_id", String, primary_key=True), + ], + True, + ) + ), + "train_config": Table( + store=TableStoreDB( + dbconn, + "train_config", + [ + Column("train_config_id", String, primary_key=True), + Column("train_config__params", String), + ], + True, + ) + ), + "pipeline": Table( + store=TableStoreDB( + dbconn, + "pipeline", + [ + Column("pipeline_id", String, primary_key=True), + Column("pipeline__attribute", String), + ], + True, + ) + ), + "pipeline__is_trained_on__frozen_dataset": Table( + store=TableStoreDB( + dbconn, + "pipeline__is_trained_on__frozen_dataset", + [ + Column("pipeline_id", String, primary_key=True), + Column("frozen_dataset_id", String, primary_key=True), + Column("train_config_id", String, primary_key=True), + ], + True, + ) + ), + } + ) + + def train( + df__frozen_dataset, + df__train_config, + df__pipeline__total, + df__pipeline__is_trained_on__frozen_dataset__total, + ): + assert len(df__frozen_dataset) == 1 and len(df__train_config) == 1 + frozen_dataset_id = df__frozen_dataset.iloc[0]["frozen_dataset_id"] + train_config_id = df__train_config.iloc[0]["train_config_id"] + df__pipeline = pd.DataFrame( + [ + { + "pipeline_id": f"new_pipeline__{frozen_dataset_id}x{train_config_id}", + "pipeline__attribute": "new_pipeline__attr", + } + ] + ) + df__pipeline__is_trained_on__frozen_dataset = pd.DataFrame( + [ + { + "pipeline_id": f"new_pipeline__{frozen_dataset_id}x{train_config_id}", + "frozen_dataset_id": frozen_dataset_id, + "train_config_id": train_config_id, + } + ] + ) + df__pipeline__total = pd.concat([df__pipeline__total, df__pipeline], ignore_index=True) + df__pipeline__is_trained_on__frozen_dataset__total = pd.concat( + [ + df__pipeline__is_trained_on__frozen_dataset__total, + df__pipeline__is_trained_on__frozen_dataset, + ], + ignore_index=True, + ) + return df__pipeline__total, df__pipeline__is_trained_on__frozen_dataset__total + + pipeline = Pipeline( + [ + BatchTransform( + func=train, + inputs=[ + "frozen_dataset", # full join (часть transform_keys) + "train_config", # full join (часть transform_keys) + "pipeline", # full join, выходная таблица + "pipeline__is_trained_on__frozen_dataset", # full join, выходная таблица + ], + outputs=["pipeline", "pipeline__is_trained_on__frozen_dataset"], + transform_keys=["frozen_dataset_id", "train_config_id"], + chunk_size=1, + use_offset_optimization=True, + ), + ] + ) + steps = build_compute(ds, catalog, pipeline) + ds.get_table("frozen_dataset").store_chunk(TEST__FROZEN_DATASET) + ds.get_table("train_config").store_chunk(TEST__TRAIN_CONFIG) + run_steps(ds, steps) + assert len(ds.get_table("pipeline").get_data()) == len(TEST__FROZEN_DATASET) * len(TEST__TRAIN_CONFIG) + assert len(ds.get_table("pipeline__is_trained_on__frozen_dataset").get_data()) == len(TEST__FROZEN_DATASET) * len( + TEST__TRAIN_CONFIG + ) + + +def complex_transform_with_many_recordings(dbconn, N: int): + ds = DataStore(dbconn, create_meta_table=True) + catalog = Catalog( + { + "tbl_image": Table( + store=TableStoreDB( + dbconn, + "tbl_image", + [ + Column("image_id", Integer, primary_key=True), + ], + True, + ) + ), + "tbl_image__attribute": Table( + store=TableStoreDB( + dbconn, + "tbl_image__attribute", + [ + Column("image_id", Integer, primary_key=True), + Column("attribute", Integer), + ], + True, + ) + ), + "tbl_prediction": Table( + store=TableStoreDB( + dbconn, + "tbl_prediction", + [ + Column("image_id", Integer, primary_key=True), + Column("model_id", Integer, primary_key=True), + Column("prediction__attribite", Integer), + ], + True, + ) + ), + "tbl_best_model": Table( + store=TableStoreDB( + dbconn, + "tbl_best_model", + [ + Column("model_id", Integer, primary_key=True), + ], + True, + ) + ), + "tbl_output": Table( + store=TableStoreDB( + dbconn, + "tbl_output", + [ + Column("image_id", Integer, primary_key=True), + Column("model_id", Integer, primary_key=True), + Column("result", Integer), + ], + True, + ) + ), + } + ) + + def gen_tbls(df1, df2, df3, df4): + yield df1, df2, df3, df4 + + test_df__image = pd.DataFrame({"image_id": range(N)}) + test_df__image__attribute = pd.DataFrame({"image_id": range(N), "attribute": [5 * x for x in range(N)]}) + test_df__prediction = pd.DataFrame( + { + "image_id": list(range(N)) * 5, + "model_id": [0] * N + [1] * N + [2] * N + [3] * N + [4] * N, + "prediction__attribite": ( + [1 * x for x in range(N)] # model_id=0 + + [2 * x for x in range(N)] # model_id=1 + + [3 * x for x in range(N)] # model_id=2 + + [4 * x for x in range(N)] # model_id=3 + + [5 * x for x in range(N)] # model_id=4 + ), + } + ) + test_df__best_model = pd.DataFrame({"model_id": [4]}) + + def get_some_prediction_only_on_best_model( + df__image: pd.DataFrame, + df__image__attribute: pd.DataFrame, + df__prediction: pd.DataFrame, + df__best_model: pd.DataFrame, + ): + df__prediction = pd.merge(df__prediction, df__best_model, on=["model_id"]) + df__image = pd.merge(df__image, df__image__attribute, on=["image_id"]) + df__result = pd.merge(df__image, df__prediction, on=["image_id"]) + df__result["result"] = df__result["attribute"] - df__result["prediction__attribite"] + return df__result[["image_id", "model_id", "result"]] + + pipeline = Pipeline( + [ + BatchGenerate( + func=gen_tbls, + outputs=[ + "tbl_image", + "tbl_image__attribute", + "tbl_prediction", + "tbl_best_model", + ], + kwargs=dict( + df1=test_df__image, + df2=test_df__image__attribute, + df3=test_df__prediction, + df4=test_df__best_model, + ), + ), + BatchTransform( + func=get_some_prediction_only_on_best_model, + inputs=[ + Required("tbl_image", join_keys={"image_id": "image_id"}), # inner join с reverse join + Required("tbl_image__attribute", join_keys={"image_id": "image_id"}), # inner join с reverse join + "tbl_prediction", # full join, содержит все transform_keys (image_id, model_id) + Required("tbl_best_model", join_keys={"model_id": "model_id"}), # inner join с reverse join + ], + outputs=["tbl_output"], + transform_keys=["image_id", "model_id"], + use_offset_optimization=True, + ), + ] + ) + steps = build_compute(ds, catalog, pipeline) + run_steps(ds, steps) + test__df_output = pd.DataFrame({"image_id": range(N), "model_id": [4] * N, "result": [0] * N}) + assert_df_equal( + ds.get_table("tbl_output").get_data(), + test__df_output, + index_cols=["image_id", "model_id"], + ) + + +def test_complex_transform_with_many_recordings_N100_with_offset_optimization(dbconn): + complex_transform_with_many_recordings(dbconn, N=100) + + +def test_complex_transform_with_many_recordings_N1000_with_offset_optimization(dbconn): + complex_transform_with_many_recordings(dbconn, N=1000) + + +@pytest.mark.skip(reason="fails on sqlite") +def test_complex_transform_with_many_recordings_N10000_with_offset_optimization(dbconn): + complex_transform_with_many_recordings(dbconn, N=10000) diff --git a/tests/test_deleted_records_with_join_keys.py b/tests/test_deleted_records_with_join_keys.py new file mode 100644 index 00000000..8f8e35a0 --- /dev/null +++ b/tests/test_deleted_records_with_join_keys.py @@ -0,0 +1,215 @@ +""" +Тест для воспроизведения бага с удаленными записями при использовании join_keys. + +Проблема: +При использовании join_keys и FK в data table, удаленные записи не попадают +в processed_idx из-за INNER JOIN, что приводит к тому что они не удаляются +из output таблицы. + +Сценарий: +1. Создать user и subscription таблицы +2. subscription имеет ДВА join_keys для filtered join: + - join_keys={"follower_id": "id"} (кто подписывается) + - join_keys={"following_id": "id"} (на кого подписываются) +3. Создать 3 subscription: + - sub1: будет удалена + - sub2: будет изменена + - sub3: без изменений +4. Запустить трансформацию → все записи в output +5. Удалить sub1, изменить sub2 +6. Запустить трансформацию → проверить что: + - sub1 удалена из output (удаленная запись) + - sub2 обновлена и присутствует (измененная запись) + - sub3 присутствует без изменений (неизмененная запись) +""" + +import time + +import pandas as pd +from sqlalchemy import Column, String + +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_deleted_records_with_join_keys(dbconn: DBConn): + """ + Воспроизводит баг: удаленные записи не попадают в processed_idx + при использовании join_keys, и не удаляются из output. + """ + ds = DataStore(dbconn, create_meta_table=True) + + # 1. Создать user таблицу + user_store = TableStoreDB( + dbconn, + "users", + [Column("id", String, primary_key=True), Column("name", String)], + create_table=True, + ) + user_dt = ds.create_table("users", user_store) + + # 2. Создать subscription таблицу + subscription_store = TableStoreDB( + dbconn, + "subscriptions", + [ + Column("id", String, primary_key=True), + Column("follower_id", String), # FK в data table! + Column("following_id", String), + ], + create_table=True, + ) + subscription_dt = ds.create_table("subscriptions", subscription_store) + + # 3. Создать output таблицу + output_store = TableStoreDB( + dbconn, + "enriched_subscriptions", + [ + Column("id", String, primary_key=True), + Column("follower_name", String), + Column("following_name", String), + ], + create_table=True, + ) + output_dt = ds.create_table("enriched_subscriptions", output_store) + + # 4. Трансформация с ДВА join_keys (получает ТРИ DataFrame) + def transform_func(subscription_df, followers_df, followings_df): + # JOIN для follower (кто подписывается) + result = subscription_df.merge( + followers_df, + left_on="follower_id", + right_on="id", + how="left", + suffixes=("", "_follower"), + ) + # JOIN для following (на кого подписываются) + result = result.merge( + followings_df, + left_on="following_id", + right_on="id", + how="left", + suffixes=("", "_following"), + ) + # После merge колонки будут: id, follower_id, following_id, id_follower, name, id_following, name_following + # name - это имя follower, name_following - имя following + return result[["id", "name", "name_following"]].rename( + columns={ + "name": "follower_name", + "name_following": "following_name", + } + ) + + step = BatchTransformStep( + ds=ds, + name="test_join_keys_delete", + func=transform_func, + input_dts=[ + ComputeInput(dt=subscription_dt, join_type="full"), + # ← КРИТИЧНО: ДВА отдельных ComputeInput для user_dt с разными join_keys! + ComputeInput( + dt=user_dt, + join_type="full", + join_keys={"follower_id": "id"}, # Follower + ), + ComputeInput( + dt=user_dt, + join_type="full", + join_keys={"following_id": "id"}, # Following + ), + ], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=10, + ) + + # === НАЧАЛЬНЫЕ ДАННЫЕ === + + t1 = time.time() + + # Создать пользователей + users_df = pd.DataFrame( + [ + {"id": "u1", "name": "Alice"}, + {"id": "u2", "name": "Bob"}, + {"id": "u3", "name": "Charlie"}, + ] + ) + user_dt.store_chunk(users_df, now=t1) + + # Создать подписки + subscriptions_df = pd.DataFrame( + [ + {"id": "sub1", "follower_id": "u1", "following_id": "u2"}, # Будет удалена + {"id": "sub2", "follower_id": "u2", "following_id": "u3"}, # Будет изменена + {"id": "sub3", "follower_id": "u3", "following_id": "u1"}, # Без изменений + ] + ) + subscription_dt.store_chunk(subscriptions_df, now=t1) + + # === ПЕРВАЯ ОБРАБОТКА === + + time.sleep(0.01) + step.run_full(ds) + + # Проверяем что все 3 записи в output + output_data = output_dt.get_data() + assert len(output_data) == 3, f"Expected 3 records in output, got {len(output_data)}" + assert set(output_data["id"]) == {"sub1", "sub2", "sub3"} + + # === ИЗМЕНЕНИЯ === + + time.sleep(0.01) + t2 = time.time() + + # Удаляем sub1 + subscription_dt.delete_by_idx(pd.DataFrame({"id": ["sub1"]}), now=t2) + + # Изменяем sub2 (новый follower) + time.sleep(0.01) + t3 = time.time() + subscription_dt.store_chunk( + pd.DataFrame( + [ + {"id": "sub2", "follower_id": "u1", "following_id": "u3"} # follower u2 → u1 + ] + ), + now=t3, + ) + + # === ВТОРАЯ ОБРАБОТКА === + + time.sleep(0.01) + step.run_full(ds) + + # === ПРОВЕРКА === + + output_data_after = output_dt.get_data() + output_ids = set(output_data_after["id"]) + + # ОЖИДАНИЕ: sub1 должна быть удалена + assert "sub1" not in output_ids, ( + f"БАГ: sub1 должна быть удалена из output, но она есть! " + f"Это происходит потому что при INNER JOIN с data table удаленные записи " + f"не попадают в processed_idx. Output IDs: {output_ids}" + ) + + # ОЖИДАНИЕ: sub2 должна быть обновлена + assert "sub2" in output_ids, f"sub2 должна быть в output" + sub2_data = output_data_after[output_data_after["id"] == "sub2"].iloc[0] + assert sub2_data["follower_name"] == "Alice", ( + f"sub2 должна иметь обновленного follower (Alice), " f"got {sub2_data['follower_name']}" + ) + assert sub2_data["following_name"] == "Charlie", ( + f"sub2 должна иметь following (Charlie), " f"got {sub2_data['following_name']}" + ) + + # ОЖИДАНИЕ: sub3 без изменений + assert "sub3" in output_ids, f"sub3 должна быть в output" + + # ОЖИДАНИЕ: Ровно 2 записи в output + assert len(output_data_after) == 2, f"Expected 2 records in output (sub2, sub3), got {len(output_data_after)}" diff --git a/tests/test_initialize_offsets.py b/tests/test_initialize_offsets.py new file mode 100644 index 00000000..61158bc8 --- /dev/null +++ b/tests/test_initialize_offsets.py @@ -0,0 +1,296 @@ +""" +Тесты для Phase 4: Инициализация offset'ов из существующих TransformMetaTable + +Проверяет функцию initialize_offsets_from_transform_meta() которая устанавливает +начальные offset'ы на основе уже обработанных данных. +""" +import time + +import pandas as pd +from sqlalchemy import Column, Integer, String + +from datapipe.compute import ComputeInput +from datapipe.datatable import DataStore +from datapipe.meta.sql_meta import initialize_offsets_from_transform_meta +from datapipe.step.batch_transform import BatchTransformStep +from datapipe.store.database import DBConn, TableStoreDB + + +def test_initialize_offsets_from_empty_transform_meta(dbconn: DBConn): + """ + Тест инициализации offset'ов когда TransformMetaTable пустая. + + Ожидаемое поведение: offset'ы не устанавливаются (нет данных для инициализации) + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "init_empty_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("init_empty_input", input_store) + + output_store = TableStoreDB( + dbconn, + "init_empty_output", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("init_empty_output", output_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + step = BatchTransformStep( + ds=ds, + name="init_empty_transform", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=False, # Не используем автообновление + ) + + # Инициализируем offset'ы (TransformMetaTable пустая) + result = initialize_offsets_from_transform_meta(ds, step) + + # Проверяем что offset'ы не установлены + assert result == {} + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert len(offsets) == 0 + + +def test_initialize_offsets_from_existing_data(dbconn: DBConn): + """ + Тест инициализации offset'ов из существующих обработанных данных. + + Сценарий: + 1. Обрабатываем данные старым методом (v1, без offset'ов) + 2. Инициализируем offset'ы на основе уже обработанных данных + 3. Проверяем что offset'ы установлены корректно + 4. Добавляем новые данные и проверяем что обрабатываются только они + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "init_existing_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("init_existing_input", input_store) + + output_store = TableStoreDB( + dbconn, + "init_existing_output", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("init_existing_output", output_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + # Создаем step БЕЗ offset optimization + step = BatchTransformStep( + ds=ds, + name="init_existing_transform", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=False, # Используем старый метод + ) + + # ========== Шаг 1: Обрабатываем данные старым методом ========== + now1 = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["A", "B", "C"], "value": [10, 20, 30]}), now=now1 + ) + + step.run_full(ds) + + # Проверяем что данные обработались + output_data = output_dt.get_data() + assert len(output_data) == 3 + + # ========== Шаг 2: Инициализируем offset'ы ========== + result = initialize_offsets_from_transform_meta(ds, step) + + # Проверяем что offset был установлен + assert len(result) == 1 + assert "init_existing_input" in result + assert result["init_existing_input"] >= now1 + + # Проверяем что offset сохранен в БД + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert len(offsets) == 1 + assert offsets["init_existing_input"] >= now1 + + # ========== Шаг 3: Включаем offset optimization ========== + step.use_offset_optimization = True + + # ========== Шаг 4: Добавляем новые данные ========== + time.sleep(0.01) + now2 = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["D", "E"], "value": [40, 50]}), now=now2 + ) + + # ========== Шаг 5: Проверяем что видны только новые записи ========== + changed_count = step.get_changed_idx_count(ds) + assert changed_count == 2, f"Expected 2 new records, got {changed_count}" + + # Запускаем обработку с offset optimization + step.run_full(ds) + + # Проверяем что обработались все 5 записей + output_data = output_dt.get_data() + assert len(output_data) == 5 + assert sorted(output_data["id"].tolist()) == ["A", "B", "C", "D", "E"] + + +def test_initialize_offsets_with_multiple_inputs(dbconn: DBConn): + """ + Тест инициализации offset'ов для трансформации с несколькими входными таблицами. + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Первая входная таблица + input1_store = TableStoreDB( + dbconn, + "init_multi_input1", + [Column("id", String, primary_key=True), Column("value1", Integer)], + create_table=True, + ) + input1_dt = ds.create_table("init_multi_input1", input1_store) + + # Вторая входная таблица + input2_store = TableStoreDB( + dbconn, + "init_multi_input2", + [Column("id", String, primary_key=True), Column("value2", Integer)], + create_table=True, + ) + input2_dt = ds.create_table("init_multi_input2", input2_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "init_multi_output", + [Column("id", String, primary_key=True), Column("sum", Integer)], + create_table=True, + ) + output_dt = ds.create_table("init_multi_output", output_store) + + def transform_func(df1, df2): + merged = df1.merge(df2, on="id", how="outer") + merged["sum"] = merged["value1"].fillna(0) + merged["value2"].fillna(0) + return merged[["id", "sum"]].astype({"sum": int}) + + step = BatchTransformStep( + ds=ds, + name="init_multi_transform", + func=transform_func, + input_dts=[ + ComputeInput(dt=input1_dt, join_type="full"), + ComputeInput(dt=input2_dt, join_type="full"), + ], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=False, + ) + + # Добавляем данные в обе таблицы с разными timestamp + now1 = time.time() + input1_dt.store_chunk(pd.DataFrame({"id": ["X", "Y"], "value1": [1, 2]}), now=now1) + + time.sleep(0.01) + now2 = time.time() + input2_dt.store_chunk(pd.DataFrame({"id": ["X", "Y"], "value2": [10, 20]}), now=now2) + + # Обрабатываем + step.run_full(ds) + + # Инициализируем offset'ы + result = initialize_offsets_from_transform_meta(ds, step) + + # Проверяем что offset'ы установлены для обеих таблиц + assert len(result) == 2 + assert "init_multi_input1" in result + assert "init_multi_input2" in result + + # Проверяем что offset'ы корректные (учитывают MIN(process_ts)) + assert result["init_multi_input1"] >= now1 + assert result["init_multi_input2"] >= now2 + + +def test_initialize_offsets_conservative_approach(dbconn: DBConn): + """ + Тест что инициализация использует консервативный подход: + MAX(input.update_ts) <= MIN(transform.process_ts) + + Это гарантирует что мы не пропустим данные которые еще не были обработаны. + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "init_conservative_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("init_conservative_input", input_store) + + output_store = TableStoreDB( + dbconn, + "init_conservative_output", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("init_conservative_output", output_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + step = BatchTransformStep( + ds=ds, + name="init_conservative_transform", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=False, + ) + + # Добавляем первую партию данных + now1 = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2"], "value": [10, 20]}), now=now1 + ) + + # Обрабатываем первую партию + step.run_full(ds) + + # Добавляем вторую партию данных (с более поздним timestamp) + time.sleep(0.01) + now2 = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["3", "4"], "value": [30, 40]}), now=now2 + ) + + # НЕ обрабатываем вторую партию! + # Инициализируем offset'ы + result = initialize_offsets_from_transform_meta(ds, step) + + # Offset должен быть <= now1 (только для первой партии) + # Потому что MIN(process_ts) соответствует первой партии + # Это гарантирует что вторая партия (с now2) НЕ будет пропущена + assert result["init_conservative_input"] <= now1 + 0.1 # небольшой запас на overhead + + # Включаем offset optimization и проверяем что видны обе непроцессированные записи + step.use_offset_optimization = True + changed_count = step.get_changed_idx_count(ds) + assert changed_count == 2 # Должны видеть записи 3 и 4 diff --git a/tests/test_multi_table_filtered_join.py b/tests/test_multi_table_filtered_join.py new file mode 100644 index 00000000..1c8b41d3 --- /dev/null +++ b/tests/test_multi_table_filtered_join.py @@ -0,0 +1,711 @@ +""" +Тесты для проверки мульти-табличных трансформаций с filtered join оптимизацией. + +Тесты проверяют: +1. Что filtered join вызывается корректно при наличии join_keys +2. Что join выполняется по правильным ключам +3. Что результаты v1 (FULL OUTER JOIN) и v2 (offset-based) идентичны +""" +import time + +import pandas as pd +from sqlalchemy import Column, Integer, String + +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_filtered_join_is_called(dbconn: DBConn): + """ + Тест 1: Проверяет что filtered join вызывается и читает только нужные данные. + + Сценарий: + - Основная таблица users с user_id + - Справочник profiles с большим количеством записей + - join_keys указывает связь user_id -> id + - Проверяем что из profiles читаются только записи для существующих пользователей + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Основная таблица users + users_store = TableStoreDB( + dbconn, + "users", + [ + Column("user_id", String, primary_key=True), + Column("name", String), + ], + create_table=True, + ) + users_dt = ds.create_table("users", users_store) + + # Справочная таблица profiles (много записей) + profiles_store = TableStoreDB( + dbconn, + "profiles", + [ + Column("id", String, primary_key=True), + Column("description", String), + ], + create_table=True, + ) + profiles_dt = ds.create_table("profiles", profiles_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "user_profiles", + [ + Column("user_id", String, primary_key=True), + Column("name", String), + Column("description", String), + ], + create_table=True, + ) + output_dt = ds.create_table("user_profiles", output_store) + + # Функция трансформации с отслеживанием вызовов + transform_calls = [] + + def transform_func(users_df, profiles_df): + transform_calls.append({ + "users_count": len(users_df), + "profiles_count": len(profiles_df), + "profiles_ids": sorted(profiles_df["id"].tolist()) if not profiles_df.empty else [] + }) + # Merge by user_id and id + merged = users_df.merge( + profiles_df, + left_on="user_id", + right_on="id", + how="left" + ) + return merged[["user_id", "name", "description"]].fillna("") + + # ОТСЛЕЖИВАНИЕ ВЫЗОВОВ get_data для profiles_dt + get_data_calls = [] + original_get_data = profiles_dt.get_data + + def tracked_get_data(idx=None, **kwargs): + if idx is not None: + get_data_calls.append({ + "idx_columns": list(idx.columns), + "idx_values": sorted(idx["id"].tolist()) if "id" in idx.columns else [], + "idx_length": len(idx) + }) + result = original_get_data(idx=idx, **kwargs) + return result + + profiles_dt.get_data = tracked_get_data # type: ignore[method-assign] + + # Создаем step с filtered join + step = BatchTransformStep( + ds=ds, + name="test_filtered_join", + func=transform_func, + input_dts=[ + ComputeInput(dt=users_dt, join_type="full"), + ComputeInput( + dt=profiles_dt, + join_type="full", + join_keys={"user_id": "id"} # user_id из idx -> id из profiles + ), + ], + output_dts=[output_dt], + transform_keys=["user_id"], + chunk_size=10, + ) + + # Добавляем данные: только 3 пользователя + now = time.time() + users_dt.store_chunk( + pd.DataFrame({ + "user_id": ["u1", "u2", "u3"], + "name": ["Alice", "Bob", "Charlie"] + }), + now=now + ) + + # Добавляем много профилей (100 записей), но только 3 связаны с пользователями + profiles_data = [] + for i in range(100): + profiles_data.append({ + "id": f"p{i}", + "description": f"Profile {i}" + }) + # Добавляем профили для наших пользователей + profiles_data.extend([ + {"id": "u1", "description": "Alice's profile"}, + {"id": "u2", "description": "Bob's profile"}, + {"id": "u3", "description": "Charlie's profile"}, + ]) + + profiles_dt.store_chunk(pd.DataFrame(profiles_data), now=now) + + # Запускаем трансформацию + step.run_full(ds) + + # Проверяем что трансформация вызвалась + assert len(transform_calls) > 0 + + # ПРОВЕРКА 1: get_data был вызван с filtered idx + assert len(get_data_calls) > 0, "get_data should be called with filtered idx" + + first_get_data_call = get_data_calls[0] + assert "id" in first_get_data_call["idx_columns"], ( + f"idx should contain 'id' column for filtered join, " + f"but got columns: {first_get_data_call['idx_columns']}" + ) + assert first_get_data_call["idx_length"] == 3, ( + f"Filtered idx should contain 3 records (u1, u2, u3), " + f"but got {first_get_data_call['idx_length']}" + ) + assert first_get_data_call["idx_values"] == ["u1", "u2", "u3"], ( + f"Filtered idx should contain correct user_ids mapped to profile ids, " + f"but got {first_get_data_call['idx_values']}" + ) + + # ПРОВЕРКА 2: filtered join должен читать только 3 профиля, + # а не все 103 записи из таблицы profiles + first_call = transform_calls[0] + assert first_call["users_count"] == 3, "Should process 3 users" + assert first_call["profiles_count"] == 3, ( + f"Filtered join should read only 3 profiles (for u1, u2, u3), " + f"but got {first_call['profiles_count']}" + ) + assert first_call["profiles_ids"] == ["u1", "u2", "u3"], ( + "Should read profiles only for existing users" + ) + + # Проверяем результат + output_data = output_dt.get_data().sort_values("user_id").reset_index(drop=True) + assert len(output_data) == 3 + assert output_data["user_id"].tolist() == ["u1", "u2", "u3"] + assert output_data["description"].tolist() == [ + "Alice's profile", + "Bob's profile", + "Charlie's profile" + ] + + +def test_join_keys_correctness(dbconn: DBConn): + """ + Тест 2: Проверяет что join происходит по правильным ключам. + + Сценарий: + - Основная таблица orders с order_id и customer_id как часть primary key + - Справочник customers с id и другими полями + - join_keys: {"customer_id": "id"} + - Проверяем что данные джойнятся правильно + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Таблица заказов с customer_id в primary key + orders_store = TableStoreDB( + dbconn, + "orders", + [ + Column("order_id", String, primary_key=True), + Column("customer_id", String, primary_key=True), + Column("amount", Integer), + ], + create_table=True, + ) + orders_dt = ds.create_table("orders", orders_store) + + # Таблица покупателей + customers_store = TableStoreDB( + dbconn, + "customers", + [ + Column("id", String, primary_key=True), + Column("name", String), + ], + create_table=True, + ) + customers_dt = ds.create_table("customers", customers_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "enriched_orders", + [ + Column("order_id", String, primary_key=True), + Column("customer_id", String, primary_key=True), + Column("customer_name", String), + Column("amount", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("enriched_orders", output_store) + + def transform_func(orders_df, customers_df): + # Join по customer_id = id + merged = orders_df.merge( + customers_df, + left_on="customer_id", + right_on="id", + how="left" + ) + return merged[["order_id", "customer_id", "name", "amount"]].rename(columns={"name": "customer_name"}) + + # Step с join_keys + step = BatchTransformStep( + ds=ds, + name="test_join_keys", + func=transform_func, + input_dts=[ + ComputeInput(dt=orders_dt, join_type="full"), + ComputeInput( + dt=customers_dt, + join_type="full", + join_keys={"customer_id": "id"} + ), + ], + output_dts=[output_dt], + transform_keys=["order_id", "customer_id"], + ) + + # Добавляем данные + now = time.time() + orders_dt.store_chunk( + pd.DataFrame({ + "order_id": ["o1", "o2", "o3"], + "customer_id": ["c1", "c2", "c1"], + "amount": [100, 200, 150] + }), + now=now + ) + + customers_dt.store_chunk( + pd.DataFrame({ + "id": ["c1", "c2", "c3"], # c3 не используется + "name": ["John", "Jane", "Bob"] + }), + now=now + ) + + # Запускаем + step.run_full(ds) + + # Проверяем результат + output_data = output_dt.get_data().sort_values("order_id").reset_index(drop=True) + assert len(output_data) == 3 + + # КЛЮЧЕВАЯ ПРОВЕРКА: join keys работают правильно + expected = pd.DataFrame({ + "order_id": ["o1", "o2", "o3"], + "customer_id": ["c1", "c2", "c1"], + "customer_name": ["John", "Jane", "John"], + "amount": [100, 200, 150] + }) + pd.testing.assert_frame_equal(output_data, expected) + + +def test_v1_vs_v2_results_identical(dbconn: DBConn): + """ + Тест 3: Сравнивает результаты v1 (FULL OUTER JOIN) и v2 (offset-based). + + Сценарий: + - Две входные таблицы с мульти-табличной трансформацией + - Запускаем одну и ту же трансформацию с v1 и v2 + - Проверяем что результаты идентичны + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Входные таблицы + table_a_store = TableStoreDB( + dbconn, + "table_a", + [ + Column("id", String, primary_key=True), + Column("value_a", Integer), + ], + create_table=True, + ) + table_a_dt = ds.create_table("table_a", table_a_store) + + table_b_store = TableStoreDB( + dbconn, + "table_b", + [ + Column("id", String, primary_key=True), + Column("value_b", Integer), + ], + create_table=True, + ) + table_b_dt = ds.create_table("table_b", table_b_store) + + # Две выходные таблицы - одна для v1, другая для v2 + output_v1_store = TableStoreDB( + dbconn, + "output_v1", + [ + Column("id", String, primary_key=True), + Column("sum_value", Integer), + ], + create_table=True, + ) + output_v1_dt = ds.create_table("output_v1", output_v1_store) + + output_v2_store = TableStoreDB( + dbconn, + "output_v2", + [ + Column("id", String, primary_key=True), + Column("sum_value", Integer), + ], + create_table=True, + ) + output_v2_dt = ds.create_table("output_v2", output_v2_store) + + def transform_func(df_a, df_b): + # Объединяем по id + merged = df_a.merge(df_b, on="id", how="outer") + merged["sum_value"] = merged["value_a"].fillna(0) + merged["value_b"].fillna(0) + return merged[["id", "sum_value"]].astype({"sum_value": int}) + + # Step v1 (без offset) + step_v1 = BatchTransformStep( + ds=ds, + name="test_transform_v1", + func=transform_func, + input_dts=[ + ComputeInput(dt=table_a_dt, join_type="full"), + ComputeInput(dt=table_b_dt, join_type="full"), + ], + output_dts=[output_v1_dt], + transform_keys=["id"], + use_offset_optimization=False, + ) + + # Step v2 (с offset) + step_v2 = BatchTransformStep( + ds=ds, + name="test_transform_v2", + func=transform_func, + input_dts=[ + ComputeInput(dt=table_a_dt, join_type="full"), + ComputeInput(dt=table_b_dt, join_type="full"), + ], + output_dts=[output_v2_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # Добавляем данные + now = time.time() + table_a_dt.store_chunk( + pd.DataFrame({ + "id": ["1", "2", "3"], + "value_a": [10, 20, 30] + }), + now=now + ) + table_b_dt.store_chunk( + pd.DataFrame({ + "id": ["2", "3", "4"], + "value_b": [100, 200, 300] + }), + now=now + ) + + # Запускаем обе трансформации + step_v1.run_full(ds) + step_v2.run_full(ds) + + # Сравниваем результаты + result_v1 = output_v1_dt.get_data().sort_values("id").reset_index(drop=True) + result_v2 = output_v2_dt.get_data().sort_values("id").reset_index(drop=True) + + # КЛЮЧЕВАЯ ПРОВЕРКА: результаты v1 и v2 должны быть идентичны + pd.testing.assert_frame_equal(result_v1, result_v2) + + # Проверяем корректность результатов + expected = pd.DataFrame({ + "id": ["1", "2", "3", "4"], + "sum_value": [10, 120, 230, 300] + }) + pd.testing.assert_frame_equal(result_v1, expected) + + # === Инкрементальная обработка === + # Добавляем новые данные + time.sleep(0.01) + now2 = time.time() + table_a_dt.store_chunk( + pd.DataFrame({ + "id": ["5"], + "value_a": [50] + }), + now=now2 + ) + + # Запускаем снова + step_v1.run_full(ds) + step_v2.run_full(ds) + + # Снова сравниваем + result_v1 = output_v1_dt.get_data().sort_values("id").reset_index(drop=True) + result_v2 = output_v2_dt.get_data().sort_values("id").reset_index(drop=True) + + pd.testing.assert_frame_equal(result_v1, result_v2) + + # Проверяем что добавилась новая запись + assert len(result_v1) == 5 + assert "5" in result_v1["id"].values + + +def test_three_tables_filtered_join(dbconn: DBConn): + """ + Тест 4: Проверяет работу filtered join с ТРЕМЯ таблицами. + + Сценарий: + - Основная таблица posts с post_id, user_id, category_id + - Справочник users с user_id + - Справочник categories с category_id + - join_keys для обоих справочников + - Проверяем что все три таблицы корректно джойнятся и filtered join работает + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Основная таблица: посты + posts_store = TableStoreDB( + dbconn, + "posts", + [ + Column("id", String, primary_key=True), + Column("user_id", String), + Column("category_id", String), + Column("content", String), + ], + create_table=True, + ) + posts_dt = ds.create_table("posts", posts_store) + + # Справочник 1: пользователи + users_store = TableStoreDB( + dbconn, + "users", + [ + Column("id", String, primary_key=True), + Column("username", String), + ], + create_table=True, + ) + users_dt = ds.create_table("users", users_store) + + # Справочник 2: категории + categories_store = TableStoreDB( + dbconn, + "categories", + [ + Column("id", String, primary_key=True), + Column("category_name", String), + ], + create_table=True, + ) + categories_dt = ds.create_table("categories", categories_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "enriched_posts", + [ + Column("id", String, primary_key=True), + Column("content", String), + Column("username", String), + Column("category_name", String), + ], + create_table=True, + ) + output_dt = ds.create_table("enriched_posts", output_store) + + # Отслеживание вызовов get_data для проверки filtered join + users_get_data_calls = [] + categories_get_data_calls = [] + + original_users_get_data = users_dt.get_data + original_categories_get_data = categories_dt.get_data + + def tracked_users_get_data(idx=None, **kwargs): + if idx is not None: + users_get_data_calls.append({ + "idx_columns": list(idx.columns), + "idx_values": sorted(idx["id"].tolist()) if "id" in idx.columns else [], + "idx_length": len(idx) + }) + return original_users_get_data(idx=idx, **kwargs) + + def tracked_categories_get_data(idx=None, **kwargs): + if idx is not None: + categories_get_data_calls.append({ + "idx_columns": list(idx.columns), + "idx_values": sorted(idx["id"].tolist()) if "id" in idx.columns else [], + "idx_length": len(idx) + }) + return original_categories_get_data(idx=idx, **kwargs) + + users_dt.get_data = tracked_users_get_data # type: ignore[method-assign] + categories_dt.get_data = tracked_categories_get_data # type: ignore[method-assign] + + # Функция трансформации: join всех трёх таблиц + def transform_func(posts_df, users_df, categories_df): + # Join с users + result = posts_df.merge( + users_df, + left_on="user_id", + right_on="id", + how="left", + suffixes=("", "_user") + ) + # Join с categories + result = result.merge( + categories_df, + left_on="category_id", + right_on="id", + how="left", + suffixes=("", "_cat") + ) + return result[["id", "content", "username", "category_name"]] + + # Step с двумя filtered joins + # join_keys работает только с v2 (offset optimization) + step = BatchTransformStep( + ds=ds, + name="test_three_tables", + func=transform_func, + input_dts=[ + ComputeInput(dt=posts_dt, join_type="full"), # Основная таблица + ComputeInput( + dt=users_dt, + join_type="full", + join_keys={"user_id": "id"} # Filtered join: posts.user_id -> users.id + ), + ComputeInput( + dt=categories_dt, + join_type="full", + join_keys={"category_id": "id"} # Filtered join: posts.category_id -> categories.id + ), + ], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, # join_keys требует v2 + ) + + # Добавляем данные + now = time.time() + + # 3 поста от 2 пользователей в 2 категориях + posts_dt.store_chunk( + pd.DataFrame({ + "id": ["p1", "p2", "p3"], + "user_id": ["u1", "u2", "u1"], + "category_id": ["cat1", "cat2", "cat1"], + "content": ["Post 1", "Post 2", "Post 3"] + }), + now=now + ) + + # 100 пользователей (но только u1, u2 используются в постах) + users_data = [{"id": f"u{i}", "username": f"User{i}"} for i in range(3, 100)] + users_data.extend([ + {"id": "u1", "username": "Alice"}, + {"id": "u2", "username": "Bob"}, + ]) + users_dt.store_chunk(pd.DataFrame(users_data), now=now) + + # 50 категорий (но только cat1, cat2 используются в постах) + categories_data = [{"id": f"cat{i}", "category_name": f"Category {i}"} for i in range(3, 50)] + categories_data.extend([ + {"id": "cat1", "category_name": "Tech"}, + {"id": "cat2", "category_name": "News"}, + ]) + categories_dt.store_chunk(pd.DataFrame(categories_data), now=now) + + # Запускаем трансформацию + step.run_full(ds) + + # ПРОВЕРКА 1: Filtered join для users должен читать только u1, u2 + assert len(users_get_data_calls) > 0, "users get_data should be called with filtered idx" + users_call = users_get_data_calls[0] + assert "id" in users_call["idx_columns"] + assert users_call["idx_length"] == 2, f"Should filter to 2 users, got {users_call['idx_length']}" + assert sorted(users_call["idx_values"]) == ["u1", "u2"], ( + f"Should filter users to u1, u2, got {users_call['idx_values']}" + ) + + # ПРОВЕРКА 2: Filtered join для categories должен читать только cat1, cat2 + assert len(categories_get_data_calls) > 0, "categories get_data should be called with filtered idx" + categories_call = categories_get_data_calls[0] + assert "id" in categories_call["idx_columns"] + assert categories_call["idx_length"] == 2, ( + f"Should filter to 2 categories, got {categories_call['idx_length']}" + ) + assert sorted(categories_call["idx_values"]) == ["cat1", "cat2"], ( + f"Should filter categories to cat1, cat2, got {categories_call['idx_values']}" + ) + + # ПРОВЕРКА 3: Результат должен содержать правильные данные + output_data = output_dt.get_data().sort_values("id").reset_index(drop=True) + assert len(output_data) == 3 + + expected = pd.DataFrame({ + "id": ["p1", "p2", "p3"], + "content": ["Post 1", "Post 2", "Post 3"], + "username": ["Alice", "Bob", "Alice"], + "category_name": ["Tech", "News", "Tech"] + }) + pd.testing.assert_frame_equal(output_data, expected) + + print("\n✅ Three tables filtered join test passed!") + print(f" - Posts: 3 records") + print(f" - Users: filtered from {len(users_data)} to 2 records (u1, u2)") + print(f" - Categories: filtered from {len(categories_data)} to 2 records (cat1, cat2)") + print(f" - Output: 3 enriched posts with correct joins") + + # === ДОПОЛНИТЕЛЬНАЯ ПРОВЕРКА: сравнение v1 vs v2 === + # Создаём отдельную выходную таблицу для v1 + output_v1_store = TableStoreDB( + dbconn, + "enriched_posts_v1", + [ + Column("id", String, primary_key=True), + Column("content", String), + Column("username", String), + Column("category_name", String), + ], + create_table=True, + ) + output_v1_dt = ds.create_table("enriched_posts_v1", output_v1_store) + + # Step v1 БЕЗ join_keys (обычный FULL OUTER JOIN всех таблиц) + step_v1 = BatchTransformStep( + ds=ds, + name="test_three_tables_v1", + func=transform_func, + input_dts=[ + ComputeInput(dt=posts_dt, join_type="full"), + ComputeInput(dt=users_dt, join_type="full"), # БЕЗ join_keys + ComputeInput(dt=categories_dt, join_type="full"), # БЕЗ join_keys + ], + output_dts=[output_v1_dt], + transform_keys=["id"], + use_offset_optimization=False, # v1 + ) + + # Запускаем v1 + step_v1.run_full(ds) + + # Сравниваем результаты v1 и v2 + result_v1 = output_v1_dt.get_data().sort_values("id").reset_index(drop=True) + result_v2 = output_data # Уже отсортирован + + # КЛЮЧЕВАЯ ПРОВЕРКА: результаты v1 и v2 должны быть идентичны + pd.testing.assert_frame_equal(result_v1, result_v2) + + print("\n✅ V1 vs V2 comparison PASSED!") + print(f" - V1 (FULL OUTER JOIN): {len(result_v1)} rows") + print(f" - V2 (offset + filtered join): {len(result_v2)} rows") + print(f" - Results are identical ✓") diff --git a/tests/test_offset_auto_update.py b/tests/test_offset_auto_update.py new file mode 100644 index 00000000..5c7f0cd7 --- /dev/null +++ b/tests/test_offset_auto_update.py @@ -0,0 +1,408 @@ +""" +Интеграционный тест для Phase 3: Автоматическое обновление offset'ов + +Проверяет что offset'ы автоматически обновляются после успешной обработки батча +и используются при последующих запусках для фильтрации уже обработанных данных. +""" +import time + +import pandas as pd +import sqlalchemy as sa +from sqlalchemy import Column, Integer, String + +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_auto_update_integration(dbconn: DBConn): + """ + Интеграционный тест: Полный цикл работы с автоматическим обновлением offset'ов + + Сценарий: + 1. Создаем входные данные (3 записи) + 2. Запускаем трансформацию с offset optimization + 3. Проверяем что все обработалось + 4. Проверяем что offset'ы обновились в базе + 5. Добавляем новые данные (2 записи) + 6. Запускаем трансформацию снова + 7. Проверяем что обработались только новые записи (благодаря offset'ам) + 8. Проверяем что offset'ы снова обновились + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Создаем входную таблицу + input_store = TableStoreDB( + dbconn, + "integration_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("integration_input", input_store) + + # Создаем выходную таблицу + output_store = TableStoreDB( + dbconn, + "integration_output", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("integration_output", output_store) + + # Трансформация: просто копируем данные с переименованием колонки + def transform_func(df): + return df.rename(columns={"value": "result"}) + + # Создаем step с ВКЛЮЧЕННОЙ offset оптимизацией + step = BatchTransformStep( + ds=ds, + name="integration_transform", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # ========== Шаг 1: Добавляем первую партию данных ========== + now1 = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["A", "B", "C"], "value": [10, 20, 30]}), now=now1 + ) + + # Проверяем что offset'ов ещё нет + offsets_before = ds.offset_table.get_offsets_for_transformation( + step.get_name() + ) + assert len(offsets_before) == 0, "Offset'ы должны быть пустыми до первого запуска" + + # ========== Шаг 2: Первый запуск трансформации ========== + step.run_full(ds) + + # Проверяем что все данные обработались + output_data = output_dt.get_data() + assert len(output_data) == 3 + assert sorted(output_data["id"].tolist()) == ["A", "B", "C"] + assert sorted(output_data["result"].tolist()) == [10, 20, 30] + + # ========== Шаг 3: Проверяем что offset обновился ========== + offsets_after_first_run = ds.offset_table.get_offsets_for_transformation( + step.get_name() + ) + assert len(offsets_after_first_run) == 1, "Должен быть offset для одной входной таблицы" + assert "integration_input" in offsets_after_first_run + + first_offset = offsets_after_first_run["integration_input"] + assert first_offset >= now1, f"Offset ({first_offset}) должен быть >= времени создания данных ({now1})" + assert first_offset <= time.time(), "Offset не должен быть из будущего" + + # ========== Шаг 4: Проверяем что повторный запуск ничего не обработает ========== + changed_count_before_new_data = step.get_changed_idx_count(ds) + assert changed_count_before_new_data == 0, "Не должно быть изменений после первого запуска" + + # ========== Шаг 5: Добавляем новые данные ========== + time.sleep(0.01) # Небольшая задержка чтобы update_ts был больше + now2 = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["D", "E"], "value": [40, 50]}), now=now2 + ) + + # ========== Шаг 6: Проверяем что видны только новые записи ========== + changed_count_after_new_data = step.get_changed_idx_count(ds) + assert changed_count_after_new_data == 2, "Должно быть видно только 2 новые записи (D, E)" + + # ========== Шаг 7: Второй запуск трансформации ========== + step.run_full(ds) + + # Проверяем что теперь всего 5 записей + output_data = output_dt.get_data() + assert len(output_data) == 5 + assert sorted(output_data["id"].tolist()) == ["A", "B", "C", "D", "E"] + assert sorted(output_data["result"].tolist()) == [10, 20, 30, 40, 50] + + # ========== Шаг 8: Проверяем что offset снова обновился ========== + offsets_after_second_run = ds.offset_table.get_offsets_for_transformation( + step.get_name() + ) + assert len(offsets_after_second_run) == 1 + + second_offset = offsets_after_second_run["integration_input"] + assert second_offset > first_offset, f"Второй offset ({second_offset}) должен быть больше первого ({first_offset})" + assert second_offset >= now2, f"Offset ({second_offset}) должен быть >= времени создания новых данных ({now2})" + + # ========== Шаг 9: Финальная проверка - нет новых изменений ========== + final_changed_count = step.get_changed_idx_count(ds) + assert final_changed_count == 0, "После второго запуска не должно быть изменений" + + +def test_offset_update_with_multiple_inputs(dbconn: DBConn): + """ + Тест автоматического обновления offset'ов для трансформации с несколькими входами + + Проверяет что offset'ы обновляются независимо для каждой входной таблицы + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Первая входная таблица + input1_store = TableStoreDB( + dbconn, + "multi_input1", + [Column("id", String, primary_key=True), Column("value1", Integer)], + create_table=True, + ) + input1_dt = ds.create_table("multi_input1", input1_store) + + # Вторая входная таблица + input2_store = TableStoreDB( + dbconn, + "multi_input2", + [Column("id", String, primary_key=True), Column("value2", Integer)], + create_table=True, + ) + input2_dt = ds.create_table("multi_input2", input2_store) + + # Выходная таблица + output_store = TableStoreDB( + dbconn, + "multi_output", + [Column("id", String, primary_key=True), Column("sum", Integer)], + create_table=True, + ) + output_dt = ds.create_table("multi_output", output_store) + + # Трансформация: суммируем значения из обеих таблиц + def transform_func(df1, df2): + merged = df1.merge(df2, on="id", how="outer") + merged["sum"] = merged["value1"].fillna(0) + merged["value2"].fillna(0) + return merged[["id", "sum"]].astype({"sum": int}) + + step = BatchTransformStep( + ds=ds, + name="multi_input_transform", + func=transform_func, + input_dts=[ + ComputeInput(dt=input1_dt, join_type="full"), + ComputeInput(dt=input2_dt, join_type="full"), + ], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # ========== Добавляем данные в обе таблицы ========== + now1 = time.time() + input1_dt.store_chunk(pd.DataFrame({"id": ["X", "Y"], "value1": [1, 2]}), now=now1) + + time.sleep(0.01) + now2 = time.time() + input2_dt.store_chunk(pd.DataFrame({"id": ["X", "Y"], "value2": [10, 20]}), now=now2) + + # ========== Первый запуск ========== + step.run_full(ds) + + output_data = output_dt.get_data() + assert len(output_data) == 2 + assert sorted(output_data["id"].tolist()) == ["X", "Y"] + assert sorted(output_data["sum"].tolist()) == [11, 22] + + # ========== Проверяем offset'ы для обеих таблиц ========== + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert len(offsets) == 2, "Должно быть 2 offset'а (по одному на каждую входную таблицу)" + assert "multi_input1" in offsets + assert "multi_input2" in offsets + + offset1 = offsets["multi_input1"] + offset2 = offsets["multi_input2"] + assert offset1 >= now1 + assert offset2 >= now2 + + # ========== Добавляем данные только в первую таблицу ========== + time.sleep(0.01) + now3 = time.time() + input1_dt.store_chunk(pd.DataFrame({"id": ["Z"], "value1": [3]}), now=now3) + + # ========== Второй запуск ========== + step.run_full(ds) + + # Проверяем результат + output_data = output_dt.get_data() + assert len(output_data) == 3 + assert "Z" in output_data["id"].tolist() + + # ========== Проверяем что offset первой таблицы обновился, а второй - нет ========== + offsets_after = ds.offset_table.get_offsets_for_transformation(step.get_name()) + new_offset1 = offsets_after["multi_input1"] + new_offset2 = offsets_after["multi_input2"] + + assert new_offset1 > offset1, "Offset первой таблицы должен обновиться" + assert new_offset1 >= now3 + assert new_offset2 == offset2, "Offset второй таблицы не должен измениться (не было новых данных)" + + +def test_offset_updated_even_when_optimization_disabled(dbconn: DBConn): + """ + Тест что offset'ы обновляются даже когда use_offset_optimization=False + + Это позволяет накапливать offset'ы для будущего использования, + чтобы при включении оптимизации они были готовы к работе. + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "disabled_opt_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("disabled_opt_input", input_store) + + output_store = TableStoreDB( + dbconn, + "disabled_opt_output", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("disabled_opt_output", output_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + # Создаем step с ВЫКЛЮЧЕННОЙ offset оптимизацией + step = BatchTransformStep( + ds=ds, + name="disabled_opt_transform", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=False, # ВЫКЛЮЧЕНО! + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2"], "value": [10, 20]}), now=now + ) + + # Запускаем трансформацию (используется v1 метод, но offset'ы должны обновиться) + step.run_full(ds) + + # Проверяем что данные обработались + output_data = output_dt.get_data() + assert len(output_data) == 2 + + # ГЛАВНАЯ ПРОВЕРКА: offset'ы должны обновиться даже при выключенной оптимизации + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert len(offsets) == 1, "Offset должен обновиться даже при use_offset_optimization=False" + assert "disabled_opt_input" in offsets + assert offsets["disabled_opt_input"] >= now + + +def test_works_without_offset_table(dbconn: DBConn): + """ + Тест что код работает корректно если таблица offset'ов не создана. + + Симулируем ситуацию: сначала создаем таблицу offset'ов, затем удаляем её, + и проверяем что код не падает при попытке использовать v2 метод и обновить offset'ы. + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "no_offset_table_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("no_offset_table_input", input_store) + + output_store = TableStoreDB( + dbconn, + "no_offset_table_output", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("no_offset_table_output", output_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + step = BatchTransformStep( + ds=ds, + name="no_offset_table_transform", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["1", "2"], "value": [10, 20]}), now=now) + + # УДАЛЯЕМ таблицу offset'ов, симулируя ситуацию когда она не создана + with dbconn.con.begin() as con: + con.execute(sa.text("DROP TABLE IF EXISTS transform_input_offsets")) + + # Запускаем трансформацию - должна работать без ошибок + # v2 метод должен обработать все данные (get_offsets_for_transformation вернет {}) + # Обновление offset'ов должно залогировать warning и продолжить работу + step.run_full(ds) + + # Проверяем что данные обработались несмотря на отсутствие offset таблицы + output_data = output_dt.get_data() + assert len(output_data) == 2 + assert sorted(output_data["id"].tolist()) == ["1", "2"] + + +def test_offset_not_updated_on_empty_result(dbconn: DBConn): + """ + Тест что offset НЕ обновляется если трансформация вернула пустой результат + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "empty_result_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("empty_result_input", input_store) + + output_store = TableStoreDB( + dbconn, + "empty_result_output", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("empty_result_output", output_store) + + # Трансформация которая всегда возвращает пустой DataFrame + def transform_func(df): + return pd.DataFrame(columns=["id", "result"]) + + step = BatchTransformStep( + ds=ds, + name="empty_result_transform", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["1"], "value": [100]}), now=now) + + # Запускаем трансформацию + step.run_full(ds) + + # Проверяем что output пустой + output_data = output_dt.get_data() + assert len(output_data) == 0 + + # Проверяем что offset НЕ обновился (потому что результат пустой) + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert len(offsets) == 0, "Offset не должен обновиться при пустом результате" diff --git a/tests/test_offset_hypotheses.py b/tests/test_offset_hypotheses.py new file mode 100644 index 00000000..b8b311b6 --- /dev/null +++ b/tests/test_offset_hypotheses.py @@ -0,0 +1,525 @@ +""" +Раздельные тесты для гипотез 1 и 2. + +ВАЖНО: Эти тесты независимы и проверяют РАЗНЫЕ проблемы! + +Гипотеза 1: Строгое неравенство update_ts > offset +Гипотеза 2: ORDER BY по transform_keys, а не по update_ts +""" +import time + +import pandas as pd +import pytest +from sqlalchemy import Column, Integer, String + +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_hypothesis_1_strict_inequality_loses_records_with_equal_update_ts(dbconn: DBConn): + """ + Тест ТОЛЬКО для гипотезы 1: Строгое неравенство update_ts > offset. + + Сценарий: + - ВСЕ записи имеют ОДИНАКОВЫЙ update_ts + - Записи сортируются по id (это не важно для этого теста) + - Первый батч: обрабатываем 5 записей с update_ts=T1 + - offset = MAX(T1) = T1 + - Второй запуск: WHERE update_ts > T1 пропускает записи с update_ts == T1 + + Результат: Все необработанные записи с update_ts == offset ПОТЕРЯНЫ! + + === КАК ЭТО МОЖЕТ ПРОИЗОЙТИ В PRODUCTION === + + 1. **Bulk insert / Batch processing**: + - Приложение получает пакет данных (например, 1000 записей из внешнего API) + - Все записи вставляются одним вызовом store_chunk(df, now=current_time) + - Результат: 1000 записей с ОДИНАКОВЫМ update_ts + + 2. **Миграция данных**: + - Перенос исторических данных из старой системы + - Данные импортируются пакетами с одним timestamp + - Результат: Тысячи записей с одинаковым update_ts + + 3. **Реальный production кейс (из hashtag_issue.md)**: + - Трансформация extract_hashtags создала записи пакетами + - Каждый пост может иметь несколько хештегов → несколько записей с одним update_ts + - Пример: пост с 5 хештегами → 5 записей с одинаковым update_ts + - При chunk_size=10 часть записей попадает в первый батч, часть остается + - offset устанавливается на update_ts первого батча + - Оставшиеся записи с тем же update_ts ТЕРЯЮТСЯ! + + 4. **High-load scenario**: + - При высокой нагрузке записи могут создаваться очень быстро + - Точность timestamp может быть до секунды или миллисекунды + - В рамках одной миллисекунды может быть создано 10-100+ записей + - Результат: Множество записей с одинаковым update_ts + + Этот тест должен ПРОЙТИ при исправлении: + - ✅ update_ts > offset → update_ts >= offset + + Этот тест НЕ должен зависеть от: + - ❌ Порядка сортировки (ORDER BY id vs ORDER BY update_ts) + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "hyp1_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("hyp1_input", input_store) + + output_store = TableStoreDB( + dbconn, + "hyp1_output", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + output_dt = ds.create_table("hyp1_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="hyp1_copy", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Создаем первую партию записей с ОДИНАКОВЫМ update_ts + # Симулируем bulk insert или batch processing + base_time = time.time() + same_timestamp = base_time + 1 + + first_batch_df = pd.DataFrame({ + "id": [f"rec_{i:02d}" for i in range(7)], + "value": list(range(7)), + }) + + # Одним вызовом store_chunk - как в production при bulk insert + input_dt.store_chunk(first_batch_df, now=same_timestamp) + time.sleep(0.001) + + print(f"\n=== ПОДГОТОВКА ===") + print(f"Создано {len(first_batch_df)} записей с update_ts = {same_timestamp:.2f}") + print("(Симуляция bulk insert или batch processing)") + + # ПЕРВЫЙ ЗАПУСК: обрабатываем ВСЕ записи через run_full + print(f"\n=== ПЕРВЫЙ ЗАПУСК (run_full) ===") + step.run_full(ds) + + # Проверяем offset + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_after_first = offsets["hyp1_input"] + + output_after_first = output_dt.get_data() + processed_ids_first = set(output_after_first["id"].tolist()) + + print(f"Обработано: {len(output_after_first)} записей") + print(f"offset = {offset_after_first:.2f}") + print(f"Обработанные id: {sorted(processed_ids_first)}") + + assert len(processed_ids_first) == 7, "Должно быть обработано 7 записей" + + # ДОБАВЛЯЕМ НОВЫЕ ЗАПИСИ с ТЕМ ЖЕ update_ts + # Это критический момент: новые записи имеют update_ts == offset + time.sleep(0.001) + second_batch_df = pd.DataFrame({ + "id": [f"rec_{i:02d}" for i in range(7, 12)], + "value": list(range(7, 12)), + }) + input_dt.store_chunk(second_batch_df, now=same_timestamp) # ТОТ ЖЕ timestamp! + + print(f"\n=== ДОБАВЛЕНЫ НОВЫЕ ЗАПИСИ ===") + print(f"Добавлено {len(second_batch_df)} записей с update_ts = {same_timestamp:.2f}") + print(f"update_ts == offset ({same_timestamp:.2f} == {offset_after_first:.2f})") + + # ВТОРОЙ ЗАПУСК: проверяем что записи с update_ts == offset обработаются + print(f"\n=== ВТОРОЙ ЗАПУСК (run_full) ===") + step.run_full(ds) + + # ПРОВЕРКА: НОВЫЕ записи с update_ts == offset должны быть обработаны + final_output = output_dt.get_data() + final_processed_ids = set(final_output["id"].tolist()) + + all_expected_ids = set([f"rec_{i:02d}" for i in range(12)]) + lost_records = all_expected_ids - final_processed_ids + + if lost_records: + print(f"\n=== 🚨 ПОТЕРЯННЫЕ ЗАПИСИ (БАГ!) ===") + print(f"Потерянные id: {sorted(lost_records)}") + print(f"Все они имеют update_ts == offset ({same_timestamp:.2f})") + + pytest.fail( + f"ГИПОТЕЗА 1 ПОДТВЕРЖДЕНА: {len(lost_records)} записей с update_ts == offset ПОТЕРЯНЫ!\n" + f"Ожидалось: {len(all_expected_ids)} записей\n" + f"Получено: {len(final_output)} записей\n" + f"Потеряно: {len(lost_records)} записей\n" + f"Потерянные id: {sorted(lost_records)}\n\n" + f"Причина: Строгое неравенство 'update_ts > offset' пропускает записи с update_ts == offset\n" + f"Исправление: datapipe/meta/sql_meta.py - заменить '>' на '>='" + ) + + print(f"\n=== ✅ ВСЕ ЗАПИСИ ОБРАБОТАНЫ ===") + print(f"Всего записей: {len(all_expected_ids)}") + print(f"Обработано: {len(final_output)}") + + +def test_hypothesis_2_order_by_transform_keys_with_mixed_update_ts(dbconn: DBConn): + """ + Тест ТОЛЬКО для гипотезы 2: ORDER BY по transform_keys, а не по update_ts. + + Сценарий: + - Записи имеют РАЗНЫЕ update_ts + - Записи сортируются по id (transform_keys), НЕ по update_ts + - В батч попадают записи с разными update_ts (например: T1, T1, T3, T3, T3) + - offset = MAX(T1, T1, T3, T3, T3) = T3 + - Но есть запись с id ПОСЛЕ последней обработанной, но с update_ts < T3 + - Второй запуск: WHERE update_ts > T3 пропускает эту запись + + ВАЖНО: Этот тест должен ПАДАТЬ даже при исправлении гипотезы 1 (> на >=)! + Для этого мы НЕ должны иметь записей с update_ts == offset в необработанных данных. + + Этот тест должен ПРОЙТИ при исправлении: + - ✅ ORDER BY transform_keys → ORDER BY update_ts + - ИЛИ другой способ обеспечить что offset не превышает MAX(update_ts обработанных записей) + + Этот тест НЕ должен пройти при исправлении: + - ❌ update_ts > offset → update_ts >= offset (гипотеза 1) + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "hyp2_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("hyp2_input", input_store) + + output_store = TableStoreDB( + dbconn, + "hyp2_output", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + output_dt = ds.create_table("hyp2_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="hyp2_copy", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Создаем записи с РАЗНЫМИ update_ts в "неправильном" порядке id + base_time = time.time() + + # Группа 1: T1 - ранний timestamp + t1 = base_time + 1 + input_dt.store_chunk( + pd.DataFrame({"id": ["rec_00", "rec_01"], "value": [0, 1]}), + now=t1 + ) + time.sleep(0.001) + + # Группа 2: T3 - ПОЗДНИЙ timestamp (специально создаем "дыру") + t3 = base_time + 3 + input_dt.store_chunk( + pd.DataFrame({"id": ["rec_02", "rec_03", "rec_04"], "value": [2, 3, 4]}), + now=t3 + ) + time.sleep(0.001) + + # Группа 3: T2 - СРЕДНИЙ timestamp (но id ПОСЛЕ первого батча) + t2 = base_time + 2 + input_dt.store_chunk( + pd.DataFrame({"id": ["rec_05", "rec_06", "rec_07"], "value": [5, 6, 7]}), + now=t2 # Старый timestamp, но id ПОСЛЕ rec_04! + ) + time.sleep(0.001) + + # Группа 4: T4 - Еще более поздний timestamp + t4 = base_time + 4 + input_dt.store_chunk( + pd.DataFrame({"id": ["rec_08", "rec_09", "rec_10"], "value": [8, 9, 10]}), + now=t4 + ) + + # Проверяем данные + all_meta = input_dt.meta_table.get_metadata() + print(f"\n=== ПОДГОТОВКА ===") + print(f"Всего записей: {len(all_meta)}") + print("Распределение по update_ts (сортировка по id):") + for idx, row in all_meta.sort_values("id").iterrows(): + ts_label = "T1" if abs(row["update_ts"] - t1) < 0.01 else \ + "T2" if abs(row["update_ts"] - t2) < 0.01 else \ + "T3" if abs(row["update_ts"] - t3) < 0.01 else "T4" + print(f" id={row['id']:10} update_ts={ts_label} ({row['update_ts']:.2f})") + + # ПЕРВЫЙ ЗАПУСК: обрабатываем ВСЕ записи через run_full + print(f"\n=== ПЕРВЫЙ ЗАПУСК (run_full) ===") + step.run_full(ds) + + # Проверяем offset + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_after_first = offsets["hyp2_input"] + + output_after_first = output_dt.get_data() + processed_ids = set(output_after_first["id"].tolist()) + + print(f"Обработано: {len(output_after_first)} записей") + print(f"offset = {offset_after_first:.2f} (должно быть T4 = {t4:.2f})") + print(f"Обработанные id: {sorted(processed_ids)}") + + all_input_ids = set(all_meta["id"].tolist()) + assert processed_ids == all_input_ids, "Все исходные записи должны быть обработаны" + + # ДОБАВЛЯЕМ НОВЫЕ ЗАПИСИ с update_ts БОЛЬШЕ offset + # Критический тест: даже если id меньше последних обработанных, + # записи должны обработаться благодаря сортировке по update_ts + time.sleep(0.001) + t_new = t4 + 1 # Timestamp больше всех предыдущих + + new_records_df = pd.DataFrame({ + "id": ["rec_new_1", "rec_new_2"], # Префикс "rec_new" сортируется после "rec_10" + "value": [11, 12], + }) + input_dt.store_chunk(new_records_df, now=t_new) + + print(f"\n=== ДОБАВЛЕНЫ НОВЫЕ ЗАПИСИ ===") + print(f"Добавлено {len(new_records_df)} записей с update_ts = {t_new:.2f}") + print(f"update_ts ({t_new:.2f}) > offset ({offset_after_first:.2f})") + print(f"Новые id: {sorted(new_records_df['id'].tolist())}") + print("С правильной сортировкой по update_ts эти записи будут обработаны") + + # ВТОРОЙ ЗАПУСК: с учетом offset и сортировки по update_ts + print(f"\n=== ВТОРОЙ ЗАПУСК (run_full) ===") + step.run_full(ds) + + # ПРОВЕРКА: НОВЫЕ записи с update_ts > offset должны быть обработаны + final_output = output_dt.get_data() + final_processed_ids = set(final_output["id"].tolist()) + + all_expected_ids = all_input_ids | set(new_records_df["id"].tolist()) + lost_records = all_expected_ids - final_processed_ids + + if lost_records: + all_meta_final = input_dt.meta_table.get_metadata() + lost_meta = all_meta_final[all_meta_final["id"].isin(lost_records)] + print(f"\n=== 🚨 ПОТЕРЯННЫЕ ЗАПИСИ (БАГ!) ===") + for idx, row in lost_meta.sort_values("id").iterrows(): + print( + f" id={row['id']:10} update_ts={row['update_ts']:.2f} " + f"> offset={offset_after_first:.2f} (но все равно пропущены!)" + ) + + pytest.fail( + f"ГИПОТЕЗА 2 ПОДТВЕРЖДЕНА: {len(lost_records)} записей ПОТЕРЯНЫ из-за ORDER BY по transform_keys!\n" + f"Ожидалось: {len(all_input_ids)} записей\n" + f"Получено: {len(final_output)} записей\n" + f"Потеряно: {len(lost_records)} записей\n" + f"Потерянные id: {sorted(lost_records)}\n\n" + f"Причина: Батчи сортируются ORDER BY transform_keys (id), но offset = MAX(update_ts).\n" + f" Записи с id ПОСЛЕ последней обработанной, но с update_ts < offset ПОТЕРЯНЫ.\n" + f"Исправление: Либо сортировать по update_ts, либо пересмотреть логику offset." + ) + + print(f"\n=== ✅ ВСЕ ЗАПИСИ ОБРАБОТАНЫ ===") + print(f"Всего записей: {len(all_input_ids)}") + print(f"Обработано: {len(final_output)}") + + +def test_antiregression_no_infinite_loop_with_equal_update_ts(dbconn: DBConn): + """ + Анти-регрессионный тест: Проверяет что после исправления > на >= не возникает зацикливание. + + ВАЖНО: Этот тест должен ПРОХОДИТЬ (не xfail) и после исправления тоже должен проходить! + + Сценарий: + 1. Создаем 12 записей с ОДИНАКОВЫМ update_ts (bulk insert) + 2. Первый запуск: обрабатываем первый батч (5 записей) + - offset = T1 + - Проверяем что обработано ровно 5 записей + 3. Второй запуск: обрабатываем следующий батч (5 записей с update_ts == T1) + - Проверяем что обработано ровно 5 НОВЫХ записей (не те же самые!) + - Проверяем что offset НЕ изменился (всё ещё T1) + 4. Третий запуск: обрабатываем последний батч (2 записи) + - Проверяем что обработано 2 новых записи + 5. Добавляем НОВЫЕ записи с update_ts > T1 + - Проверяем что новые записи будут обработаны + + Критично: + - Каждый запуск должен обрабатывать НОВЫЕ записи, не зацикливаться на одних и тех же + - После исправления >= система должна корректно обрабатывать записи с update_ts == offset + - process_ts должен обновляться для обработанных записей + """ + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "antiregr_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("antiregr_input", input_store) + + output_store = TableStoreDB( + dbconn, + "antiregr_output", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + output_dt = ds.create_table("antiregr_output", output_store) + + def copy_func(df): + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="antiregr_copy", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=5, + ) + + # Создаем первую партию с ОДИНАКОВЫМ update_ts (bulk insert) + base_time = time.time() + same_timestamp = base_time + 1 + + first_batch_df = pd.DataFrame({ + "id": [f"rec_{i:02d}" for i in range(12)], + "value": list(range(12)), + }) + + input_dt.store_chunk(first_batch_df, now=same_timestamp) + time.sleep(0.001) + + print(f"\n=== ПОДГОТОВКА ===") + print(f"Создано {len(first_batch_df)} записей с одинаковым update_ts = {same_timestamp:.2f}") + + # ========== ПЕРВЫЙ ЗАПУСК (run_full): обработать все ========== + print(f"\n=== ПЕРВЫЙ ЗАПУСК (run_full) ===") + step.run_full(ds) + + output_1 = output_dt.get_data() + processed_ids_1 = set(output_1["id"].tolist()) + offsets_1 = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_1 = offsets_1["antiregr_input"] + + print(f"Обработано: {len(output_1)} записей") + print(f"offset = {offset_1:.2f}") + print(f"Обработанные id: {sorted(processed_ids_1)}") + + assert len(output_1) == 12, f"Ожидалось 12 записей, получено {len(output_1)}" + first_batch_offset = offset_1 + + # ========== ДОБАВЛЯЕМ НОВЫЕ ЗАПИСИ с ТЕМ ЖЕ update_ts == offset ========== + # Критический момент: проверяем что нет зацикливания на записях с update_ts == offset + time.sleep(0.001) + second_batch_df = pd.DataFrame({ + "id": [f"same_{i:02d}" for i in range(5)], + "value": list(range(50, 55)), + }) + input_dt.store_chunk(second_batch_df, now=same_timestamp) # ТОТ ЖЕ timestamp! + + print(f"\n=== ДОБАВЛЕНЫ НОВЫЕ ЗАПИСИ ===") + print(f"Добавлено {len(second_batch_df)} записей с update_ts = {same_timestamp:.2f}") + print(f"update_ts == offset ({same_timestamp:.2f} == {offset_1:.2f})") + + # ========== ВТОРОЙ ЗАПУСК: проверяем что нет зацикливания ========== + print(f"\n=== ВТОРОЙ ЗАПУСК (run_full) ===") + step.run_full(ds) + + output_2 = output_dt.get_data() + processed_ids_2 = set(output_2["id"].tolist()) + new_ids_2 = processed_ids_2 - processed_ids_1 + offsets_2 = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_2 = offsets_2["antiregr_input"] + + print(f"Всего обработано: {len(output_2)} записей") + print(f"Новых записей: {len(new_ids_2)}") + print(f"Новые id: {sorted(new_ids_2)}") + print(f"offset = {offset_2:.2f}") + + # Критичная проверка: должны обработать НОВЫЕ записи (с update_ts == offset) + assert len(new_ids_2) == 5, ( + f"Ожидалось 5 НОВЫХ записей с update_ts == offset, получено {len(new_ids_2)}!\n" + f"Возможно баг: записи с update_ts == offset не обработаны." + ) + assert len(output_2) == 17, f"Всего должно быть 17 записей, получено {len(output_2)}" + + # Проверяем что это действительно НОВЫЕ записи, а не зацикливание + intersection = processed_ids_1 & new_ids_2 + assert len(intersection) == 0, ( + f"ЗАЦИКЛИВАНИЕ: Повторно обрабатываем те же записи: {sorted(intersection)}" + ) + assert all(id.startswith("same_") for id in new_ids_2), ( + f"Новые записи должны начинаться с 'same_', получено: {sorted(new_ids_2)}" + ) + + # ========== ДОБАВЛЯЕМ ЗАПИСИ с update_ts > offset ========== + time.sleep(0.01) + new_timestamp = same_timestamp + 1 # Гарантированно больше чем offset + third_batch_df = pd.DataFrame({ + "id": [f"new_{i:02d}" for i in range(5)], + "value": list(range(100, 105)), + }) + input_dt.store_chunk(third_batch_df, now=new_timestamp) + + print(f"\n=== ДОБАВЛЕНЫ НОВЫЕ ЗАПИСИ ===") + print(f"Добавлено {len(third_batch_df)} записей с update_ts = {new_timestamp:.2f} > offset = {offset_2:.2f}") + + # ========== ТРЕТИЙ ЗАПУСК: новые записи с update_ts > offset ========== + print(f"\n=== ТРЕТИЙ ЗАПУСК (run_full) ===") + step.run_full(ds) + + output_3 = output_dt.get_data() + processed_ids_3 = set(output_3["id"].tolist()) + new_ids_3 = processed_ids_3 - processed_ids_2 + offsets_3 = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_3 = offsets_3["antiregr_input"] + + print(f"Всего обработано: {len(output_3)} записей") + print(f"Новых записей: {len(new_ids_3)}") + print(f"Новые id: {sorted(new_ids_3)}") + print(f"offset = {offset_3:.2f}") + + assert len(output_3) == 22, f"Всего должно быть 22 записи (12 + 5 + 5), получено {len(output_3)}" + assert len(new_ids_3) == 5, f"Ожидалось 5 новых записей, получено {len(new_ids_3)}" + assert offset_3 > offset_2, ( + f"offset должен обновиться для записей с update_ts > offset! " + f"Был {offset_2:.2f}, остался {offset_3:.2f}" + ) + + # Проверяем что новые записи действительно новые + assert all(id.startswith("new_") for id in new_ids_3), ( + f"Новые записи должны начинаться с 'new_', получено: {sorted(new_ids_3)}" + ) + + print(f"\n=== ✅ ВСЕ ПРОВЕРКИ ПРОШЛИ ===") + print("1. Нет зацикливания на одних и тех же записях") + print("2. Записи с update_ts == offset обрабатываются корректно (>=)") + print("3. Записи с update_ts > offset также обрабатываются") + print("3. Записи с update_ts == offset корректно обрабатываются") + print("4. Новые записи с update_ts > offset корректно обрабатываются") + print("5. offset корректно обновляется") diff --git a/tests/test_offset_hypothesis_3_multi_step.py b/tests/test_offset_hypothesis_3_multi_step.py new file mode 100644 index 00000000..8504f959 --- /dev/null +++ b/tests/test_offset_hypothesis_3_multi_step.py @@ -0,0 +1,424 @@ +""" +Тест для гипотезы 3: Рассинхронизация update_ts и process_ts в multi-step pipeline. + +ГИПОТЕЗА 3 (из README.md): +"Другая трансформация обновляет process_ts, но НЕ update_ts" + +СЦЕНАРИЙ ИЗ PRODUCTION: +1. Transform_extract_hashtags создает записи в hashtag_table (update_ts=16:21) +2. Transform_hashtag_stats обрабатывает hashtag_table спустя 4 часа (20:04) + - process_ts в Transform_hashtag_stats.meta_table = 20:04 + - update_ts в hashtag_table (input) остается = 16:21 +3. offset(Transform_hashtag_stats, hashtag_table) = 16:21 (MAX update_ts из input) +4. Временной разрыв: update_ts=16:21, process_ts=20:04 + +ВОПРОС: +Влияет ли эта рассинхронизация на offset optimization? +Может ли это вызвать потерю данных? + +АРХИТЕКТУРА: +- У каждой трансформации СВОЯ TransformMetaTable с СВОИМ process_ts +- Transform_A.meta_table хранит process_ts для Transform_A +- Transform_B.meta_table хранит process_ts для Transform_B +- Они не пересекаются! + +ОЖИДАНИЕ: +Рассинхронизация НЕ должна влиять на корректность offset optimization, +так как каждая трансформация работает со своим process_ts. +""" +import time + +import pandas as pd +import pytest +import sqlalchemy as sa +from sqlalchemy import Column, Integer, String + +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_hypothesis_3_multi_step_pipeline_update_ts_vs_process_ts_desync(dbconn: DBConn): + """ + Проверка гипотезы 3: Рассинхронизация update_ts и process_ts в multi-step pipeline. + + Pipeline: + source_table → Transform_A → intermediate_table → Transform_B → final_table + + Сценарий: + 1. В T1 (16:21): Transform_A создает данные в intermediate_table (update_ts=T1) + 2. В T2 (20:04): Transform_B обрабатывает intermediate_table + - process_ts в Transform_B.meta = T2 + - update_ts в intermediate_table остается = T1 + - offset(Transform_B, intermediate_table) = T1 + 3. Добавляем новые данные в source_table + 4. В T3: Transform_A обрабатывает новые данные (update_ts=T3 в intermediate) + 5. В T4: Transform_B обрабатывает новые данные + - Проверяем что offset optimization работает корректно + - Проверяем что старые данные не обрабатываются повторно + """ + ds = DataStore(dbconn, create_meta_table=True) + + # ========== СОЗДАНИЕ ТАБЛИЦ ========== + # Source table + source_store = TableStoreDB( + dbconn, + "hyp3_source", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + source_table = ds.create_table("hyp3_source", source_store) + + # Intermediate table (output Transform_A, input Transform_B) + intermediate_store = TableStoreDB( + dbconn, + "hyp3_intermediate", + [Column("id", String, primary_key=True), Column("value_doubled", Integer)], + create_table=True, + ) + intermediate_table = ds.create_table("hyp3_intermediate", intermediate_store) + + # Final table (output Transform_B) + final_store = TableStoreDB( + dbconn, + "hyp3_final", + [Column("id", String, primary_key=True), Column("value_squared", Integer)], + create_table=True, + ) + final_table = ds.create_table("hyp3_final", final_store) + + # ========== СОЗДАНИЕ ТРАНСФОРМАЦИЙ ========== + def double_func(df): + """Transform_A: удваивает значения""" + return df.assign(value_doubled=df["value"] * 2)[["id", "value_doubled"]] + + transform_a = BatchTransformStep( + ds=ds, + name="transform_a_double", + func=double_func, + input_dts=[ComputeInput(dt=source_table, join_type="full")], + output_dts=[intermediate_table], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=10, + ) + + def square_func(df): + """Transform_B: возводит в квадрат""" + return df.assign(value_squared=df["value_doubled"] ** 2)[["id", "value_squared"]] + + transform_b = BatchTransformStep( + ds=ds, + name="transform_b_square", + func=square_func, + input_dts=[ComputeInput(dt=intermediate_table, join_type="full")], + output_dts=[final_table], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=10, + ) + + # ========== ФАЗА 1: T1 (16:21) - Transform_A создает данные ========== + print("\n" + "=" * 80) + print("ФАЗА 1: T1 (16:21) - Transform_A создает первую партию данных") + print("=" * 80) + + base_time = time.time() + t1 = base_time + 1 # 16:21 в production + + # Загружаем данные в source + source_data_1 = pd.DataFrame({ + "id": [f"rec_{i:02d}" for i in range(5)], + "value": [1, 2, 3, 4, 5], + }) + source_table.store_chunk(source_data_1, now=t1) + time.sleep(0.01) + + # Transform_A обрабатывает source → intermediate + transform_a.run_full(ds=ds, run_config=None) + + # Проверяем результат + intermediate_data = intermediate_table.get_data() + intermediate_meta = intermediate_table.meta_table.get_metadata() + + print(f"\nIntermediate table после Transform_A:") + print(f" Записей: {len(intermediate_data)}") + print(f" update_ts: {intermediate_meta['update_ts'].unique()}") + + # Сохраняем update_ts первой партии + intermediate_update_ts_phase1 = intermediate_meta['update_ts'].iloc[0] + + assert len(intermediate_data) == 5 + assert (intermediate_data["value_doubled"] == [2, 4, 6, 8, 10]).all() + + # ========== ФАЗА 2: T2 (20:04) - Transform_B обрабатывает (С ЗАДЕРЖКОЙ!) ========== + print("\n" + "=" * 80) + print("ФАЗА 2: T2 (20:04) - Transform_B обрабатывает intermediate (4 часа спустя!)") + print("=" * 80) + + # Симулируем задержку 4 часа + time.sleep(0.05) # В тесте - маленькая задержка + t2 = base_time + 2 # 20:04 в production + + # Запоминаем текущее время перед обработкой + before_transform_b = time.time() + + # Transform_B обрабатывает intermediate → final + transform_b.run_full(ds=ds, run_config=None) + + after_transform_b = time.time() + + # Проверяем результат + final_data = final_table.get_data() + + # Получаем данные из Transform_B meta table через SQL + with ds.meta_dbconn.con.begin() as con: + transform_b_meta = pd.read_sql( + sa.select(transform_b.meta_table.sql_table), + con + ) + + # КРИТИЧНО: проверяем update_ts в intermediate table - он НЕ должен измениться! + intermediate_meta_after_b = intermediate_table.meta_table.get_metadata() + + print(f"\nIntermediate table после Transform_B:") + print(f" update_ts: {intermediate_meta_after_b['update_ts'].unique()}") + print(f" (должен остаться = T1, потому что Transform_B читает но НЕ пишет в intermediate)") + + print(f"\nTransform_B meta table:") + print(f" process_ts: {transform_b_meta['process_ts'].unique()}") + print(f" (должен быть ≈ T2, текущее время обработки)") + + # Проверяем offset + offsets_b = ds.offset_table.get_offsets_for_transformation(transform_b.get_name()) + offset_b_intermediate = offsets_b["hyp3_intermediate"] + + print(f"\nOffset для Transform_B:") + print(f" offset(Transform_B, intermediate_table) = {offset_b_intermediate:.2f}") + print(f" (должен быть = MAX(update_ts из intermediate) ≈ T1)") + + print(f"\n🔍 РАССИНХРОНИЗАЦИЯ:") + print(f" update_ts в intermediate = {intermediate_update_ts_phase1:.2f} (T1)") + print(f" process_ts в Transform_B = {transform_b_meta['process_ts'].iloc[0]:.2f} (T2)") + print(f" Разница: {transform_b_meta['process_ts'].iloc[0] - intermediate_update_ts_phase1:.2f} сек") + + # ПРОВЕРКИ + assert len(final_data) == 5, f"Должно быть 5 записей, получено {len(final_data)}" + + # КРИТИЧНО: Проверяем что ВСЕ записи обработаны (ничего не потеряно!) + expected_ids_phase2 = {f"rec_{i:02d}" for i in range(5)} + actual_ids_phase2 = set(final_data["id"].tolist()) + lost_ids_phase2 = expected_ids_phase2 - actual_ids_phase2 + + if lost_ids_phase2: + pytest.fail( + f"🚨 ПОТЕРЯ ДАННЫХ В ФАЗЕ 2!\n" + f"Ожидалось: {len(expected_ids_phase2)} записей\n" + f"Обработано: {len(actual_ids_phase2)} записей\n" + f"ПОТЕРЯНО: {len(lost_ids_phase2)} записей: {sorted(lost_ids_phase2)}\n" + f"Это означает что гипотеза 3 влияет на потерю данных!" + ) + + print(f"\n✓ ВСЕ записи обработаны: {len(actual_ids_phase2)}/{len(expected_ids_phase2)}") + + # Проверяем значения + assert (final_data["value_squared"] == [4, 16, 36, 64, 100]).all() + + # update_ts в intermediate НЕ должен измениться (Transform_B только читает) + assert (intermediate_meta_after_b['update_ts'] == intermediate_update_ts_phase1).all(), \ + "update_ts в intermediate table НЕ должен измениться после Transform_B" + + # process_ts в Transform_B должен быть ≈ текущее время + assert all(before_transform_b <= ts <= after_transform_b + for ts in transform_b_meta['process_ts']), \ + "process_ts в Transform_B должен быть текущим временем обработки" + + # offset должен быть = MAX(update_ts из intermediate) ≈ T1 + assert abs(offset_b_intermediate - intermediate_update_ts_phase1) < 0.1, \ + "offset должен быть равен MAX(update_ts из intermediate)" + + # ========== ФАЗА 3: T3 - Добавляем новые данные через Transform_A ========== + print("\n" + "=" * 80) + print("ФАЗА 3: T3 - Добавляем новые данные и обрабатываем через Transform_A") + print("=" * 80) + + time.sleep(0.01) + t3 = base_time + 3 + + # Добавляем новые данные в source + source_data_2 = pd.DataFrame({ + "id": [f"rec_{i:02d}" for i in range(5, 10)], + "value": [6, 7, 8, 9, 10], + }) + source_table.store_chunk(source_data_2, now=t3) + time.sleep(0.01) + + # Transform_A обрабатывает новые данные + transform_a.run_full(ds=ds, run_config=None) + + intermediate_data_after_new = intermediate_table.get_data() + print(f"\nIntermediate table после добавления новых данных:") + print(f" Записей: {len(intermediate_data_after_new)}") + + assert len(intermediate_data_after_new) == 10 + + # ========== ФАЗА 4: T4 - Transform_B обрабатывает новые данные ========== + print("\n" + "=" * 80) + print("ФАЗА 4: T4 - Transform_B обрабатывает ТОЛЬКО новые данные") + print("=" * 80) + + time.sleep(0.01) + t4 = base_time + 4 + + # Получаем количество записей в transform_b.meta ДО обработки + with ds.meta_dbconn.con.begin() as con: + transform_b_meta_before = pd.read_sql( + sa.select(transform_b.meta_table.sql_table), + con + ) + process_ts_before = dict(zip(transform_b_meta_before['id'], transform_b_meta_before['process_ts'])) + + print(f"\nTransform_B meta ДО обработки новых данных:") + print(f" Записей: {len(transform_b_meta_before)}") + print(f" process_ts для старых записей (rec_00): {process_ts_before.get('rec_00', 'N/A')}") + + # Transform_B обрабатывает новые данные с использованием offset + transform_b.run_full(ds=ds, run_config=None) + + final_data_after_new = final_table.get_data() + with ds.meta_dbconn.con.begin() as con: + transform_b_meta_after = pd.read_sql( + sa.select(transform_b.meta_table.sql_table), + con + ) + process_ts_after = dict(zip(transform_b_meta_after['id'], transform_b_meta_after['process_ts'])) + + print(f"\nTransform_B meta ПОСЛЕ обработки новых данных:") + print(f" Записей: {len(transform_b_meta_after)}") + print(f" process_ts для старых записей (rec_00): {process_ts_after.get('rec_00', 'N/A')}") + print(f" process_ts для новых записей (rec_05): {process_ts_after.get('rec_05', 'N/A')}") + + # КРИТИЧНАЯ ПРОВЕРКА: старые записи НЕ должны обработаться повторно + old_ids = [f"rec_{i:02d}" for i in range(5)] + reprocessed_ids = [] + + for old_id in old_ids: + if old_id in process_ts_before and old_id in process_ts_after: + if abs(process_ts_after[old_id] - process_ts_before[old_id]) > 0.001: + reprocessed_ids.append(old_id) + + if reprocessed_ids: + print(f"\n🚨 ПРОБЛЕМА: Старые записи обработаны ПОВТОРНО!") + print(f" Записи: {reprocessed_ids}") + for rid in reprocessed_ids: + print(f" {rid}: process_ts ДО={process_ts_before[rid]:.6f}, " + f"ПОСЛЕ={process_ts_after[rid]:.6f}") + pytest.fail( + f"ГИПОТЕЗА 3: Рассинхронизация update_ts и process_ts вызвала повторную обработку!\n" + f"Старые записи ({reprocessed_ids}) были обработаны повторно.\n" + f"Это указывает на проблему с offset optimization в multi-step pipeline." + ) + + # КРИТИЧНО: Проверяем что ВСЕ записи обработаны (ничего не потеряно!) + expected_ids_phase4 = {f"rec_{i:02d}" for i in range(10)} # rec_00..rec_09 + actual_ids_phase4 = set(final_data_after_new["id"].tolist()) + lost_ids_phase4 = expected_ids_phase4 - actual_ids_phase4 + + if lost_ids_phase4: + pytest.fail( + f"🚨 ПОТЕРЯ ДАННЫХ В ФАЗЕ 4!\n" + f"Ожидалось: {len(expected_ids_phase4)} записей\n" + f"Обработано: {len(actual_ids_phase4)} записей\n" + f"ПОТЕРЯНО: {len(lost_ids_phase4)} записей: {sorted(lost_ids_phase4)}\n\n" + f"Детали:\n" + f" Старые записи (должны остаться): rec_00..rec_04\n" + f" Новые записи (должны добавиться): rec_05..rec_09\n" + f" Фактически потерянные: {sorted(lost_ids_phase4)}\n\n" + f"Это означает что гипотеза 3 (рассинхронизация update_ts/process_ts)\n" + f"влияет на потерю данных в multi-step pipeline!" + ) + + print(f"\n✓ ВСЕ записи обработаны: {len(actual_ids_phase4)}/{len(expected_ids_phase4)}") + print(f" Старые записи (rec_00..rec_04): {all(f'rec_{i:02d}' in actual_ids_phase4 for i in range(5))}") + print(f" Новые записи (rec_05..rec_09): {all(f'rec_{i:02d}' in actual_ids_phase4 for i in range(5, 10))}") + + # ПРОВЕРКИ + assert len(final_data_after_new) == 10, f"Должно быть 10 записей, получено {len(final_data_after_new)}" + + # Проверяем что все значения корректны + expected_values = [4, 16, 36, 64, 100, 144, 196, 256, 324, 400] # (value*2)^2 + actual_values_sorted = sorted(final_data_after_new["value_squared"].tolist()) + assert actual_values_sorted == expected_values, \ + f"Значения не совпадают. Ожидалось: {expected_values}, получено: {actual_values_sorted}" + + print(f"\n✅ ГИПОТЕЗА 3: Рассинхронизация update_ts и process_ts НЕ вызывает проблем") + print(f"✓ Старые записи НЕ обработаны повторно") + print(f"✓ Новые записи обработаны корректно") + print(f"✓ Все записи обработаны, НИЧЕГО не потеряно") + print(f"✓ offset optimization работает корректно в multi-step pipeline") + print(f"\nОбъяснение:") + print(f" - У каждой трансформации СВОЯ meta table с СВОИМ process_ts") + print(f" - Transform_B использует offset на основе update_ts из intermediate table") + print(f" - process_ts в Transform_B.meta НЕ связан с update_ts в intermediate") + print(f" - Рассинхронизация НЕ влияет на корректность offset optimization") + + +def test_hypothesis_3_explanation(): + """ + Документация и объяснение гипотезы 3. + + ВОПРОС: Почему рассинхронизация update_ts и process_ts не влияет на offset optimization? + + ОТВЕТ: + 1. У каждой трансформации СВОЯ TransformMetaTable с СВОИМ process_ts + 2. offset(Transform_B, TableA) = MAX(update_ts из TableA) + 3. process_ts в Transform_B.meta относится к обработке Transform_B + 4. Они не пересекаются! + + АРХИТЕКТУРА: + + source_table → Transform_A → intermediate_table → Transform_B → final_table + [Meta_A] [Meta_B] + + - Meta_A.process_ts = когда Transform_A обработал записи + - intermediate_table.update_ts = когда Transform_A записал данные + - offset(Transform_B, intermediate) = MAX(intermediate_table.update_ts) + - Meta_B.process_ts = когда Transform_B обработал записи + + ВАЖНО: + - Transform_B НЕ смотрит на Meta_A.process_ts + - Transform_B использует intermediate_table.update_ts для offset + - Рассинхронизация между Meta_A.process_ts и intermediate_table.update_ts не важна + + КОГДА НУЖНА ПРОВЕРКА process_ts: + Проверка process_ts нужна для ОДНОЙ трансформации, чтобы не обработать + одни и те же данные дважды при изменении > на >=. + + Но это проверка СВОЕГО process_ts (Transform_B.meta.process_ts), + а не process_ts других трансформаций! + """ + pass + + +if __name__ == "__main__": + from datapipe.store.database import DBConn + from sqlalchemy import create_engine, text + + DBCONNSTR = "postgresql://postgres:password@localhost:5432/postgres" + DB_TEST_SCHEMA = "test_hypothesis_3_multi_step" + + eng = create_engine(DBCONNSTR) + try: + with eng.begin() as conn: + conn.execute(text(f"DROP SCHEMA {DB_TEST_SCHEMA} CASCADE")) + except Exception: + pass + + with eng.begin() as conn: + conn.execute(text(f"CREATE SCHEMA {DB_TEST_SCHEMA}")) + + test_dbconn = DBConn(DBCONNSTR, DB_TEST_SCHEMA) + + print("Запуск теста гипотезы 3 (multi-step pipeline)...") + test_hypothesis_3_multi_step_pipeline_update_ts_vs_process_ts_desync(test_dbconn) diff --git a/tests/test_offset_joinspec.py b/tests/test_offset_joinspec.py new file mode 100644 index 00000000..d8cc6c46 --- /dev/null +++ b/tests/test_offset_joinspec.py @@ -0,0 +1,292 @@ +""" +Тест для проверки что offset'ы создаются для JoinSpec таблиц (с join_keys). + +Воспроизводит баг где offset создавался только для главной таблицы (posts), +но не для справочной таблицы (profiles) с join_keys. +""" + +import time + +import pandas as pd +from sqlalchemy import Column, Integer, String + +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_created_for_joinspec_tables(dbconn: DBConn): + """ + Проверяет что offset создается для таблиц с join_keys (JoinSpec). + + Сценарий: + 1. Создаём posts и profiles (profiles с join_keys={'user_id': 'id'}) + 2. Запускаем трансформацию с offset optimization + 3. Проверяем что offset создан ДЛЯ ОБЕИХ таблиц: posts И profiles + """ + ds = DataStore(dbconn, create_meta_table=True) + + # 1. Создать posts таблицу (используем String для id чтобы совпадать с мета-таблицей) + posts_store = TableStoreDB( + dbconn, + "posts", + [ + Column("id", String, primary_key=True), + Column("user_id", String), + Column("content", String), + ], + create_table=True, + ) + posts = ds.create_table("posts", posts_store) + + # 2. Создать profiles таблицу (справочник) + profiles_store = TableStoreDB( + dbconn, + "profiles", + [Column("id", String, primary_key=True), Column("username", String)], + create_table=True, + ) + profiles = ds.create_table("profiles", profiles_store) + + # 3. Создать output таблицу (id - primary key, остальное - данные) + output_store = TableStoreDB( + dbconn, + "posts_with_username", + [ + Column("id", String, primary_key=True), + Column("user_id", String), # Обычная колонка, не primary key + Column("content", String), + Column("username", String), + ], + create_table=True, + ) + output_dt = ds.create_table("posts_with_username", output_store) + + # 4. Добавить данные + process_ts = time.time() + + # 3 поста от 2 пользователей + posts_df = pd.DataFrame([ + {"id": "1", "user_id": "1", "content": "Post 1"}, + {"id": "2", "user_id": "1", "content": "Post 2"}, + {"id": "3", "user_id": "2", "content": "Post 3"}, + ]) + posts.store_chunk(posts_df, now=process_ts) + + # 2 профиля + profiles_df = pd.DataFrame([ + {"id": "1", "username": "alice"}, + {"id": "2", "username": "bob"}, + ]) + profiles.store_chunk(profiles_df, now=process_ts) + + # 5. Создать трансформацию с join_keys + def transform_func(posts_df, profiles_df): + # JOIN posts + profiles + result = posts_df.merge(profiles_df, left_on="user_id", right_on="id", suffixes=("", "_profile")) + return result[["id", "user_id", "content", "username"]] + + step = BatchTransformStep( + ds=ds, + name="test_transform", + func=transform_func, + input_dts=[ + ComputeInput(dt=posts, join_type="full"), # Главная таблица + ComputeInput(dt=profiles, join_type="inner", join_keys={"user_id": "id"}), # JoinSpec таблица + ], + output_dts=[output_dt], + transform_keys=["id"], # Primary key первой таблицы (posts) + use_offset_optimization=True, # ВАЖНО: используем offset optimization + ) + + # 6. Запустить трансформацию + print("\n🚀 Running initial transformation...") + step.run_full(ds) + + # Проверяем результаты трансформации + output_data = output_dt.get_data() + print(f"✅ Output rows created: {len(output_data)}") + print(f"Output data:\n{output_data}") + + # 7. Проверить что offset'ы созданы для ОБЕИХ таблиц + print("\n🔍 Checking offsets...") + # Используем step.get_name() чтобы получить имя с хэшем + transform_name = step.get_name() + print(f"🔑 Transform name with hash: {transform_name}") + offsets = ds.offset_table.get_offsets_for_transformation(transform_name) + + print(f"📊 Offsets created: {offsets}") + + # КРИТИЧЕСКИ ВАЖНО: offset должен быть для posts И для profiles! + assert "posts" in offsets, "Offset for 'posts' table not found!" + assert "profiles" in offsets, "Offset for 'profiles' table not found! (БАГ!)" + + # Оба offset'а должны быть >= process_ts + assert offsets["posts"] >= process_ts, f"posts offset {offsets['posts']} < process_ts {process_ts}" + assert offsets["profiles"] >= process_ts, f"profiles offset {offsets['profiles']} < process_ts {process_ts}" + + # Проверяем что были созданы 3 записи в output + output_data = output_dt.get_data() + assert len(output_data) == 3, f"Expected 3 output rows, got {len(output_data)}" + + # 8. Добавим новые данные и проверим инкрементальную обработку + time.sleep(0.01) # Небольшая задержка для различения timestamp'ов + process_ts2 = time.time() + + # Добавляем 1 новый пост + new_posts_df = pd.DataFrame([ + {"id": "4", "user_id": "1", "content": "New Post 4"}, + ]) + posts.store_chunk(new_posts_df, now=process_ts2) + + # Добавляем 1 новый профиль + new_profiles_df = pd.DataFrame([ + {"id": "3", "username": "charlie"}, + ]) + profiles.store_chunk(new_profiles_df, now=process_ts2) + + # 9. Запускаем инкрементальную обработку + step.run_full(ds) + + # 10. Проверяем что offset'ы обновились + new_offsets = ds.offset_table.get_offsets_for_transformation(transform_name) + + print(f"\n📊 New offsets after incremental run: {new_offsets}") + + # Оба offset'а должны обновиться до process_ts2 + assert new_offsets["posts"] >= process_ts2, f"posts offset not updated: {new_offsets['posts']} < {process_ts2}" + assert new_offsets["profiles"] >= process_ts2, f"profiles offset not updated: {new_offsets['profiles']} < {process_ts2}" + + # Проверяем что теперь 4 записи в output (3 старых + 1 новый пост) + output_data = output_dt.get_data() + assert len(output_data) == 4, f"Expected 4 output rows, got {len(output_data)}" + + print("\n✅ SUCCESS: Offsets created and updated for both posts AND profiles (including JoinSpec table)!") + + +def test_joinspec_update_ts_from_meta_table_not_null(dbconn: DBConn): + """ + Проблема: Для join_keys (reverse join) update_ts подставляется как NULL + из primary_data_tbl, что приводит к переобработке на каждом запуске. + + Сценарий: + 1. Создать основную таблицу (posts) и справочную таблицу (profiles) + 2. Связать через join_keys + 3. Изменить справочную таблицу + 4. Обработать через reverse join + 5. Повторно запустить - выборка должна быть ПУСТОЙ (нет новых изменений) + + Ожидание: update_ts должен браться из мета-таблицы справочника (tbl.c.update_ts), + а не подставляться как NULL из primary_data_tbl + + Код исправления в sql_meta.py (строки 930-936): + - update_ts берется из tbl.c.update_ts вместо NULL + - update_ts добавлен в GROUP BY + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Создаем основную таблицу (posts) + posts_store = TableStoreDB( + dbconn, + "posts_table", + [ + Column("post_id", String, primary_key=True), + Column("user_id", String), + Column("content", String), + ], + create_table=True, + ) + posts_dt = ds.create_table("posts_table", posts_store) + + # Создаем справочную таблицу (profiles) + profiles_store = TableStoreDB( + dbconn, + "profiles_table", + [ + Column("id", String, primary_key=True), + Column("name", String), + ], + create_table=True, + ) + profiles_dt = ds.create_table("profiles_table", profiles_store) + + # Создаем выходную таблицу + output_store = TableStoreDB( + dbconn, + "enriched_posts", + [ + Column("post_id", String, primary_key=True), + Column("user_id", String), + Column("content", String), + ], + create_table=True, + ) + output_dt = ds.create_table("enriched_posts", output_store) + + def join_func(posts_df, profiles_df): + # Обогащаем посты данными из profiles (хотя в этом тесте просто возвращаем posts) + return posts_df[["post_id", "user_id", "content"]] + + step = BatchTransformStep( + ds=ds, + name="join_test", + func=join_func, + input_dts=[ + ComputeInput(dt=posts_dt, join_type="full"), # Основная таблица + ComputeInput( + dt=profiles_dt, + join_type="full", + join_keys={"user_id": "id"}, # Reverse join + ), + ], + output_dts=[output_dt], + transform_keys=["post_id"], + use_offset_optimization=True, + chunk_size=10, + ) + + # 1. Создать данные в основной таблице + t1 = time.time() + posts_df = pd.DataFrame({ + "post_id": ["post_1", "post_2"], + "user_id": ["user_1", "user_2"], + "content": ["Hello", "World"], + }) + posts_dt.store_chunk(posts_df, now=t1) + + # 2. Создать данные в справочной таблице + time.sleep(0.01) + t2 = time.time() + profiles_df = pd.DataFrame({ + "id": ["user_1", "user_2"], + "name": ["Alice", "Bob"], + }) + profiles_dt.store_chunk(profiles_df, now=t2) + + # 3. Первый прогон - должен обработать все записи + step.run_full(ds) + + output_data = output_dt.get_data() + assert len(output_data) == 2 + assert set(output_data["post_id"]) == {"post_1", "post_2"} + + # 4. Проверить offset + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert "posts_table" in offsets + assert "profiles_table" in offsets + posts_offset = offsets["posts_table"] + profiles_offset = offsets["profiles_table"] + + # 5. Повторный прогон - выборка должна быть ПУСТОЙ (нет новых изменений) + # Получаем количество батчей для обработки + idx_count, idx_gen = step.get_full_process_ids(ds=ds, run_config=None) + + assert idx_count == 0, ( + f"После первого успешного run_full выборка должна быть пустой. " + f"Получено {idx_count} батчей для обработки. " + f"Это указывает на то, что offset-оптимизация для JoinSpec НЕ работает: " + f"записи с join_keys переобрабатываются на каждом запуске из-за update_ts = NULL." + ) + + print("\n✅ SUCCESS: JoinSpec update_ts correctly taken from meta table, no reprocessing!") diff --git a/tests/test_offset_optimization_runtime_switch.py b/tests/test_offset_optimization_runtime_switch.py new file mode 100644 index 00000000..46b3bfd8 --- /dev/null +++ b/tests/test_offset_optimization_runtime_switch.py @@ -0,0 +1,340 @@ +""" +Тесты для динамического переключения между v1 и v2 через RunConfig.labels +""" +import time + +import pandas as pd +from sqlalchemy import Column, Integer, String + +from datapipe.compute import ComputeInput +from datapipe.datatable import DataStore +from datapipe.run_config import RunConfig +from datapipe.step.batch_transform import BatchTransformStep +from datapipe.store.database import DBConn, TableStoreDB + + +def test_runtime_switch_via_labels_default_false(dbconn: DBConn): + """Тест переключения режима через RunConfig.labels при use_offset_optimization=False""" + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "test_input", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("test_input", input_store) + + output_store = TableStoreDB( + dbconn, + "test_output", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("test_output", output_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + # Создаем step с ВЫКЛЮЧЕННОЙ offset оптимизацией + step = BatchTransformStep( + ds=ds, + name="test_transform", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=False, # По умолчанию ВЫКЛЮЧЕНО + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2", "3"], "value": [10, 20, 30]}), now=now + ) + + # Первый запуск БЕЗ offset (дефолтное поведение) + step.run_full(ds) + + output_data = output_dt.get_data() + assert len(output_data) == 3 + + # Добавляем новые данные + time.sleep(0.01) + now2 = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["4", "5"], "value": [40, 50]}), now=now2 + ) + + # Второй запуск С ВКЛЮЧЕНИЕМ offset через RunConfig.labels + run_config = RunConfig(labels={"use_offset_optimization": True}) + step.run_full(ds, run_config=run_config) + + # Проверяем результат + output_data = output_dt.get_data() + assert len(output_data) == 5 + assert sorted(output_data["id"].tolist()) == ["1", "2", "3", "4", "5"] + + +def test_runtime_switch_via_labels_default_true(dbconn: DBConn): + """Тест переключения режима через RunConfig.labels при use_offset_optimization=True""" + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "test_input2", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("test_input2", input_store) + + output_store = TableStoreDB( + dbconn, + "test_output2", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("test_output2", output_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + # Создаем step с ВКЛЮЧЕННОЙ offset оптимизацией + step = BatchTransformStep( + ds=ds, + name="test_transform2", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, # По умолчанию ВКЛЮЧЕНО + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2", "3"], "value": [10, 20, 30]}), now=now + ) + + # Первый запуск С offset (дефолтное поведение) + step.run_full(ds) + + output_data = output_dt.get_data() + assert len(output_data) == 3 + + # Добавляем новые данные + time.sleep(0.01) + now2 = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["4", "5"], "value": [40, 50]}), now=now2 + ) + + # Второй запуск С ВЫКЛЮЧЕНИЕМ offset через RunConfig.labels + run_config = RunConfig(labels={"use_offset_optimization": False}) + step.run_full(ds, run_config=run_config) + + # Проверяем результат + output_data = output_dt.get_data() + assert len(output_data) == 5 + assert sorted(output_data["id"].tolist()) == ["1", "2", "3", "4", "5"] + + +def test_runtime_switch_multiple_runs(dbconn: DBConn): + """Тест множественных переключений между режимами""" + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "test_input3", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("test_input3", input_store) + + output_store = TableStoreDB( + dbconn, + "test_output3", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("test_output3", output_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + step = BatchTransformStep( + ds=ds, + name="test_transform3", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=False, + ) + + # Запуск 1: v1 режим (по умолчанию) + now = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["1"], "value": [10]}), now=now) + step.run_full(ds) + assert len(output_dt.get_data()) == 1 + + # Запуск 2: переключаемся на v2 через labels + time.sleep(0.01) + now = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["2"], "value": [20]}), now=now) + step.run_full(ds, run_config=RunConfig(labels={"use_offset_optimization": True})) + assert len(output_dt.get_data()) == 2 + + # Запуск 3: обратно на v1 через labels + time.sleep(0.01) + now = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["3"], "value": [30]}), now=now) + step.run_full(ds, run_config=RunConfig(labels={"use_offset_optimization": False})) + assert len(output_dt.get_data()) == 3 + + # Запуск 4: снова v2 + time.sleep(0.01) + now = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["4"], "value": [40]}), now=now) + step.run_full(ds, run_config=RunConfig(labels={"use_offset_optimization": True})) + assert len(output_dt.get_data()) == 4 + + # Проверяем финальный результат + output_data = output_dt.get_data() + assert sorted(output_data["id"].tolist()) == ["1", "2", "3", "4"] + assert sorted(output_data["result"].tolist()) == [10, 20, 30, 40] + + +def test_get_changed_idx_count_respects_label_override(dbconn: DBConn): + """Тест что get_changed_idx_count также учитывает переопределение через labels""" + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "test_input4", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("test_input4", input_store) + + output_store = TableStoreDB( + dbconn, + "test_output4", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("test_output4", output_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + # Создаем step с включенной offset оптимизацией для корректной работы теста + step = BatchTransformStep( + ds=ds, + name="test_transform4", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, # Включаем для обновления offset'ов + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2", "3"], "value": [10, 20, 30]}), now=now + ) + + # Проверяем get_changed_idx_count с v2 (дефолт) + count_v2 = step.get_changed_idx_count(ds) + assert count_v2 == 3 + + # Проверяем get_changed_idx_count с override на v1 + count_v1 = step.get_changed_idx_count( + ds, run_config=RunConfig(labels={"use_offset_optimization": False}) + ) + assert count_v1 == 3 # Должно быть то же самое количество + + # Обрабатываем с offset оптимизацией (обновляются offset'ы) + step.run_full(ds) + + # Добавляем новые данные + time.sleep(0.01) + now2 = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["4"], "value": [40]}), now=now2) + + # v2 увидит только новую запись (фильтрация по offset) + count_v2_after = step.get_changed_idx_count(ds) + assert count_v2_after == 1 + + # v1 также увидит только новую запись (метаданные уже обработаны) + count_v1_after = step.get_changed_idx_count( + ds, run_config=RunConfig(labels={"use_offset_optimization": False}) + ) + assert count_v1_after == 1 + + +def test_get_full_process_ids_respects_label_override(dbconn: DBConn): + """Тест что get_full_process_ids также учитывает переопределение через labels""" + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "test_input5", + [Column("id", String, primary_key=True), Column("value", Integer)], + create_table=True, + ) + input_dt = ds.create_table("test_input5", input_store) + + output_store = TableStoreDB( + dbconn, + "test_output5", + [Column("id", String, primary_key=True), Column("result", Integer)], + create_table=True, + ) + output_dt = ds.create_table("test_output5", output_store) + + def transform_func(df): + return df.rename(columns={"value": "result"}) + + # Создаем step с включенной offset оптимизацией для корректной работы теста + step = BatchTransformStep( + ds=ds, + name="test_transform5", + func=transform_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, # Включаем для обновления offset'ов + ) + + # Добавляем данные + now = time.time() + input_dt.store_chunk( + pd.DataFrame({"id": ["1", "2", "3"], "value": [10, 20, 30]}), now=now + ) + + # Обрабатываем с offset оптимизацией (обновляются offset'ы) + step.run_full(ds) + + # Добавляем новые данные + time.sleep(0.01) + now2 = time.time() + input_dt.store_chunk(pd.DataFrame({"id": ["4"], "value": [40]}), now=now2) + + # v2 (дефолт) должен вернуть только новую запись + count_v2, ids_v2 = step.get_full_process_ids(ds) + ids_v2_list = pd.concat(list(ids_v2)) + assert count_v2 == 1 + assert len(ids_v2_list) == 1 + assert ids_v2_list["id"].iloc[0] == "4" + + # v1 (через override) также должен вернуть только новую запись, + # так как метаданные уже обработаны + count_v1, ids_v1 = step.get_full_process_ids( + ds, run_config=RunConfig(labels={"use_offset_optimization": False}) + ) + ids_v1_list = pd.concat(list(ids_v1)) + assert count_v1 == 1 + assert len(ids_v1_list) == 1 + assert ids_v1_list["id"].iloc[0] == "4" diff --git a/tests/test_offset_pipeline_integration.py b/tests/test_offset_pipeline_integration.py new file mode 100644 index 00000000..1b85fe00 --- /dev/null +++ b/tests/test_offset_pipeline_integration.py @@ -0,0 +1,322 @@ +""" +Интеграционный тест для пайплайна с offset-оптимизацией. + +Проверяет полный цикл работы простого пайплайна копирования данных: +- Добавление новых записей +- Изменение существующих записей +- Удаление записей +- Смешанные операции +""" +import time + +import pandas as pd +import pytest +from sqlalchemy import Column, Integer, String + +from datapipe.compute import ComputeInput +from datapipe.datatable import DataStore, DataTable +from datapipe.step.batch_transform import BatchTransformStep +from datapipe.store.database import DBConn, TableStoreDB + + +@pytest.fixture +def setup_copy_pipeline(dbconn: DBConn): + """ + Фикстура для создания простого пайплайна копирования данных. + + Возвращает: + (ds, source_dt, target_dt, step) - готовый пайплайн для тестирования + """ + ds = DataStore(dbconn, create_meta_table=True) + + # Создаем source таблицу + source_store = TableStoreDB( + dbconn, + "pipeline_source", + [ + Column("id", String, primary_key=True), + Column("name", String), + Column("value", Integer), + ], + create_table=True, + ) + source_dt = ds.create_table("pipeline_source", source_store) + + # Создаем target таблицу + target_store = TableStoreDB( + dbconn, + "pipeline_target", + [ + Column("id", String, primary_key=True), + Column("name", String), + Column("value", Integer), + ], + create_table=True, + ) + target_dt = ds.create_table("pipeline_target", target_store) + + # Функция копирования (просто возвращает данные как есть) + def copy_transform(df): + return df[["id", "name", "value"]] + + # Создаем step с offset-оптимизацией + step = BatchTransformStep( + ds=ds, + name="copy_pipeline", + func=copy_transform, + input_dts=[ComputeInput(dt=source_dt, join_type="full")], + output_dts=[target_dt], + transform_keys=["id"], + use_offset_optimization=True, + ) + + return ds, source_dt, target_dt, step + + +def test_offset_pipeline_insert_records(setup_copy_pipeline): + """ + Тест добавления новых записей в пайплайн с offset-оптимизацией. + + Проверяет: + 1. Начальное добавление данных + 2. Добавление новых записей (только новые обрабатываются) + 3. Корректность данных в target + 4. Обновление offset'ов + """ + ds, source_dt, target_dt, step = setup_copy_pipeline + + # ========== Добавляем начальные данные ========== + now1 = time.time() + initial_data = pd.DataFrame({ + "id": ["A", "B", "C"], + "name": ["Alice", "Bob", "Charlie"], + "value": [10, 20, 30], + }) + source_dt.store_chunk(initial_data, now=now1) + + # Запускаем трансформацию + step.run_full(ds) + + # Проверяем что данные скопировались в target + target_data = target_dt.get_data().sort_values("id").reset_index(drop=True) + expected = initial_data.sort_values("id").reset_index(drop=True) + pd.testing.assert_frame_equal(target_data, expected) + + # Проверяем что offset установлен + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert len(offsets) == 1 + assert "pipeline_source" in offsets + assert offsets["pipeline_source"] >= now1 + + # ========== Добавляем новые записи ========== + time.sleep(0.01) + now2 = time.time() + new_data = pd.DataFrame({ + "id": ["D", "E"], + "name": ["David", "Eve"], + "value": [40, 50], + }) + source_dt.store_chunk(new_data, now=now2) + + # Проверяем что changed_count показывает только новые записи + changed_count = step.get_changed_idx_count(ds) + assert changed_count == 2, f"Expected 2 new records, got {changed_count}" + + # Запускаем трансформацию + step.run_full(ds) + + # Проверяем что в target теперь все 5 записей + target_data = target_dt.get_data().sort_values("id").reset_index(drop=True) + assert len(target_data) == 5 + expected_all = pd.concat([initial_data, new_data]).sort_values("id").reset_index(drop=True) + pd.testing.assert_frame_equal(target_data, expected_all) + + # Проверяем что offset обновился + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert offsets["pipeline_source"] >= now2 + + +def test_offset_pipeline_update_records(setup_copy_pipeline): + """ + Тест изменения существующих записей в пайплайне с offset-оптимизацией. + + Проверяет: + 1. Обновление существующих записей + 2. Только измененные записи обрабатываются + 3. Корректность обновленных данных в target + 4. Обновление offset'ов + """ + ds, source_dt, target_dt, step = setup_copy_pipeline + + # Добавляем начальные данные + now1 = time.time() + initial_data = pd.DataFrame({ + "id": ["A", "B", "C"], + "name": ["Alice", "Bob", "Charlie"], + "value": [10, 20, 30], + }) + source_dt.store_chunk(initial_data, now=now1) + step.run_full(ds) + + # ========== Изменяем существующие записи ========== + time.sleep(0.01) + now2 = time.time() + updated_data = pd.DataFrame({ + "id": ["B", "C"], + "name": ["Bob Updated", "Charlie Updated"], + "value": [200, 300], + }) + source_dt.store_chunk(updated_data, now=now2) + + # Проверяем что changed_count показывает измененные записи + changed_count = step.get_changed_idx_count(ds) + assert changed_count == 2, f"Expected 2 updated records, got {changed_count}" + + # Запускаем трансформацию + step.run_full(ds) + + # Проверяем что в target обновились записи B и C + target_data = target_dt.get_data().sort_values("id").reset_index(drop=True) + assert len(target_data) == 3 + + # Проверяем конкретные значения + row_a = target_data[target_data["id"] == "A"].iloc[0] + assert row_a["name"] == "Alice" + assert row_a["value"] == 10 + + row_b = target_data[target_data["id"] == "B"].iloc[0] + assert row_b["name"] == "Bob Updated" + assert row_b["value"] == 200 + + row_c = target_data[target_data["id"] == "C"].iloc[0] + assert row_c["name"] == "Charlie Updated" + assert row_c["value"] == 300 + + # Проверяем что offset обновился + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert offsets["pipeline_source"] >= now2 + + +def test_offset_pipeline_delete_records(setup_copy_pipeline): + """ + Тест удаления записей в пайплайне с offset-оптимизацией. + + Проверяет: + 1. Удаление записей из source + 2. Только удаленные записи обрабатываются + 3. Записи удаляются из target + 4. Обновление offset'ов + """ + ds, source_dt, target_dt, step = setup_copy_pipeline + + # Добавляем начальные данные + now1 = time.time() + initial_data = pd.DataFrame({ + "id": ["A", "B", "C", "D"], + "name": ["Alice", "Bob", "Charlie", "David"], + "value": [10, 20, 30, 40], + }) + source_dt.store_chunk(initial_data, now=now1) + step.run_full(ds) + + # Проверяем что все данные в target + assert len(target_dt.get_data()) == 4 + + # ========== Удаляем записи ========== + time.sleep(0.01) + now2 = time.time() + delete_idx = pd.DataFrame({"id": ["A", "C"]}) + source_dt.delete_by_idx(delete_idx, now=now2) + + # Проверяем что changed_count показывает удаленные записи + changed_count = step.get_changed_idx_count(ds) + assert changed_count == 2, f"Expected 2 deleted records, got {changed_count}" + + # Запускаем трансформацию + step.run_full(ds) + + # Проверяем что в target остались только B и D + target_data = target_dt.get_data().sort_values("id").reset_index(drop=True) + assert len(target_data) == 2 + assert sorted(target_data["id"].tolist()) == ["B", "D"] + + # Проверяем что значения корректные + row_b = target_data[target_data["id"] == "B"].iloc[0] + assert row_b["name"] == "Bob" + assert row_b["value"] == 20 + + row_d = target_data[target_data["id"] == "D"].iloc[0] + assert row_d["name"] == "David" + assert row_d["value"] == 40 + + +def test_offset_pipeline_mixed_operations(setup_copy_pipeline): + """ + Тест смешанных операций (добавление + изменение + удаление) в одном батче. + + Проверяет: + 1. Одновременное добавление новых, изменение существующих и удаление записей + 2. Все операции обрабатываются корректно в одном запуске + 3. Финальное состояние target корректно + """ + ds, source_dt, target_dt, step = setup_copy_pipeline + + # Добавляем начальные данные + now1 = time.time() + initial_data = pd.DataFrame({ + "id": ["1", "2", "3"], + "name": ["One", "Two", "Three"], + "value": [100, 200, 300], + }) + source_dt.store_chunk(initial_data, now=now1) + step.run_full(ds) + + # Проверяем начальное состояние + assert len(target_dt.get_data()) == 3 + + # ========== Одновременно: добавляем "4", изменяем "2", удаляем "3" ========== + time.sleep(0.01) + now2 = time.time() + + # Добавляем новую запись + source_dt.store_chunk( + pd.DataFrame({"id": ["4"], "name": ["Four"], "value": [400]}), now=now2 + ) + + # Изменяем существующую + source_dt.store_chunk( + pd.DataFrame({"id": ["2"], "name": ["Two Updated"], "value": [222]}), now=now2 + ) + + # Удаляем + delete_idx = pd.DataFrame({"id": ["3"]}) + source_dt.delete_by_idx(delete_idx, now=now2) + + # Проверяем что changed_count = 3 (новая + измененная + удаленная) + changed_count = step.get_changed_idx_count(ds) + assert changed_count == 3, f"Expected 3 changed records, got {changed_count}" + + # Запускаем трансформацию + step.run_full(ds) + + # Проверяем результат: должны остаться "1", "2", "4" + target_data = target_dt.get_data().sort_values("id").reset_index(drop=True) + assert len(target_data) == 3 + assert sorted(target_data["id"].tolist()) == ["1", "2", "4"] + + # Проверяем значения каждой записи + row_1 = target_data[target_data["id"] == "1"].iloc[0] + assert row_1["name"] == "One" + assert row_1["value"] == 100 + + row_2 = target_data[target_data["id"] == "2"].iloc[0] + assert row_2["name"] == "Two Updated" # изменился + assert row_2["value"] == 222 + + row_4 = target_data[target_data["id"] == "4"].iloc[0] + assert row_4["name"] == "Four" # новый + assert row_4["value"] == 400 + + # Проверяем что offset обновился + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + assert offsets["pipeline_source"] >= now2 diff --git a/tests/test_offset_production_bug_main.py b/tests/test_offset_production_bug_main.py new file mode 100644 index 00000000..031a0df5 --- /dev/null +++ b/tests/test_offset_production_bug_main.py @@ -0,0 +1,416 @@ +""" +🚨 КРИТИЧЕСКИЙ PRODUCTION БАГ: Offset Optimization теряет данные + +ПРОБЛЕМА В PRODUCTION: +- Дата: 08.12.2025 +- Потеряно: 48,915 из 82,000 записей (60%) +- Причина: Строгое неравенство в WHERE update_ts > offset + +КОРНЕВАЯ ПРИЧИНА: +Код: datapipe/meta/sql_meta.py:967 + tbl.c.update_ts > offset # ❌ ОШИБКА! Должно быть >= + +МЕХАНИЗМ БАГА: +1. Батч обрабатывает записи в ORDER BY (id, hashtag) - НЕ по времени! +2. Батч содержит записи с РАЗНЫМИ update_ts +3. offset = MAX(update_ts) из обработанного батча +4. Следующий запуск: WHERE update_ts > offset (строгое неравенство!) +5. Записи с update_ts == offset но не вошедшие в батч ТЕРЯЮТСЯ + +ВИЗУАЛИЗАЦИЯ ПРОБЛЕМЫ: +┌─────────────────────────────────────────────────────────────────────┐ +│ Временная шкала (update_ts): │ +│ │ +│ T1 (16:21) T2 (17:00) T3 (18:00) T4 (19:00) T5 (20:29) │ +│ │ │ │ │ │ │ +│ ▼ ▼ ▼ ▼ ▼ │ +│ rec_00 rec_08 rec_13 rec_18 rec_22 │ +│ rec_01 rec_09 rec_14 rec_19 rec_23 │ +│ ... rec_10 rec_15 rec_20 rec_24 │ +│ rec_07 rec_11 rec_16 rec_21 │ +│ rec_12 rec_17 │ +│ │ +│ ORDER BY id (сортировка для обработки): │ +│ rec_00 → rec_01 → ... → rec_07 → rec_08 → rec_09 → rec_10 → ... │ +│ T1 T1 T1 T2 T2 T2 │ +│ │ +│ ┌──────────────────────────┐ │ +│ │ БАТЧ 1 (chunk_size=10) │ │ +│ │ rec_00 до rec_09 │ │ +│ │ update_ts: T1...T1, T2 │ │ +│ └──────────────────────────┘ │ +│ ↓ │ +│ offset = MAX(T1, T2) = T2 │ +│ │ +│ Следующий запуск: │ +│ WHERE update_ts > T2 ← СТРОГОЕ НЕРАВЕНСТВО! │ +│ │ +│ 🚨 ПОТЕРЯНЫ: │ +│ rec_10, rec_11, rec_12 (update_ts = T2 == offset) │ +│ │ +└─────────────────────────────────────────────────────────────────────┘ + +РЕШЕНИЕ: +Заменить > на >= в sql_meta.py:967: + tbl.c.update_ts >= offset # ✅ ПРАВИЛЬНО +""" + +import time +from typing import List, Tuple + +import pandas as pd +import pytest +from sqlalchemy import Column, Integer, String + +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 prepare_test_data() -> List[Tuple[str, str, float]]: + """ + Подготовка тестовых данных для воспроизведения production бага. + + Данные имитируют накопление записей в течение нескольких часов + с РАЗНЫМИ update_ts (как в реальной системе). + + Returns: + List[(record_id, label, update_ts_offset)] + + Временная шкала: + T1 = base_time + 1 (16:21 в production) + T2 = base_time + 2 (17:00) + T3 = base_time + 3 (18:00) + T4 = base_time + 4 (19:00) + T5 = base_time + 5 (20:29 в production) + """ + # Формат: (id, label для логов, смещение timestamp от base_time) + test_data = [ + # Группа 1: T1 (16:21) - 8 записей + ("rec_00", "T1", 1.0), + ("rec_01", "T1", 1.0), + ("rec_02", "T1", 1.0), + ("rec_03", "T1", 1.0), + ("rec_04", "T1", 1.0), + ("rec_05", "T1", 1.0), + ("rec_06", "T1", 1.0), + ("rec_07", "T1", 1.0), + + # Группа 2: T2 (17:00) - 5 записей + # ⚠️ КРИТИЧНО: rec_08 и rec_09 войдут в ПЕРВЫЙ батч + # rec_10, rec_11, rec_12 останутся на ВТОРОЙ батч + ("rec_08", "T2", 2.0), + ("rec_09", "T2", 2.0), + ("rec_10", "T2", 2.0), # 🚨 БУДЕТ ПОТЕРЯНА + ("rec_11", "T2", 2.0), # 🚨 БУДЕТ ПОТЕРЯНА + ("rec_12", "T2", 2.0), # 🚨 БУДЕТ ПОТЕРЯНА + + # Группа 3: T3 (18:00) - 5 записей + ("rec_13", "T3", 3.0), + ("rec_14", "T3", 3.0), + ("rec_15", "T3", 3.0), + ("rec_16", "T3", 3.0), + ("rec_17", "T3", 3.0), + + # Группа 4: T4 (19:00) - 4 записи + ("rec_18", "T4", 4.0), + ("rec_19", "T4", 4.0), + ("rec_20", "T4", 4.0), + ("rec_21", "T4", 4.0), + + # Группа 5: T5 (20:29) - 3 записи + ("rec_22", "T5", 5.0), + ("rec_23", "T5", 5.0), + ("rec_24", "T5", 5.0), + ] + + return test_data + + +def print_test_data_visualization(test_data: List[Tuple[str, str, float]], base_time: float): + """Визуализация тестовых данных для отладки""" + print("\n" + "=" * 80) + print("ПОДГОТОВЛЕННЫЕ ТЕСТОВЫЕ ДАННЫЕ") + print("=" * 80) + print("\nВсего записей:", len(test_data)) + print("\nРаспределение по временным меткам:") + + # Группируем по timestamp + by_timestamp = {} + for record_id, label, offset in test_data: + ts = base_time + offset + if ts not in by_timestamp: + by_timestamp[ts] = [] + by_timestamp[ts].append((record_id, label)) + + for ts in sorted(by_timestamp.keys()): + records = by_timestamp[ts] + label = records[0][1] + ids = [r[0] for r in records] + print(f" {label}: {len(records)} записей - {', '.join(ids)}") + + print("\nОжидаемое распределение по батчам (chunk_size=10, ORDER BY id):") + print(" Батч 1: rec_00 до rec_09 (10 записей)") + print(" update_ts: T1(8 записей), T2(2 записи)") + print(" offset после батча = MAX(T1, T2) = T2") + print() + print(" 🚨 КРИТИЧНО: Следующий запуск WHERE update_ts > T2") + print(" ПРОПУСТИТ: rec_10, rec_11, rec_12 (update_ts == T2)") + print() + + +# ============================================================================ +# PRODUCTION БАГ ТЕСТ +# ============================================================================ + +def test_production_bug_offset_loses_records_with_equal_update_ts(dbconn: DBConn): + """ + Тест >= неравенства в offset фильтре (исправление production бага). + + ИСХОДНЫЙ PRODUCTION БАГ (до атомарного commit): + - 82,000 записей накоплено, chunk_size=1000 + - Обработан только ПЕРВЫЙ батч (частичная обработка) + - offset = MAX(update_ts) из батча + - Оставшиеся записи с update_ts == offset ПОТЕРЯНЫ (48,915 записей, 60%) + - Причина: WHERE update_ts > offset (строгое >) вместо >= + + ПОЧЕМУ СТАРЫЙ ТЕСТ НЕ РАБОТАЕТ: + С новым атомарным commit механизмом невозможно симулировать частичную + обработку через run_idx() - offset коммитится только после полного run_full(). + + НОВАЯ ВЕРСИЯ ТЕСТА (совместима с атомарным commit): + 1. Загружаем 25 записей с разными update_ts + 2. ПЕРВЫЙ run_full() обрабатывает ВСЕ записи, offset = MAX(update_ts) + 3. Добавляем НОВЫЕ записи с update_ts == offset (критический случай!) + 4. ВТОРОЙ run_full() должен обработать эти записи (тест >= вместо >) + 5. Проверяем что НЕТ потерь данных + + КОГДА ВОЗМОЖЕН СЦЕНАРИЙ "update_ts == offset между запусками"? + - Clock skew между серверами (разные системные часы) + - Backfill старых данных с прошлыми timestamp + - Delayed records из очереди с задержкой + - Ручное добавление записей с кастомным timestamp + + СУТЬ ТЕСТА: + Проверяем что >= работает корректно и записи с update_ts == offset + обрабатываются, а не теряются (независимо от сценария возникновения). + """ + # ========== SETUP ========== + ds = DataStore(dbconn, create_meta_table=True) + + input_store = TableStoreDB( + dbconn, + "production_bug_input", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + input_dt = ds.create_table("production_bug_input", input_store) + + output_store = TableStoreDB( + dbconn, + "production_bug_output", + [ + Column("id", String, primary_key=True), + Column("value", Integer), + ], + create_table=True, + ) + output_dt = ds.create_table("production_bug_output", output_store) + + def copy_func(df): + """Простая функция копирования (как copy_to_online в production)""" + return df[["id", "value"]] + + step = BatchTransformStep( + ds=ds, + name="production_bug_copy", + func=copy_func, + input_dts=[ComputeInput(dt=input_dt, join_type="full")], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=10, # Маленький для быстрого теста (в production=1000) + ) + + # ========== ПОДГОТОВКА НАЧАЛЬНЫХ ДАННЫХ ========== + base_time = time.time() + test_data = prepare_test_data() + + # Визуализация данных + print_test_data_visualization(test_data, base_time) + + # Загружаем начальные данные группами по timestamp + for record_id, label, offset in test_data: + ts = base_time + offset + input_dt.store_chunk( + pd.DataFrame({"id": [record_id], "value": [int(offset * 100)]}), + now=ts + ) + time.sleep(0.001) # Небольшая задержка для корректности timestamp + + # Проверяем метаданные + all_meta = input_dt.meta_table.get_metadata() + print(f"\n✓ Всего записей загружено: {len(all_meta)}") + + # ========== ПЕРВЫЙ ЗАПУСК (обработка всех начальных записей) ========== + print("\n" + "=" * 80) + print("ПЕРВЫЙ ЗАПУСК ТРАНСФОРМАЦИИ (run_full)") + print("=" * 80) + + step.run_full(ds) + + # Проверяем offset после первого запуска + offsets = ds.offset_table.get_offsets_for_transformation(step.get_name()) + offset_after_first = offsets["production_bug_input"] + + output_after_first = output_dt.get_data() + + print(f"\n✓ Обработано записей: {len(output_after_first)}") + print(f"✓ offset установлен на: {offset_after_first:.2f}") + + # Показываем какие записи обработаны + processed_ids = sorted(output_after_first["id"].tolist()) + print(f"✓ Обработанные id: {', '.join(processed_ids[:5])}...{', '.join(processed_ids[-2:])}") + + # Проверяем что все начальные записи обработаны + assert len(output_after_first) == len(test_data), ( + f"ОШИБКА: Первый run_full должен обработать все записи. " + f"Ожидалось {len(test_data)}, получено {len(output_after_first)}" + ) + + # ========== КРИТИЧЕСКИЙ СЦЕНАРИЙ: Добавляем записи с update_ts == offset ========== + print("\n" + "=" * 80) + print("КРИТИЧЕСКИЙ СЦЕНАРИЙ: Добавление записей с update_ts == offset") + print("=" * 80) + + # Добавляем НОВЫЕ записи с timestamp РАВНЫМ offset + # Это воспроизводит production баг: записи с update_ts == offset должны обрабатываться! + critical_timestamp = offset_after_first + critical_records = [ + ("rec_critical_01", 999), + ("rec_critical_02", 998), + ("rec_critical_03", 997), + ] + + print(f"\nДобавляем {len(critical_records)} записей с update_ts == {critical_timestamp:.2f}") + for record_id, value in critical_records: + input_dt.store_chunk( + pd.DataFrame({"id": [record_id], "value": [value]}), + now=critical_timestamp + ) + time.sleep(0.001) + + # Проверяем что записи действительно имеют update_ts == offset + critical_meta = input_dt.meta_table.get_metadata( + pd.DataFrame({"id": [rec[0] for rec in critical_records]}) + ) + for idx, row in critical_meta.iterrows(): + assert abs(row["update_ts"] - critical_timestamp) < 0.01, ( + f"ОШИБКА В ТЕСТЕ: Запись {row['id']} должна иметь update_ts == offset" + ) + print(f" {row['id']}: update_ts={row['update_ts']:.2f} == offset={critical_timestamp:.2f}") + + # ========== ВТОРОЙ ЗАПУСК (тестируем >= вместо >) ========== + print("\n" + "=" * 80) + print("ВТОРОЙ ЗАПУСК ТРАНСФОРМАЦИИ (проверка >= вместо >)") + print("=" * 80) + + # Проверяем сколько записей будет обработано + changed_count = step.get_changed_idx_count(ds) + print(f"Записей для обработки: {changed_count}") + + if changed_count == 0: + pytest.fail( + f"\n🚨 КРИТИЧЕСКИЙ БАГ В OFFSET OPTIMIZATION!\n" + f"{'=' * 50}\n" + f"Добавлено {len(critical_records)} НОВЫХ записей с update_ts == offset={critical_timestamp:.2f}\n" + f"Но get_changed_idx_count вернул 0 - записи НЕ ВИДНЫ для обработки!\n\n" + f"МЕХАНИЗМ БАГА:\n" + f"WHERE update_ts > offset (строгое неравенство!) пропускает записи с update_ts == offset\n" + f"Должно быть: WHERE update_ts >= offset\n\n" + f"В PRODUCTION: Этот баг привел к потере 48,915 из 82,000 записей (60%)\n" + f"{'=' * 50}" + ) + + # NOTE: changed_count может быть > len(critical_records) потому что >= включает + # старые записи с update_ts == offset. Система отфильтрует их по process_ts. + # Главное - чтобы критические записи были видны и обработаны! + print(f" (может включать старые записи с update_ts == offset, они будут отфильтрованы по process_ts)") + + # Запускаем обработку + step.run_full(ds) + + # ========== ПРОВЕРКА РЕЗУЛЬТАТА ========== + print("\n" + "=" * 80) + print("ПРОВЕРКА РЕЗУЛЬТАТА") + print("=" * 80) + + final_output = output_dt.get_data() + final_processed_ids = set(final_output["id"].tolist()) + + # Проверяем что все критические записи обработаны + all_critical_processed = all(rec[0] in final_processed_ids for rec in critical_records) + + print(f"\nФинальный результат:") + print(f" Начальных записей: {len(test_data)}") + print(f" Критических записей: {len(critical_records)}") + print(f" ВСЕГО ожидается: {len(test_data) + len(critical_records)}") + print(f" Обработано в output: {len(final_output)}") + + if not all_critical_processed: + lost_critical = [rec[0] for rec in critical_records if rec[0] not in final_processed_ids] + print(f"\n🚨 ПОТЕРЯНЫ КРИТИЧЕСКИЕ ЗАПИСИ: {lost_critical}") + + pytest.fail( + f"\n🚨 КРИТИЧЕСКИЙ БАГ В OFFSET OPTIMIZATION!\n" + f"{'=' * 50}\n" + f"Критические записи с update_ts == offset НЕ обработаны!\n" + f"Потеряно: {len(lost_critical)} из {len(critical_records)}\n" + f"Потерянные id: {lost_critical}\n\n" + f"МЕХАНИЗМ БАГА:\n" + f"WHERE update_ts > offset (строгое >) вместо >=\n" + f"Записи с update_ts == offset пропускаются!\n\n" + f"В PRODUCTION: 82,000 записей, потеряно 48,915 (60%)\n" + f"{'=' * 50}" + ) + + # Финальная проверка: все записи обработаны + expected_total = len(test_data) + len(critical_records) + assert len(final_output) == expected_total, ( + f"Ожидалось {expected_total} записей, получено {len(final_output)}" + ) + + print(f"\n✅ Все записи обработаны корректно!") + print(f"✅ Записи с update_ts == offset обработаны (>= работает правильно)") + + +if __name__ == "__main__": + # Для ручного запуска и отладки + from datapipe.store.database import DBConn + from sqlalchemy import create_engine, text + + DBCONNSTR = "postgresql://postgres:password@localhost:5432/postgres" + DB_TEST_SCHEMA = "test_production_bug" + + eng = create_engine(DBCONNSTR) + try: + with eng.begin() as conn: + conn.execute(text(f"DROP SCHEMA {DB_TEST_SCHEMA} CASCADE")) + except Exception: + pass + + with eng.begin() as conn: + conn.execute(text(f"CREATE SCHEMA {DB_TEST_SCHEMA}")) + + test_dbconn = DBConn(DBCONNSTR, DB_TEST_SCHEMA) + + print("Запуск теста воспроизведения production бага...") + test_production_bug_offset_loses_records_with_equal_update_ts(test_dbconn) diff --git a/tests/test_offset_table.py b/tests/test_offset_table.py new file mode 100644 index 00000000..509ae5dc --- /dev/null +++ b/tests/test_offset_table.py @@ -0,0 +1,234 @@ +import time + +import sqlalchemy as sa + +from datapipe.meta.sql_meta import TransformInputOffsetTable +from datapipe.store.database import DBConn + + +def test_offset_table_create(dbconn: DBConn): + """Тест создания таблицы и индексов""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Проверяем, что таблица создана + assert offset_table.sql_table is not None + + # Проверяем, что таблица существует в БД + inspector = sa.inspect(dbconn.con) + schema = offset_table.sql_table.schema + assert offset_table.sql_table.name in inspector.get_table_names(schema=schema) + + +def test_offset_table_get_offset_empty(dbconn: DBConn): + """Тест получения offset'а для несуществующей трансформации""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + offset = offset_table.get_offset("test_transform", "test_table") + assert offset is None + + +def test_offset_table_update_and_get(dbconn: DBConn): + """Тест обновления и получения offset'а""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Обновляем offset + offset_table.update_offset("test_transform", "test_table", 123.45) + + # Получаем offset + offset = offset_table.get_offset("test_transform", "test_table") + assert offset == 123.45 + + +def test_offset_table_update_existing(dbconn: DBConn): + """Тест обновления существующего offset'а""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Создаем первый offset + offset_table.update_offset("test_transform", "test_table", 100.0) + + # Обновляем его + offset_table.update_offset("test_transform", "test_table", 200.0) + + # Проверяем, что offset обновился + offset = offset_table.get_offset("test_transform", "test_table") + assert offset == 200.0 + + +def test_offset_table_multiple_inputs(dbconn: DBConn): + """Тест работы с несколькими входными таблицами""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Создаем offset'ы для разных входных таблиц одной трансформации + offset_table.update_offset("test_transform", "table1", 100.0) + offset_table.update_offset("test_transform", "table2", 200.0) + + # Проверяем, что каждая таблица имеет свой offset + assert offset_table.get_offset("test_transform", "table1") == 100.0 + assert offset_table.get_offset("test_transform", "table2") == 200.0 + + +def test_offset_table_bulk_update(dbconn: DBConn): + """Тест bulk обновления offset'ов""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Обновляем несколько offset'ов за раз + offsets = { + ("transform1", "table1"): 100.0, + ("transform1", "table2"): 200.0, + ("transform2", "table1"): 300.0, + } + offset_table.update_offsets_bulk(offsets) + + # Проверяем, что все offset'ы установлены + assert offset_table.get_offset("transform1", "table1") == 100.0 + assert offset_table.get_offset("transform1", "table2") == 200.0 + assert offset_table.get_offset("transform2", "table1") == 300.0 + + +def test_offset_table_bulk_update_existing(dbconn: DBConn): + """Тест bulk обновления существующих offset'ов""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Создаем начальные offset'ы + offset_table.update_offset("transform1", "table1", 100.0) + offset_table.update_offset("transform1", "table2", 200.0) + + # Обновляем их через bulk + offsets = { + ("transform1", "table1"): 150.0, + ("transform1", "table2"): 250.0, + } + offset_table.update_offsets_bulk(offsets) + + # Проверяем, что offset'ы обновились + assert offset_table.get_offset("transform1", "table1") == 150.0 + assert offset_table.get_offset("transform1", "table2") == 250.0 + + +def test_offset_table_reset_single_table(dbconn: DBConn): + """Тест сброса offset'а для одной входной таблицы""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Создаем offset'ы + offset_table.update_offset("test_transform", "table1", 100.0) + offset_table.update_offset("test_transform", "table2", 200.0) + + # Сбрасываем offset для одной таблицы + offset_table.reset_offset("test_transform", "table1") + + # Проверяем + assert offset_table.get_offset("test_transform", "table1") is None + assert offset_table.get_offset("test_transform", "table2") == 200.0 + + +def test_offset_table_reset_all_tables(dbconn: DBConn): + """Тест сброса всех offset'ов для трансформации""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Создаем offset'ы + offset_table.update_offset("test_transform", "table1", 100.0) + offset_table.update_offset("test_transform", "table2", 200.0) + + # Сбрасываем все offset'ы для трансформации + offset_table.reset_offset("test_transform") + + # Проверяем + assert offset_table.get_offset("test_transform", "table1") is None + assert offset_table.get_offset("test_transform", "table2") is None + + +def test_offset_table_get_statistics(dbconn: DBConn): + """Тест получения статистики по offset'ам""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Создаем offset'ы + now = time.time() + offset_table.update_offset("transform1", "table1", now - 100) + offset_table.update_offset("transform1", "table2", now - 200) + offset_table.update_offset("transform2", "table1", now - 300) + + # Получаем статистику для всех трансформаций + stats = offset_table.get_statistics() + assert len(stats) == 3 + + # Проверяем структуру + for stat in stats: + assert "transformation_id" in stat + assert "input_table_name" in stat + assert "update_ts_offset" in stat + assert "offset_age_seconds" in stat + assert isinstance(stat["offset_age_seconds"], float) + assert stat["offset_age_seconds"] >= 0 + + # Получаем статистику для конкретной трансформации + stats_t1 = offset_table.get_statistics("transform1") + assert len(stats_t1) == 2 + assert all(s["transformation_id"] == "transform1" for s in stats_t1) + + +def test_offset_table_get_offset_count(dbconn: DBConn): + """Тест подсчета количества offset'ов""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Изначально offset'ов нет + assert offset_table.get_offset_count() == 0 + + # Добавляем offset'ы + offset_table.update_offset("transform1", "table1", 100.0) + assert offset_table.get_offset_count() == 1 + + offset_table.update_offset("transform1", "table2", 200.0) + assert offset_table.get_offset_count() == 2 + + offset_table.update_offset("transform2", "table1", 300.0) + assert offset_table.get_offset_count() == 3 + + # Обновление существующего offset'а не увеличивает счетчик + offset_table.update_offset("transform1", "table1", 150.0) + assert offset_table.get_offset_count() == 3 + + # Сброс offset'а уменьшает счетчик + offset_table.reset_offset("transform1", "table1") + assert offset_table.get_offset_count() == 2 + + +def test_offset_table_multiple_transformations(dbconn: DBConn): + """Тест изоляции offset'ов между трансформациями""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Создаем offset'ы для разных трансформаций с одинаковыми именами таблиц + offset_table.update_offset("transform1", "common_table", 100.0) + offset_table.update_offset("transform2", "common_table", 200.0) + offset_table.update_offset("transform3", "common_table", 300.0) + + # Проверяем, что offset'ы не смешиваются + assert offset_table.get_offset("transform1", "common_table") == 100.0 + assert offset_table.get_offset("transform2", "common_table") == 200.0 + assert offset_table.get_offset("transform3", "common_table") == 300.0 + + +def test_offset_table_get_offsets_for_transformation(dbconn: DBConn): + """Тест получения всех offset'ов для трансформации одним запросом""" + offset_table = TransformInputOffsetTable(dbconn, create_table=True) + + # Создаем offset'ы для одной трансформации с разными входными таблицами + offset_table.update_offset("test_transform", "table1", 100.0) + offset_table.update_offset("test_transform", "table2", 200.0) + offset_table.update_offset("test_transform", "table3", 300.0) + + # Создаем offset'ы для другой трансформации (не должны попасть в результат) + offset_table.update_offset("other_transform", "table1", 999.0) + + # Получаем все offset'ы для test_transform одним запросом + offsets = offset_table.get_offsets_for_transformation("test_transform") + + # Проверяем результат + assert offsets == { + "table1": 100.0, + "table2": 200.0, + "table3": 300.0, + } + + # Проверяем для несуществующей трансформации + empty_offsets = offset_table.get_offsets_for_transformation("nonexistent") + assert empty_offsets == {} diff --git a/tests/test_reverse_join_with_fk_in_meta.py b/tests/test_reverse_join_with_fk_in_meta.py new file mode 100644 index 00000000..8e959ba8 --- /dev/null +++ b/tests/test_reverse_join_with_fk_in_meta.py @@ -0,0 +1,187 @@ +""" +Тест для reverse_join когда FK находятся в meta table, а не в data table. + +Проблема 1: +_build_reverse_join_cte всегда делает JOIN с primary_data_tbl, +но если FK есть в primary_meta_tbl, то JOIN с data table не нужен. + +Проблема 2: +Строки 956-963 делают literal(None) для колонок которые не в primary_data_cols, +но если колонка в meta_cols справочной таблицы, нужно брать из tbl.c[k]. + +Сценарий: +- posts.meta содержит [id, author_id, update_ts, delete_ts] (FK в META!) +- posts.data содержит [id, content] (минимум) +- users.meta содержит [id, name, update_ts, delete_ts] +- При изменении user должен сработать reverse_join для posts +""" + +import time + +import pandas as pd +from sqlalchemy import Column, String + +from datapipe.compute import ComputeInput +from datapipe.datatable import DataStore +from datapipe.step.batch_transform import BatchTransformStep +from datapipe.store.database import DBConn, MetaKey, TableStoreDB + + +def test_reverse_join_fk_in_meta_simple(dbconn: DBConn): + """ + Простой тест: posts с author_id в meta, users с name в meta. + """ + ds = DataStore(dbconn, create_meta_table=True) + + # 1. Posts таблица: author_id в МЕТА через MetaKey + post_store = TableStoreDB( + dbconn, + "posts", + [ + Column("id", String, primary_key=True), + Column("author_id", String, MetaKey()), # FK в meta table! + Column("content", String), # В data table + ], + create_table=True, + ) + post_dt = ds.create_table("posts", post_store) + + # 2. Users таблица: name в мета через MetaKey + user_store = TableStoreDB( + dbconn, + "users", + [ + Column("id", String, primary_key=True), + Column("name", String, MetaKey()), # В meta table! + Column("email", String), # В data table + ], + create_table=True, + ) + user_dt = ds.create_table("users", user_store) + + # 3. Output таблица + output_store = TableStoreDB( + dbconn, + "posts_with_author", + [ + Column("id", String, primary_key=True), + Column("author_name", String), + Column("content", String), + ], + create_table=True, + ) + output_dt = ds.create_table("posts_with_author", output_store) + + # 4. Трансформация с join_keys + def transform_func(post_df, user_df): + result = post_df.merge( + user_df[["id", "name"]], + left_on="author_id", + right_on="id", + how="left", + suffixes=("", "_user"), + ) + return result[["id", "content", "name"]].rename(columns={"name": "author_name"}) + + step = BatchTransformStep( + ds=ds, + name="test_reverse_join_simple", + func=transform_func, + input_dts=[ + ComputeInput(dt=post_dt, join_type="full"), + # КРИТИЧНО: join_keys связывают posts.meta.author_id ← users.meta.id + ComputeInput( + dt=user_dt, + join_type="full", + join_keys={"author_id": "id"}, # author_id в posts.META! + ), + ], + output_dts=[output_dt], + transform_keys=["id"], + use_offset_optimization=True, + chunk_size=10, + ) + + # === НАЧАЛЬНЫЕ ДАННЫЕ === + + t1 = time.time() + + # Создать пользователей + users_df = pd.DataFrame( + [ + {"id": "u1", "name": "Alice", "email": "alice@example.com"}, + {"id": "u2", "name": "Bob", "email": "bob@example.com"}, + ] + ) + user_dt.store_chunk(users_df, now=t1) + + # Создать посты (author_id в meta через idx_columns) + posts_df = pd.DataFrame( + [ + {"id": "p1", "author_id": "u1", "content": "Hello World"}, + {"id": "p2", "author_id": "u2", "content": "Test Post"}, + ] + ) + post_dt.store_chunk(posts_df, now=t1) + + # === ПЕРВАЯ ОБРАБОТКА === + + time.sleep(0.01) + step.run_full(ds) + + output_data = output_dt.get_data() + assert len(output_data) == 2, f"Expected 2 records in output, got {len(output_data)}" + + p1 = output_data[output_data["id"] == "p1"].iloc[0] + assert p1["author_name"] == "Alice", f"p1.author_name should be Alice, got {p1['author_name']}" + + # === ИЗМЕНЕНИЕ USER (должен сработать reverse_join) === + + time.sleep(0.01) + t2 = time.time() + + # Изменить Alice → Alicia + user_dt.store_chunk( + pd.DataFrame([{"id": "u1", "name": "Alicia", "email": "alice@example.com"}]), + now=t2, + ) + + # === ВТОРАЯ ОБРАБОТКА === + + time.sleep(0.01) + + # Проверяем что reverse_join сработал (users.meta изменился → posts должны пересчитаться) + idx_count, idx_gen = step.get_full_process_ids(ds, run_config=None) + print(f"\n=== Записей для обработки после изменения user: {idx_count} ===") + + # Посмотрим на SQL запрос + sql_keys, sql = step._build_changed_idx_sql(ds, run_config=None) + print(f"\n=== SQL запрос для changed_idx ===") + print(sql) + + # Должен быть минимум p1 (зависит от u1 который изменился) + assert idx_count >= 1, ( + f"БАГ: должно быть минимум 1 запись для обработки (p1 зависит от u1), got {idx_count}. " + "Это означает что reverse_join не сработал!" + ) + + step.run_full(ds) + + # === ПРОВЕРКА === + + output_data_after = output_dt.get_data() + assert len(output_data_after) == 2, f"Expected 2 records, got {len(output_data_after)}" + + # p1: author=u1 (Alicia) - должно обновиться через reverse_join + p1_after = output_data_after[output_data_after["id"] == "p1"].iloc[0] + assert p1_after["author_name"] == "Alicia", ( + f"БАГ: p1.author_name должно быть 'Alicia', got '{p1_after['author_name']}'. " + "Проблема 1: reverse_join делает JOIN с posts.data, хотя author_id в posts.meta. " + "Проблема 2: колонка 'name' из users.meta заменяется на literal(None)." + ) + + # p2: author=u2 (Bob) - без изменений + p2_after = output_data_after[output_data_after["id"] == "p2"].iloc[0] + assert p2_after["author_name"] == "Bob", f"p2.author_name should be Bob, got {p2_after['author_name']}" + + print("✓ Reverse_join корректно работает с FK в meta table")