tracepipe 0.2.0__py3-none-any.whl

tracepipe/storage/row_identity.py

@@ -0,0 +1,217 @@
+# tracepipe/storage/row_identity.py
+"""
+Row identity tracking for pandas DataFrames.
+
+Uses: Registry + Hidden Column fallback.
+"""
+
+import warnings
+import weakref
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+
+from ..core import TracePipeConfig
+
+_TRACEPIPE_ROW_ID_COL = "__tracepipe_row_id__"
+
+
+class PandasRowIdentity:
+    """
+    Hybrid row identity tracking for pandas DataFrames.
+
+    Implements: RowIdentityStrategy protocol
+
+    Handles:
+    - Standard operations (filter, sort, copy)
+    - reset_index(drop=True)
+    - Duplicate indices (with warning)
+    - Chained operations
+
+    Future alternatives:
+    - PolarsRowIdentity: Uses Polars row numbers and lazy evaluation
+    - SparkRowIdentity: Uses monotonically_increasing_id() or RDD zipWithIndex
+    """
+
+    def __init__(self, config: TracePipeConfig):
+        self.config = config
+        self._registry: dict[int, pd.Series] = {}
+        self._df_refs: weakref.WeakValueDictionary = weakref.WeakValueDictionary()
+        self._next_row_id: int = 0
+
+    def register(
+        self,
+        df: pd.DataFrame,
+        row_ids: Optional[pd.Series] = None,
+        warn_duplicate_index: bool = True,
+    ) -> pd.Series:
+        """
+        Register a DataFrame and assign row IDs.
+
+        Args:
+            df: DataFrame to register
+            row_ids: Optional pre-assigned IDs (for propagation)
+            warn_duplicate_index: Warn if index has duplicates
+
+        Returns:
+            Series of row IDs aligned to df.index
+        """
+        # Check for duplicate index
+        if warn_duplicate_index and self.config.warn_on_duplicate_index:
+            if df.index.has_duplicates:
+                warnings.warn(
+                    "TracePipe: DataFrame has duplicate index values. "
+                    "Row identity may be ambiguous for duplicates.",
+                    UserWarning,
+                )
+
+        if row_ids is None:
+            # Generate new sequential IDs
+            new_ids = list(range(self._next_row_id, self._next_row_id + len(df)))
+            self._next_row_id += len(df)
+            row_ids = pd.Series(new_ids, index=df.index, dtype="int64")
+        else:
+            # Ensure alignment
+            if not row_ids.index.equals(df.index):
+                row_ids = row_ids.copy()
+                row_ids.index = df.index
+
+        obj_id = id(df)
+        self._registry[obj_id] = row_ids
+        self._df_refs[obj_id] = df
+
+        # Optionally embed in DataFrame
+        if self.config.use_hidden_column:
+            df[_TRACEPIPE_ROW_ID_COL] = row_ids.values
+
+        return row_ids
+
+    def get_ids(self, df: pd.DataFrame) -> Optional[pd.Series]:
+        """Get row IDs for a DataFrame."""
+        # 1. Try registry (fast path)
+        obj_id = id(df)
+        if obj_id in self._registry:
+            stored = self._registry[obj_id]
+            # Verify alignment still valid
+            if len(stored) == len(df) and stored.index.equals(df.index):
+                return stored
+
+        # 2. Try hidden column (fallback)
+        if _TRACEPIPE_ROW_ID_COL in df.columns:
+            row_ids = df[_TRACEPIPE_ROW_ID_COL].copy()
+            row_ids.index = df.index
+            # Re-register for future lookups
+            self._registry[obj_id] = row_ids
+            self._df_refs[obj_id] = df
+            return row_ids
+
+        # 3. Not tracked
+        return None
+
+    def propagate(self, source_df: pd.DataFrame, result_df: pd.DataFrame) -> Optional[pd.Series]:
+        """
+        Propagate row IDs from source to result DataFrame.
+
+        Handles:
+        - Filtering (fewer rows)
+        - Reordering (same rows, different order)
+        - Mixed operations
+        """
+        source_ids = self.get_ids(source_df)
+        if source_ids is None:
+            return None
+
+        if result_df is source_df:
+            return source_ids
+
+        try:
+            if result_df.index.equals(source_df.index):
+                # Same index - direct copy
+                result_ids = source_ids.copy()
+            elif result_df.index.isin(source_df.index).all():
+                # Result index is subset/reorder of source
+                result_ids = source_ids.loc[result_df.index].copy()
+            else:
+                # Partial overlap or new indices
+                result_ids = source_ids.reindex(result_df.index)
+                # Rows not in source get new IDs
+                new_mask = result_ids.isna()
+                if new_mask.any():
+                    new_count = new_mask.sum()
+                    new_row_ids = list(range(self._next_row_id, self._next_row_id + new_count))
+                    self._next_row_id += new_count
+                    result_ids.loc[new_mask] = new_row_ids
+                result_ids = result_ids.astype("int64")
+        except Exception:
+            # Fallback: positional alignment
+            if len(result_df) <= len(source_df):
+                result_ids = pd.Series(
+                    source_ids.values[: len(result_df)], index=result_df.index, dtype="int64"
+                )
+            else:
+                # Result is larger - assign new IDs to extras
+                base_ids = list(source_ids.values)
+                extra_count = len(result_df) - len(source_df)
+                extra_ids = list(range(self._next_row_id, self._next_row_id + extra_count))
+                self._next_row_id += extra_count
+                result_ids = pd.Series(base_ids + extra_ids, index=result_df.index, dtype="int64")
+
+        return self.register(result_df, result_ids, warn_duplicate_index=False)
+
+    def realign_for_reset_index(
+        self, original_df: pd.DataFrame, new_df: pd.DataFrame
+    ) -> Optional[pd.Series]:
+        """Handle reset_index(drop=True) which changes index."""
+        old_ids = self.get_ids(original_df)
+        if old_ids is None:
+            return None
+
+        # Same values, new index
+        new_ids = pd.Series(old_ids.values, index=new_df.index, dtype="int64")
+        return self.register(new_df, new_ids, warn_duplicate_index=False)
+
+    def get_dropped_ids(self, source_df: pd.DataFrame, result_df: pd.DataFrame) -> np.ndarray:
+        """
+        Get row IDs that were dropped between source and result.
+
+        Uses numpy's setdiff1d for vectorized performance (~50x faster
+        than Python set operations for large DataFrames).
+
+        Returns:
+            numpy array of dropped row IDs (empty array if none dropped)
+        """
+        source_ids = self.get_ids(source_df)
+        result_ids = self.get_ids(result_df)
+
+        if source_ids is None:
+            return np.array([], dtype="int64")
+        if result_ids is None:
+            return np.asarray(source_ids.values, dtype="int64")
+
+        # Vectorized set difference - O(n log n) in C instead of O(n) in Python
+        return np.setdiff1d(source_ids.values, result_ids.values)
+
+    def strip_hidden_column(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Remove hidden column for export."""
+        if _TRACEPIPE_ROW_ID_COL in df.columns:
+            return df.drop(columns=[_TRACEPIPE_ROW_ID_COL])
+        return df
+
+    def cleanup(self) -> None:
+        """Remove stale entries."""
+        stale = [k for k in list(self._registry.keys()) if k not in self._df_refs]
+        for k in stale:
+            del self._registry[k]
+
+    def all_registered_ids(self) -> list[int]:
+        """
+        Get all row IDs that have ever been registered.
+
+        Returns:
+            List of all registered row IDs.
+        """
+        all_ids = set()
+        for row_ids in self._registry.values():
+            all_ids.update(row_ids.values.tolist())
+        return sorted(all_ids)

tracepipe/utils/__init__.py

@@ -0,0 +1,6 @@
+# tracepipe/utils/__init__.py
+"""Utility functions for TracePipe."""
+
+from .value_capture import capture_typed_value, values_equal
+
+__all__ = ["capture_typed_value", "values_equal"]

tracepipe/utils/value_capture.py

@@ -0,0 +1,137 @@
+# tracepipe/utils/value_capture.py
+"""
+Value capture and comparison utilities with complete NA handling.
+"""
+
+from typing import Any
+
+import numpy as np
+import pandas as pd
+
+# Interned type strings (avoid allocating same strings repeatedly)
+_TYPE_NULL = "null"
+_TYPE_BOOL = "bool"
+_TYPE_INT = "int"
+_TYPE_FLOAT = "float"
+_TYPE_STR = "str"
+_TYPE_DATETIME = "datetime"
+_TYPE_OTHER = "other"
+
+
+def capture_typed_value(value: Any) -> tuple[Any, str]:
+    """
+    Convert value to (python_native, type_string) for storage.
+
+    Handles:
+    - None, np.nan, pd.NA, pd.NaT
+    - numpy scalars (not JSON-serializable)
+    - Standard Python types
+    - Datetime types
+
+    Returns:
+        Tuple of (native_value, type_string)
+    """
+    # Handle all NA types first (pd.isna handles None, np.nan, pd.NA, pd.NaT)
+    try:
+        if pd.isna(value):
+            return None, _TYPE_NULL
+    except (ValueError, TypeError):
+        # pd.isna can fail on some types (e.g., lists)
+        pass
+
+    # numpy scalar -> Python native (CRITICAL for JSON serialization)
+    if hasattr(value, "item"):
+        try:
+            value = value.item()
+        except (ValueError, AttributeError):
+            pass
+
+    # Type mapping (order matters: bool before int)
+    if isinstance(value, bool):
+        return value, _TYPE_BOOL
+    elif isinstance(value, (int, np.integer)):
+        return int(value), _TYPE_INT
+    elif isinstance(value, (float, np.floating)):
+        return float(value), _TYPE_FLOAT
+    elif isinstance(value, str):
+        return value, _TYPE_STR
+    elif isinstance(value, (pd.Timestamp, np.datetime64)):
+        return str(value), _TYPE_DATETIME
+    else:
+        # Fallback: stringify for storage
+        return str(value), _TYPE_OTHER
+
+
+def values_equal(a: Any, b: Any) -> bool:
+    """
+    Compare two values, handling NA correctly.
+
+    pd.isna(x) == pd.isna(y) handles the case where both are NA.
+    """
+    try:
+        a_na = pd.isna(a)
+        b_na = pd.isna(b)
+
+        if a_na and b_na:
+            return True
+        if a_na or b_na:
+            return False
+        return a == b
+    except (ValueError, TypeError):
+        # Fallback for unhashable types
+        return str(a) == str(b)
+
+
+def find_changed_indices_vectorized(old_series: pd.Series, new_series: pd.Series) -> np.ndarray:
+    """
+    Find indices where values changed, using vectorized operations.
+
+    ~50-100x faster than row-by-row .loc[] access for large DataFrames.
+
+    Args:
+        old_series: Series of old values (must be aligned with new_series)
+        new_series: Series of new values
+
+    Returns:
+        Boolean mask array where True indicates the value changed
+    """
+    old_arr = old_series.values
+    new_arr = new_series.values
+    n = len(old_arr)
+
+    if n == 0:
+        return np.array([], dtype=bool)
+
+    # Vectorized NA detection
+    old_na = pd.isna(old_arr)
+    new_na = pd.isna(new_arr)
+
+    # NA status changed (one is NA, other isn't)
+    na_status_changed = old_na != new_na
+
+    # For non-NA values, check if they differ
+    # Handle mixed types safely by comparing element-by-element for non-NA
+    both_not_na = ~old_na & ~new_na
+
+    # Initialize values_differ as False
+    values_differ = np.zeros(n, dtype=bool)
+
+    # Only compare where both are non-NA
+    non_na_indices = np.where(both_not_na)[0]
+    if len(non_na_indices) > 0:
+        # Try vectorized comparison first (fast path for homogeneous arrays)
+        try:
+            with np.errstate(invalid="ignore"):
+                values_differ[non_na_indices] = old_arr[non_na_indices] != new_arr[non_na_indices]
+        except (TypeError, ValueError):
+            # Fallback for mixed types: element-by-element comparison
+            for i in non_na_indices:
+                try:
+                    values_differ[i] = old_arr[i] != new_arr[i]
+                except (TypeError, ValueError):
+                    # Different types that can't be compared - treat as different
+                    values_differ[i] = True
+
+    changed_mask = (both_not_na & values_differ) | na_status_changed
+
+    return changed_mask

tracepipe/visualization/__init__.py

@@ -0,0 +1,6 @@
+# tracepipe/visualization/__init__.py
+"""Visualization exports for TracePipe."""
+
+from .html_export import save
+
+__all__ = ["save"]