feat: implement strict mode for max_records_per_run to prevent offset overshoot with composite keys

.github/workflows/test.yaml

@@ -59,8 +59,10 @@ jobs:
 
           - python-version: "3.12"
             test-db-env: "postgres"
+            pip-extra: '"pandas<2.2"'
           - python-version: "3.12"
             test-db-env: "sqlite"
+            pip-extra: '"pandas<2.2"'
 
     services:
       # Label used to access the service container

datapipe/index_metadata.py

@@ -0,0 +1,133 @@
+"""
+Wrapper для IndexDF с дополнительными метаданными.
+
+Используется для передачи исходного changed_idx (с полным PK) вместе
+со сгруппированным idx (по transform_keys) для поддержки strict mode
+при использовании max_records_per_run.
+"""
+
+from typing import Optional
+
+from datapipe.types import IndexDF
+
+
+class IndexDFWithMetadata:
+    """
+    Wrapper для IndexDF с дополнительными метаданными.
+
+    Позволяет передавать исходный changed_idx (с полным PK) вместе
+    со сгруппированным idx (по transform_keys) для поддержки strict mode
+    при использовании max_records_per_run.
+
+    Ведёт себя как обычный IndexDF для обратной совместимости через делегирование.
+
+    Example:
+        >>> idx_grouped = data_to_index(chunk_df, transform_keys)
+        >>> idx_with_meta = IndexDFWithMetadata(
+        ...     idx=idx_grouped,
+        ...     changed_idx=chunk_df  # Исходный с полным PK
+        ... )
+        >>> len(idx_with_meta)  # Работает как обычный DataFrame
+        3
+        >>> idx_with_meta.has_changed_idx()
+        True
+    """
+
+    def __init__(
+        self,
+        idx: IndexDF,
+        changed_idx: Optional[IndexDF] = None
+    ):
+        """
+        Args:
+            idx: Основной индекс (обычно сгруппированный по transform_keys)
+            changed_idx: Исходный индекс с полным PK из changed_idx query
+                        (опционально, используется для strict mode)
+        """
+        self._idx = idx
+        self._changed_idx = changed_idx
+
+    @property
+    def idx(self) -> IndexDF:
+        """Основной индекс (сгруппированный по transform_keys)"""
+        return self._idx
+
+    @property
+    def changed_idx(self) -> Optional[IndexDF]:
+        """Исходный индекс с полным PK (если есть)"""
+        return self._changed_idx
+
+    def has_changed_idx(self) -> bool:
+        """Проверка наличия changed_idx"""
+        return self._changed_idx is not None
+
+    # Делегирование к pandas DataFrame для обратной совместимости
+    def __getattr__(self, name):
+        """Делегируем все неизвестные атрибуты к idx"""
+        return getattr(self._idx, name)
+
+    def __getitem__(self, key):
+        """Поддержка индексации df[key]"""
+        return self._idx[key]
+
+    def __len__(self):
+        """len(idx_with_meta) возвращает длину основного индекса"""
+        return len(self._idx)
+
+    def __repr__(self):
+        return (
+            f"IndexDFWithMetadata("
+            f"idx={len(self._idx)} rows, "
+            f"has_changed_idx={self.has_changed_idx()}"
+            f")"
+        )
+
+    def __str__(self):
+        return str(self._idx)
+
+    # Поддержка итерации
+    def __iter__(self):
+        return iter(self._idx)
+
+    # Поддержка pandas операций
+    @property
+    def columns(self):
+        return self._idx.columns
+
+    @property
+    def shape(self):
+        return self._idx.shape
+
+    def to_records(self, *args, **kwargs):
+        """Делегируем to_records к основному DataFrame"""
+        return self._idx.to_records(*args, **kwargs)
+
+
+def unwrap_idx(idx_or_wrapped) -> tuple[IndexDF, Optional[IndexDF]]:
+    """
+    Извлекает idx и changed_idx из обёртки или возвращает как есть.
+
+    Этот хелпер позволяет функциям принимать как обычный IndexDF,
+    так и IndexDFWithMetadata без явной проверки типа.
+
+    Args:
+        idx_or_wrapped: IndexDF или IndexDFWithMetadata
+
+    Returns:
+        Кортеж (idx, changed_idx):
+            - idx: основной индекс (IndexDF)
+            - changed_idx: опциональный исходный индекс с полным PK
+
+    Example:
+        >>> idx, changed_idx = unwrap_idx(some_idx)
+        >>> if changed_idx is not None:
+        ...     # Используем strict mode
+        ...     data = dt.get_data(changed_idx)
+        >>> else:
+        ...     # Legacy mode
+        ...     data = dt.get_data(idx)
+    """
+    if isinstance(idx_or_wrapped, IndexDFWithMetadata):
+        return idx_or_wrapped.idx, idx_or_wrapped.changed_idx
+    else:
+        return idx_or_wrapped, None

datapipe/meta/sql_meta.py

@@ -1642,22 +1642,37 @@ def _build_error_records_cte(
     transform_keys: List[str],
     all_select_keys: List[str],
     filters_idx: Optional[IndexDF],
+    input_dts: Optional[List["ComputeInput"]] = None,
 ) -> Any:
     """
     Строит CTE для записей с ошибками из TransformMetaTable.
 
     Returns:
         CTE object для error_records
     """
+    # Получаем типы колонок из первой входной таблицы для дополнительных колонок
+    column_types = {}
+    if input_dts and len(input_dts) > 0:
+        first_input = input_dts[0]
+        if hasattr(first_input.dt, 'meta_table') and first_input.dt.meta_table:
+            # Используем sql_table.c для получения колонок
+            sql_table = first_input.dt.meta_table.sql_table
+            for col_name in all_select_keys:
+                if col_name in sql_table.c:
+                    column_types[col_name] = sql_table.c[col_name].type
+
     tr_tbl = meta_table.sql_table
     error_select_cols: List[Any] = []
     for k in all_select_keys:
         if k in transform_keys:
             error_select_cols.append(sa.column(k))
         else:
-            # Для дополнительных колонок (включая update_ts) используем NULL
+            # Для дополнительных колонок (включая update_ts) используем NULL с правильным типом
             if k == 'update_ts':
                 error_select_cols.append(sa.cast(sa.literal(None), sa.Float).label(k))
+            elif k in column_types:
+                # Используем тип из входной таблицы
+                error_select_cols.append(sa.cast(sa.literal(None), column_types[k]).label(k))
             else:
                 error_select_cols.append(sa.literal(None).label(k))
 
@@ -1923,6 +1938,7 @@ def build_changed_idx_sql_v2(
         transform_keys=transform_keys,
         all_select_keys=all_select_keys,
         filters_idx=filters_idx,
+        input_dts=input_dts,
     )
 
     # 4. Объединить через UNION или CROSS JOIN