kontra 0.5.2__py3-none-any.whl

kontra/probes/compare.py

@@ -0,0 +1,400 @@
+# src/kontra/probes/compare.py
+"""
+Compare probe: Measure transformation effects between before/after datasets.
+
+This probe answers: "Did my transformation preserve rows and keys as expected?"
+
+It does NOT answer: whether the transformation is "correct".
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Union
+
+import polars as pl
+
+from kontra.api.compare import CompareResult
+
+
+def compare(
+    before: Union[pl.DataFrame, str],
+    after: Union[pl.DataFrame, str],
+    key: Union[str, List[str]],
+    *,
+    sample_limit: int = 5,
+    save: bool = False,
+) -> CompareResult:
+    """
+    Compare two datasets to measure transformation effects.
+
+    Answers: "Did my transformation preserve rows and keys as expected?"
+
+    Does NOT answer: whether the transformation is "correct".
+
+    This probe provides deterministic, structured measurements that allow
+    agents (and humans) to reason about transformation effects with confidence.
+
+    Args:
+        before: Dataset before transformation (DataFrame or path/URI)
+        after: Dataset after transformation (DataFrame or path/URI)
+        key: Column(s) to use as row identifier
+        sample_limit: Max samples per category (default 5)
+        save: Persist result to state backend (not yet implemented)
+
+    Returns:
+        CompareResult with row_stats, key_stats, change_stats,
+        column_stats, and bounded samples.
+
+    Example:
+        # Basic comparison
+        result = compare(raw_df, transformed_df, key="order_id")
+
+        # With composite key
+        result = compare(before, after, key=["customer_id", "date"])
+
+        # Check for issues
+        if result.duplicated_after > 0:
+            print(f"Warning: {result.duplicated_after} keys are duplicated")
+            print(f"Sample: {result.samples_duplicated_keys}")
+
+        # Get structured output for LLM
+        print(result.to_llm())
+    """
+    # Normalize key to list
+    if isinstance(key, str):
+        key = [key]
+
+    # Load data if paths provided
+    before_df = _load_data(before)
+    after_df = _load_data(after)
+
+    # Compute the comparison
+    result = _compute_compare(before_df, after_df, key, sample_limit)
+
+    # TODO: Implement save functionality
+    if save:
+        pass
+
+    return result
+
+
+def _load_data(data: Union[pl.DataFrame, str]) -> pl.DataFrame:
+    """
+    Load data from DataFrame or path/URI.
+
+    For MVP, only Polars DataFrames are fully supported.
+    File paths are loaded via Polars read functions.
+    """
+    if isinstance(data, pl.DataFrame):
+        return data
+
+    if isinstance(data, str):
+        # Simple file loading for MVP
+        if data.lower().endswith(".parquet"):
+            return pl.read_parquet(data)
+        elif data.lower().endswith(".csv"):
+            return pl.read_csv(data)
+        elif data.startswith("s3://"):
+            return pl.read_parquet(data)
+        else:
+            # Try parquet first, then CSV
+            try:
+                return pl.read_parquet(data)
+            except Exception:
+                return pl.read_csv(data)
+
+    raise ValueError(f"Unsupported data type: {type(data)}")
+
+
+def _compute_compare(
+    before: pl.DataFrame,
+    after: pl.DataFrame,
+    key: List[str],
+    sample_limit: int,
+) -> CompareResult:
+    """
+    Compute all comparison metrics between before and after datasets.
+
+    This is the core algorithm implementing the MVP schema.
+    """
+    # Validate key columns exist
+    for k in key:
+        if k not in before.columns:
+            raise ValueError(f"Key column '{k}' not found in before dataset")
+        if k not in after.columns:
+            raise ValueError(f"Key column '{k}' not found in after dataset")
+
+    # ==========================================================================
+    # 1. Row stats
+    # ==========================================================================
+    before_rows = len(before)
+    after_rows = len(after)
+    row_delta = after_rows - before_rows
+    row_ratio = after_rows / before_rows if before_rows > 0 else float('inf')
+
+    # ==========================================================================
+    # 2. Key stats
+    # ==========================================================================
+    before_keys = before.select(key).unique()
+    after_keys = after.select(key).unique()
+
+    unique_before = len(before_keys)
+    unique_after = len(after_keys)
+
+    # Keys in both (preserved)
+    preserved_keys = before_keys.join(after_keys, on=key, how="inner")
+    preserved = len(preserved_keys)
+
+    # Keys dropped (in before but not in after)
+    dropped = unique_before - preserved
+
+    # Keys added (in after but not in before)
+    added = unique_after - preserved
+
+    # Duplicated after: count of keys appearing >1x in after
+    # (This is key count, not row count)
+    after_key_counts = after.group_by(key).agg(pl.len().alias("_count"))
+    duplicated_keys_df = after_key_counts.filter(pl.col("_count") > 1)
+    duplicated_after = len(duplicated_keys_df)
+
+    # ==========================================================================
+    # 3. Change stats (for preserved keys only)
+    # ==========================================================================
+    # Join before and after on key to find matching rows
+    # Use suffix to disambiguate columns
+    non_key_cols_before = [c for c in before.columns if c not in key]
+    non_key_cols_after = [c for c in after.columns if c not in key]
+    common_non_key_cols = set(non_key_cols_before) & set(non_key_cols_after)
+
+    unchanged_rows = 0
+    changed_rows = 0
+
+    if preserved > 0 and common_non_key_cols:
+        # For each preserved key, compare values
+        # Join on key, suffix the after columns
+        merged = before.join(after, on=key, how="inner", suffix="_after")
+
+        # Build a change mask: True if any common non-key column differs
+        # Handle NULL comparison: NULL != value should be True
+        change_exprs = []
+        for col in common_non_key_cols:
+            after_col = f"{col}_after"
+            if after_col in merged.columns:
+                # Use ne_missing to treat NULLs as different from values
+                # but NULL == NULL as same
+                change_exprs.append(
+                    (pl.col(col).ne(pl.col(after_col))) |
+                    (pl.col(col).is_null() != pl.col(after_col).is_null())
+                )
+
+        if change_exprs:
+            # Combine all expressions with OR
+            combined_mask = change_exprs[0]
+            for expr in change_exprs[1:]:
+                combined_mask = combined_mask | expr
+
+            # Count changed and unchanged
+            changed_df = merged.filter(combined_mask)
+            changed_rows = len(changed_df)
+            unchanged_rows = len(merged) - changed_rows
+        else:
+            # No common columns to compare
+            unchanged_rows = len(merged)
+            changed_rows = 0
+    elif preserved > 0:
+        # No common non-key columns, so no changes possible
+        unchanged_rows = preserved
+        changed_rows = 0
+
+    # ==========================================================================
+    # 4. Column stats
+    # ==========================================================================
+    before_cols = set(before.columns)
+    after_cols = set(after.columns)
+
+    columns_added = sorted(after_cols - before_cols)
+    columns_removed = sorted(before_cols - after_cols)
+
+    # Modified columns: columns in both where at least one value differs
+    # Also compute modified_fraction
+    columns_modified = []
+    modified_fraction: Dict[str, float] = {}
+
+    if preserved > 0:
+        merged = before.join(after, on=key, how="inner", suffix="_after")
+        preserved_count = len(merged)
+
+        for col in sorted(common_non_key_cols):
+            after_col = f"{col}_after"
+            if after_col in merged.columns and preserved_count > 0:
+                # Count rows where this column changed
+                changed_count = len(merged.filter(
+                    (pl.col(col).ne(pl.col(after_col))) |
+                    (pl.col(col).is_null() != pl.col(after_col).is_null())
+                ))
+                if changed_count > 0:
+                    columns_modified.append(col)
+                    modified_fraction[col] = changed_count / preserved_count
+
+    # Nullability delta
+    nullability_delta: Dict[str, Dict[str, Optional[float]]] = {}
+
+    # For modified columns, compute before and after null rates
+    for col in columns_modified:
+        before_null = before[col].null_count() / before_rows if before_rows > 0 else 0.0
+        after_null = after[col].null_count() / after_rows if after_rows > 0 else 0.0
+        nullability_delta[col] = {"before": before_null, "after": after_null}
+
+    # For added columns, only after rate
+    for col in columns_added:
+        after_null = after[col].null_count() / after_rows if after_rows > 0 else 0.0
+        nullability_delta[col] = {"before": None, "after": after_null}
+
+    # ==========================================================================
+    # 5. Samples
+    # ==========================================================================
+
+    # Sample duplicated keys
+    samples_duplicated_keys = _extract_key_samples(
+        duplicated_keys_df.select(key),
+        key,
+        sample_limit
+    )
+
+    # Sample dropped keys (in before but not in after)
+    dropped_keys_df = before_keys.join(after_keys, on=key, how="anti")
+    samples_dropped_keys = _extract_key_samples(dropped_keys_df, key, sample_limit)
+
+    # Sample changed rows (with before/after values)
+    samples_changed_rows = _extract_changed_row_samples(
+        before, after, key, common_non_key_cols, sample_limit
+    )
+
+    # ==========================================================================
+    # Build result
+    # ==========================================================================
+    return CompareResult(
+        # Meta
+        before_rows=before_rows,
+        after_rows=after_rows,
+        key=key,
+        execution_tier="polars",
+
+        # Row stats
+        row_delta=row_delta,
+        row_ratio=row_ratio,
+
+        # Key stats
+        unique_before=unique_before,
+        unique_after=unique_after,
+        preserved=preserved,
+        dropped=dropped,
+        added=added,
+        duplicated_after=duplicated_after,
+
+        # Change stats
+        unchanged_rows=unchanged_rows,
+        changed_rows=changed_rows,
+
+        # Column stats
+        columns_added=columns_added,
+        columns_removed=columns_removed,
+        columns_modified=columns_modified,
+        modified_fraction=modified_fraction,
+        nullability_delta=nullability_delta,
+
+        # Samples
+        samples_duplicated_keys=samples_duplicated_keys,
+        samples_dropped_keys=samples_dropped_keys,
+        samples_changed_rows=samples_changed_rows,
+
+        # Config
+        sample_limit=sample_limit,
+    )
+
+
+def _extract_key_samples(
+    keys_df: pl.DataFrame,
+    key: List[str],
+    limit: int,
+) -> List[Any]:
+    """
+    Extract sample key values from a DataFrame.
+
+    Returns list of key values (single value if single key, tuple if composite).
+    """
+    if len(keys_df) == 0:
+        return []
+
+    samples = keys_df.head(limit)
+
+    if len(key) == 1:
+        # Single key - return list of values
+        return samples[key[0]].to_list()
+    else:
+        # Composite key - return list of dicts
+        return samples.to_dicts()
+
+
+def _extract_changed_row_samples(
+    before: pl.DataFrame,
+    after: pl.DataFrame,
+    key: List[str],
+    common_cols: set,
+    limit: int,
+) -> List[Dict[str, Any]]:
+    """
+    Extract sample changed rows with before/after values.
+
+    Returns list of dicts with key, before values, and after values
+    for columns that changed.
+    """
+    if not common_cols:
+        return []
+
+    # Join on key
+    merged = before.join(after, on=key, how="inner", suffix="_after")
+
+    if len(merged) == 0:
+        return []
+
+    samples = []
+    for row in merged.head(limit * 2).iter_rows(named=True):
+        # Check if any column changed
+        changes_before = {}
+        changes_after = {}
+        has_change = False
+
+        for col in common_cols:
+            after_col = f"{col}_after"
+            if after_col in row:
+                before_val = row[col]
+                after_val = row[after_col]
+
+                # Check for change (handle NULL)
+                is_changed = (before_val != after_val) or (
+                    (before_val is None) != (after_val is None)
+                )
+
+                if is_changed:
+                    has_change = True
+                    changes_before[col] = before_val
+                    changes_after[col] = after_val
+
+        if has_change:
+            # Extract key value(s)
+            if len(key) == 1:
+                key_val = row[key[0]]
+            else:
+                key_val = {k: row[k] for k in key}
+
+            samples.append({
+                "key": key_val,
+                "before": changes_before,
+                "after": changes_after,
+            })
+
+            if len(samples) >= limit:
+                break
+
+    return samples

kontra/probes/relationship.py

@@ -0,0 +1,283 @@
+# src/kontra/probes/relationship.py
+"""
+Relationship probe: Measure JOIN viability between two datasets.
+
+This probe answers: "What is the shape of this join?"
+
+It does NOT answer: which join type to use, or whether the join is correct.
+"""
+
+from __future__ import annotations
+
+from typing import Any, List, Union
+
+import polars as pl
+
+from kontra.api.compare import RelationshipProfile
+
+
+def profile_relationship(
+    left: Union[pl.DataFrame, str],
+    right: Union[pl.DataFrame, str],
+    on: Union[str, List[str]],
+    *,
+    sample_limit: int = 5,
+    save: bool = False,
+) -> RelationshipProfile:
+    """
+    Profile the structural relationship between two datasets.
+
+    Answers: "What is the shape of this join?"
+
+    Does NOT answer: which join type to use, or whether the join is correct.
+
+    This probe provides deterministic, structured measurements that allow
+    agents (and humans) to understand JOIN viability before writing SQL.
+
+    Args:
+        left: Left dataset (DataFrame or path/URI)
+        right: Right dataset (DataFrame or path/URI)
+        on: Column(s) to join on
+        sample_limit: Max samples per category (default 5)
+        save: Persist result to state backend (not yet implemented)
+
+    Returns:
+        RelationshipProfile with key_stats, cardinality, coverage,
+        and bounded samples.
+
+    Example:
+        # Profile before writing JOIN
+        profile = profile_relationship(orders, customers, on="customer_id")
+
+        # Check for issues
+        if profile.right_key_multiplicity_max > 1:
+            print("Warning: right side has duplicates, JOIN may explode rows")
+            print(f"Sample duplicated keys: {profile.samples_right_duplicates}")
+
+        if profile.left_keys_without_match > 0:
+            print(f"Warning: {profile.left_keys_without_match} keys won't match")
+
+        # Get structured output for LLM
+        print(profile.to_llm())
+    """
+    # Normalize on to list
+    if isinstance(on, str):
+        on = [on]
+
+    # Load data if paths provided
+    left_df = _load_data(left)
+    right_df = _load_data(right)
+
+    # Compute the profile
+    result = _compute_relationship(left_df, right_df, on, sample_limit)
+
+    # TODO: Implement save functionality
+    if save:
+        pass
+
+    return result
+
+
+def _load_data(data: Union[pl.DataFrame, str]) -> pl.DataFrame:
+    """
+    Load data from DataFrame or path/URI.
+
+    For MVP, only Polars DataFrames are fully supported.
+    File paths are loaded via Polars read functions.
+    """
+    if isinstance(data, pl.DataFrame):
+        return data
+
+    if isinstance(data, str):
+        # Simple file loading for MVP
+        if data.lower().endswith(".parquet"):
+            return pl.read_parquet(data)
+        elif data.lower().endswith(".csv"):
+            return pl.read_csv(data)
+        elif data.startswith("s3://"):
+            return pl.read_parquet(data)
+        else:
+            # Try parquet first, then CSV
+            try:
+                return pl.read_parquet(data)
+            except Exception:
+                return pl.read_csv(data)
+
+    raise ValueError(f"Unsupported data type: {type(data)}")
+
+
+def _compute_relationship(
+    left: pl.DataFrame,
+    right: pl.DataFrame,
+    on: List[str],
+    sample_limit: int,
+) -> RelationshipProfile:
+    """
+    Compute all relationship metrics between left and right datasets.
+
+    This is the core algorithm implementing the MVP schema.
+    """
+    # Validate key columns exist
+    for col in on:
+        if col not in left.columns:
+            raise ValueError(f"Key column '{col}' not found in left dataset")
+        if col not in right.columns:
+            raise ValueError(f"Key column '{col}' not found in right dataset")
+
+    # ==========================================================================
+    # 1. Basic counts
+    # ==========================================================================
+    left_rows = len(left)
+    right_rows = len(right)
+
+    # ==========================================================================
+    # 2. Key stats - left
+    # ==========================================================================
+    # Null rate: fraction of rows with any NULL in join key columns
+    if left_rows > 0:
+        null_mask = pl.lit(False)
+        for col in on:
+            null_mask = null_mask | pl.col(col).is_null()
+        left_null_count = len(left.filter(null_mask))
+        left_null_rate = left_null_count / left_rows
+    else:
+        left_null_rate = 0.0
+
+    # Unique keys (excluding NULLs)
+    left_keys = left.select(on).drop_nulls().unique()
+    left_unique_keys = len(left_keys)
+
+    # Duplicate keys: count of keys appearing >1x
+    left_key_counts = left.drop_nulls(subset=on).group_by(on).agg(pl.len().alias("_count"))
+    left_duplicate_keys = len(left_key_counts.filter(pl.col("_count") > 1))
+
+    # ==========================================================================
+    # 3. Key stats - right
+    # ==========================================================================
+    if right_rows > 0:
+        null_mask = pl.lit(False)
+        for col in on:
+            null_mask = null_mask | pl.col(col).is_null()
+        right_null_count = len(right.filter(null_mask))
+        right_null_rate = right_null_count / right_rows
+    else:
+        right_null_rate = 0.0
+
+    # Unique keys (excluding NULLs)
+    right_keys = right.select(on).drop_nulls().unique()
+    right_unique_keys = len(right_keys)
+
+    # Duplicate keys: count of keys appearing >1x
+    right_key_counts = right.drop_nulls(subset=on).group_by(on).agg(pl.len().alias("_count"))
+    right_duplicate_keys = len(right_key_counts.filter(pl.col("_count") > 1))
+
+    # ==========================================================================
+    # 4. Cardinality (rows per key)
+    # ==========================================================================
+    # Left multiplicity
+    if len(left_key_counts) > 0:
+        left_key_multiplicity_min = left_key_counts["_count"].min()
+        left_key_multiplicity_max = left_key_counts["_count"].max()
+    else:
+        left_key_multiplicity_min = 0
+        left_key_multiplicity_max = 0
+
+    # Right multiplicity
+    if len(right_key_counts) > 0:
+        right_key_multiplicity_min = right_key_counts["_count"].min()
+        right_key_multiplicity_max = right_key_counts["_count"].max()
+    else:
+        right_key_multiplicity_min = 0
+        right_key_multiplicity_max = 0
+
+    # ==========================================================================
+    # 5. Coverage
+    # ==========================================================================
+    # Left keys with match in right
+    left_matched = left_keys.join(right_keys, on=on, how="inner")
+    left_keys_with_match = len(left_matched)
+    left_keys_without_match = left_unique_keys - left_keys_with_match
+
+    # Right keys with match in left
+    right_matched = right_keys.join(left_keys, on=on, how="inner")
+    right_keys_with_match = len(right_matched)
+    right_keys_without_match = right_unique_keys - right_keys_with_match
+
+    # ==========================================================================
+    # 6. Samples
+    # ==========================================================================
+    # Sample left unmatched keys
+    left_unmatched = left_keys.join(right_keys, on=on, how="anti")
+    samples_left_unmatched = _extract_key_samples(left_unmatched, on, sample_limit)
+
+    # Sample right unmatched keys
+    right_unmatched = right_keys.join(left_keys, on=on, how="anti")
+    samples_right_unmatched = _extract_key_samples(right_unmatched, on, sample_limit)
+
+    # Sample right duplicate keys
+    right_duplicates = right_key_counts.filter(pl.col("_count") > 1).select(on)
+    samples_right_duplicates = _extract_key_samples(right_duplicates, on, sample_limit)
+
+    # ==========================================================================
+    # Build result
+    # ==========================================================================
+    return RelationshipProfile(
+        # Meta
+        on=on,
+        left_rows=left_rows,
+        right_rows=right_rows,
+        execution_tier="polars",
+
+        # Key stats - left
+        left_null_rate=left_null_rate,
+        left_unique_keys=left_unique_keys,
+        left_duplicate_keys=left_duplicate_keys,
+
+        # Key stats - right
+        right_null_rate=right_null_rate,
+        right_unique_keys=right_unique_keys,
+        right_duplicate_keys=right_duplicate_keys,
+
+        # Cardinality
+        left_key_multiplicity_min=left_key_multiplicity_min,
+        left_key_multiplicity_max=left_key_multiplicity_max,
+        right_key_multiplicity_min=right_key_multiplicity_min,
+        right_key_multiplicity_max=right_key_multiplicity_max,
+
+        # Coverage
+        left_keys_with_match=left_keys_with_match,
+        left_keys_without_match=left_keys_without_match,
+        right_keys_with_match=right_keys_with_match,
+        right_keys_without_match=right_keys_without_match,
+
+        # Samples
+        samples_left_unmatched=samples_left_unmatched,
+        samples_right_unmatched=samples_right_unmatched,
+        samples_right_duplicates=samples_right_duplicates,
+
+        # Config
+        sample_limit=sample_limit,
+    )
+
+
+def _extract_key_samples(
+    keys_df: pl.DataFrame,
+    on: List[str],
+    limit: int,
+) -> List[Any]:
+    """
+    Extract sample key values from a DataFrame.
+
+    Returns list of key values (single value if single key, dict if composite).
+    """
+    if len(keys_df) == 0:
+        return []
+
+    samples = keys_df.head(limit)
+
+    if len(on) == 1:
+        # Single key - return list of values
+        return samples[on[0]].to_list()
+    else:
+        # Composite key - return list of dicts
+        return samples.to_dicts()