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 5b634796..02104b00 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 5f4f553f..a5f9bf2f 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 @@ -165,6 +165,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 c07185ee..bb74a195 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,6 +37,7 @@ from datapipe.compute import ComputeInput from datapipe.datatable import DataStore +logger = logging.getLogger("datapipe.meta.sql_meta") TABLE_META_SCHEMA: List[sa.Column] = [ sa.Column("hash", sa.Integer), @@ -719,7 +721,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"], @@ -728,15 +730,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, ) @@ -744,7 +758,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", ) @@ -769,12 +783,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( @@ -809,4 +825,581 @@ 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) + + +# Обратная совместимость: алиас для старой версии +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, + ) + + +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 + + # 1. Получить все offset'ы одним запросом для избежания N+1 + offsets = offset_table.get_offsets_for_transformation(transformation_id) + # Для таблиц без offset используем 0.0 (обрабатываем все данные) + for inp in input_dts: + if inp.dt.name not in offsets: + offsets[inp.dt.name] = 0.0 + + # 2. Построить CTE для каждой входной таблицы с фильтром по offset + # Для таблиц с join_keys нужен обратный JOIN к основной таблице + changed_ctes = [] + + # Сначала находим "основную" таблицу - первую без join_keys + primary_inp = None + for inp in input_dts: + if not inp.join_keys: + primary_inp = inp + break + + for inp in input_dts: + tbl = inp.dt.meta_table.sql_table + + # Разделяем ключи на те, что есть в meta table, и те, что нужны из data 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] + keys_in_data_only = [k for k in all_select_keys if k not in meta_cols] + + if len(keys_in_meta) == 0: + continue + + offset = offsets[inp.dt.name] + + # ОБРАТНЫЙ JOIN для справочных таблиц с join_keys + # Когда изменяется справочная таблица, нужно найти все записи основной таблицы, + # которые на нее ссылаются + if inp.join_keys and primary_inp and hasattr(primary_inp.dt.table_store, 'data_table'): + # Справочная таблица изменилась - нужен обратный JOIN к основной + primary_data_tbl = primary_inp.dt.table_store.data_table + + # Строим SELECT для всех колонок из all_select_keys основной таблицы + primary_data_cols = [c.name for c in primary_data_tbl.columns] + select_cols = [ + primary_data_tbl.c[k] if k in primary_data_cols else sa.literal(None).label(k) + for k in all_select_keys + ] + + # Обратный JOIN: primary_table.join_key = reference_table.id + # Например: posts.user_id = profiles.id + # inp.join_keys = {'user_id': 'id'} означает: + # 'user_id' - колонка в основной таблице (posts) + # 'id' - колонка в справочной таблице (profiles) + join_conditions = [] + for primary_col, ref_col in inp.join_keys.items(): + if primary_col in primary_data_cols and ref_col in meta_cols: + join_conditions.append(primary_data_tbl.c[primary_col] == tbl.c[ref_col]) + + if len(join_conditions) == 0: + # Не можем построить JOIN - пропускаем эту таблицу + continue + + join_condition = sa.and_(*join_conditions) if len(join_conditions) > 1 else join_conditions[0] + + # SELECT primary_cols FROM reference_meta + # JOIN primary_data ON primary.join_key = reference.id + # WHERE reference.update_ts > offset + changed_sql = sa.select(*select_cols).select_from( + tbl.join(primary_data_tbl, join_condition) + ).where( + sa.or_( + tbl.c.update_ts > offset, + sa.and_( + tbl.c.delete_ts.isnot(None), + tbl.c.delete_ts > offset + ) + ) + ) + + # Применить filters и group by + changed_sql = sql_apply_filters_idx_to_subquery(changed_sql, all_select_keys, filters_idx) + # run_config фильтры применяются к справочной таблице + changed_sql = sql_apply_runconfig_filter(changed_sql, tbl, inp.dt.primary_keys, run_config) + + if len(select_cols) > 0: + changed_sql = changed_sql.group_by(*select_cols) + + changed_ctes.append(changed_sql.cte(name=f"{inp.dt.name}_changes")) + continue + + # Если все ключи есть в meta table - используем простой запрос + if len(keys_in_data_only) == 0: + select_cols = [sa.column(k) for k in keys_in_meta] + + # SELECT keys FROM input_meta WHERE update_ts > offset OR delete_ts > offset + changed_sql = sa.select(*select_cols).select_from(tbl).where( + sa.or_( + tbl.c.update_ts > offset, + sa.and_( + tbl.c.delete_ts.isnot(None), + tbl.c.delete_ts > offset + ) + ) + ) + + # Применить filters_idx и run_config + changed_sql = sql_apply_filters_idx_to_subquery(changed_sql, keys_in_meta, filters_idx) + changed_sql = sql_apply_runconfig_filter(changed_sql, tbl, inp.dt.primary_keys, run_config) + + if len(select_cols) > 0: + changed_sql = changed_sql.group_by(*select_cols) + else: + # Есть колонки только в data table - нужен JOIN с data table + # Проверяем что у table_store есть data_table (для TableStoreDB) + if not hasattr(inp.dt.table_store, 'data_table'): + # Fallback: если нет data_table, используем только meta keys + select_cols = [sa.column(k) for k in keys_in_meta] + changed_sql = sa.select(*select_cols).select_from(tbl).where( + sa.or_( + tbl.c.update_ts > offset, + sa.and_( + tbl.c.delete_ts.isnot(None), + tbl.c.delete_ts > offset + ) + ) + ) + changed_sql = sql_apply_filters_idx_to_subquery(changed_sql, keys_in_meta, filters_idx) + changed_sql = sql_apply_runconfig_filter(changed_sql, tbl, inp.dt.primary_keys, run_config) + if len(select_cols) > 0: + changed_sql = changed_sql.group_by(*select_cols) + else: + # JOIN meta table с data table для получения дополнительных колонок + data_tbl = inp.dt.table_store.data_table + + # Проверяем какие дополнительные колонки действительно есть в data table + 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] + + if len(keys_in_data_available) == 0: + # Fallback: если нужных колонок нет в data table, используем только meta keys + select_cols = [sa.column(k) for k in keys_in_meta] + changed_sql = sa.select(*select_cols).select_from(tbl).where( + sa.or_( + tbl.c.update_ts > offset, + sa.and_( + tbl.c.delete_ts.isnot(None), + tbl.c.delete_ts > offset + ) + ) + ) + changed_sql = sql_apply_filters_idx_to_subquery(changed_sql, keys_in_meta, filters_idx) + changed_sql = sql_apply_runconfig_filter(changed_sql, tbl, inp.dt.primary_keys, run_config) + if len(select_cols) > 0: + changed_sql = changed_sql.group_by(*select_cols) + changed_ctes.append(changed_sql.cte(name=f"{inp.dt.name}_changes")) + continue + + # SELECT meta_keys, data_keys FROM meta JOIN data ON primary_keys + # WHERE update_ts > offset OR delete_ts > offset + select_cols = [tbl.c[k] for k in keys_in_meta] + [data_tbl.c[k] for k in keys_in_data_available] + + # Строим JOIN condition по primary keys + if len(inp.dt.primary_keys) == 1: + join_condition = tbl.c[inp.dt.primary_keys[0]] == data_tbl.c[inp.dt.primary_keys[0]] + else: + join_condition = sa.and_(*[ + tbl.c[pk] == data_tbl.c[pk] for pk in inp.dt.primary_keys + ]) + + changed_sql = sa.select(*select_cols).select_from( + tbl.join(data_tbl, join_condition) + ).where( + sa.or_( + tbl.c.update_ts > offset, + sa.and_( + tbl.c.delete_ts.isnot(None), + tbl.c.delete_ts > offset + ) + ) + ) + + # Применить filters_idx и run_config + all_keys = keys_in_meta + keys_in_data_available + changed_sql = sql_apply_filters_idx_to_subquery(changed_sql, all_keys, filters_idx) + changed_sql = sql_apply_runconfig_filter(changed_sql, tbl, inp.dt.primary_keys, run_config) + + if len(select_cols) > 0: + changed_sql = changed_sql.group_by(*select_cols) + + changed_ctes.append(changed_sql.cte(name=f"{inp.dt.name}_changes")) + + # 3. Получить записи с ошибками из TransformMetaTable + # Важно: error_records должен иметь все колонки из all_select_keys для UNION + # Для additional_columns используем NULL, так как их нет в transform meta table + tr_tbl = meta_table.sql_table + error_select_cols: List[Any] = [sa.column(k) for k in transform_keys] + [ + sa.literal(None).label(k) for k in additional_columns + ] + 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]) + + error_records_cte = error_records_sql.cte(name="error_records") + + # 4. Объединить все изменения и ошибки через UNION + 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: + # UNION всех изменений и ошибок + # Важно: UNION должен включать все колонки из all_select_keys + # Для отсутствующих колонок используем NULL + union_parts = [] + for cte in changed_ctes: + # Для каждой колонки из all_select_keys: берем из CTE если есть, иначе NULL + 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) + + # 5. Применить сортировку + # Нам нужно join с transform meta для получения priority + union_cte = union_sql.cte(name="changed_union") + + 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]) + + # Используем `out` для консистентности с v1 + # Важно: Включаем все колонки (transform_keys + additional_columns) + 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) + ) + + if order_by is None: + out = out.order_by( + tr_tbl.c.priority.desc().nullslast(), + *[union_cte.c[k] for k in transform_keys], + ) + else: + if order == "desc": + out = out.order_by( + *[sa.desc(union_cte.c[k]) for k in order_by], + tr_tbl.c.priority.desc().nullslast(), + ) + elif order == "asc": + out = out.order_by( + *[sa.asc(union_cte.c[k]) for k in order_by], + tr_tbl.c.priority.desc().nullslast(), + ) + + 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..412a30f1 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() - for k, v in extra_filters.items(): - df[k] = v + 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] - yield cast(IndexDF, df) + for k, v in extra_filters.items(): + df[k] = v + + 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,42 @@ def store_batch_result( self.meta_table.mark_rows_processed_success(idx, process_ts=process_ts, run_config=run_config) + # НОВОЕ: Обновление offset'ов для каждой входной таблицы (Phase 3) + # Обновляем offset'ы всегда при успешной обработке, независимо от use_offset_optimization + # Это позволяет накапливать offset'ы для будущего использования + 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) + + offsets_to_update = {} + + 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: + offsets_to_update[(self.get_name(), inp.dt.name)] = max_update_ts + + # Batch update всех offset'ов за одну транзакцию + if offsets_to_update: + try: + ds.offset_table.update_offsets_bulk(offsets_to_update) + 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)" + ) + return changes def store_batch_err( @@ -405,7 +648,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, @@ -584,6 +875,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 +885,7 @@ def __init__( transform_keys=transform_keys, chunk_size=chunk_size, labels=labels, + use_offset_optimization=use_offset_optimization, ) self.func = func @@ -632,12 +925,13 @@ def pipeline_input_to_compute_input(self, ds: DataStore, catalog: Catalog, input 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") @@ -681,6 +975,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 +989,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 9fb62685..76baa2f1 100644 --- a/datapipe/types.py +++ b/datapipe/types.py @@ -9,6 +9,7 @@ Dict, List, NewType, + Optional, Set, Tuple, Type, @@ -60,6 +61,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 diff --git a/tests/conftest.py b/tests/conftest.py index 0f545247..017149fc 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: 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..b1d15589 --- /dev/null +++ b/tests/test_build_changed_idx_sql_v2.py @@ -0,0 +1,276 @@ +""" +Тесты для 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 компилируется + assert transform_keys == ["id"] + 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_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_joinspec.py b/tests/test_offset_joinspec.py new file mode 100644 index 00000000..0b6dc857 --- /dev/null +++ b/tests/test_offset_joinspec.py @@ -0,0 +1,165 @@ +""" +Тест для проверки что 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)!") 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_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 == {}