kontra 0.5.2__py3-none-any.whl

kontra/preplan/planner.py

@@ -0,0 +1,253 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import date, datetime
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import pyarrow.fs as pafs  # <-- Added
+import pyarrow.parquet as pq
+
+from .types import PrePlan, Decision
+
+# NOTE: The preplan consumes simple, metadata-usable predicates only.
+# Shape: (rule_id, column, op, value)
+#   op ∈ {"==","!=",">=",">","<=","<","^=","not_null"}
+#   "^=" means "string prefix"
+Predicate = Tuple[str, str, str, Any]  # (rule_id, column, op, value)
+
+
+# ---------- small helpers ----------
+
+def _iso(v: Any) -> Any:
+    if isinstance(v, (date, datetime)):
+        return v.isoformat()
+    return v
+
+
+def _schema_names(md_schema) -> List[str]:
+    # Compatible with various pyarrow versions
+    try:
+        return list(md_schema.names)
+    except Exception:
+        try:
+            return [f.name for f in md_schema.to_arrow_schema()]
+        except Exception:
+            return []
+
+
+def _rg_col_stats(rg, j) -> Optional[Dict[str, Any]]:
+    """Return a safe dict of min/max/null_count for a row-group column j."""
+    col = rg.column(j)
+    stats = col.statistics
+    if stats is None:
+        return None
+    out: Dict[str, Any] = {
+        "min": _iso(getattr(stats, "min", None)) if getattr(stats, "has_min_max", True) else None,
+        "max": _iso(getattr(stats, "max", None)) if getattr(stats, "has_min_max", True) else None,
+    }
+    if getattr(stats, "has_null_count", True):
+        out["null_count"] = getattr(stats, "null_count", None)
+    return out
+
+
+def _name_for_rg_col(rg, j, fallback: str) -> str:
+    try:
+        # path_in_schema handles nested names properly
+        return str(rg.column(j).path_in_schema)
+    except Exception:
+        return fallback
+
+
+# ---------- metadata reasoning (per predicate, per row group) ----------
+
+def _verdict_overlaps(op: str, val: Any, stats: Optional[Dict[str, Any]]) -> Optional[bool]:
+    """
+    Return:
+      - True  -> group MAY satisfy the predicate (cannot be ruled out by min/max)
+      - False -> group CANNOT satisfy predicate (disjoint by min/max)
+      - None  -> unknown (no stats)
+    """
+    if not stats or (stats.get("min") is None and stats.get("max") is None):
+        return None
+    mn, mx = stats.get("min"), stats.get("max")
+
+    # Normalize type for string columns
+    if isinstance(mn, str) and not isinstance(val, str):
+        val = str(val)
+
+    if op == "==":
+        if mn is not None and mx is not None and (val < mn or val > mx):
+            return False
+        return True
+    if op == "!=":
+        return True  # min/max alone cannot rule "!=" out
+    if op == ">=":
+        return False if (mx is not None and mx < val) else True
+    if op == "<=":
+        return False if (mn is not None and mn > val) else True
+    if op == ">":
+        return False if (mx is not None and mx <= val) else True
+    if op == "<":
+        return False if (mn is not None and mn >= val) else True
+    if op == "^=":  # string prefix: keep if ranges overlap the prefix window
+        if not isinstance(mn, str) or not isinstance(mx, str):
+            return None
+        upper = str(val) + "\uffff"
+        return not (upper < mn or str(val) > mx)
+    if op == "not_null":
+        # Overlap sense isn't meaningful; we handle not_null via _decide_fail/_decide_pass
+        return None
+    return None
+
+
+def _decide_pass(op: str, val: Any, rg_stats_iter: Iterable[Optional[Dict[str, Any]]]) -> bool:
+    """
+    Can we *prove* that EVERY row in the file satisfies the predicate using only RG stats?
+    (dataset-level "PASS" for that rule)
+    """
+    # For >= c: if for all rgs min >= c → pass
+    # For <= c: if for all rgs max <= c → pass
+    # For == c: if for all rgs (min==max==c) → pass
+    # For not_null: if for all rgs null_count == 0 → pass
+    ok_all = True
+    for s in rg_stats_iter:
+        if s is None:
+            return False
+        mn, mx = s.get("min"), s.get("max")
+        if op == ">=":
+            if mn is None or mn < val:
+                ok_all = False; break
+        elif op == "<=":
+            if mx is None or mx > val:
+                ok_all = False; break
+        elif op == "==":
+            if mn is None or mx is None or not (mn == val and mx == val):
+                ok_all = False; break
+        elif op == "not_null":
+            # Can only prove PASS if null_count is exactly 0 for all row groups
+            # null_count > 0 means violations exist; None means unknown (can't prove)
+            if s.get("null_count") != 0:
+                ok_all = False; break
+        else:
+            # For >, <, !=, ^= we don't try to prove dataset-level PASS via min/max
+            ok_all = False; break
+    return ok_all
+
+
+def _decide_fail(op: str, val: Any, rg_stats_iter: Iterable[Optional[Dict[str, Any]]]) -> bool:
+    """
+    Can we *prove* that AT LEAST ONE row violates the predicate using RG stats?
+    (dataset-level "FAIL" for that rule)
+    """
+    for s in rg_stats_iter:
+        if s is None:
+            continue
+        mn, mx = s.get("min"), s.get("max")
+        if op == ">=":
+            # If an RG has mx < val ⇒ all rows in that RG violate ⇒ dataset FAIL
+            if mx is not None and mx < val:
+                return True
+        elif op == "<=":
+            if mn is not None and mn > val:
+                return True
+        elif op == "==":
+            # If an RG has range entirely not equal to val ⇒ all rows in that RG violate
+            if mn is not None and mx is not None and (mx < val or mn > val or (mn == mx and mn != val)):
+                return True
+        elif op == "not_null":
+            # Any rg with null_count > 0 proves at least one violation
+            nulls = s.get("null_count")
+            if isinstance(nulls, int) and nulls > 0:
+                return True
+        # For >, <, !=, ^= we typically cannot prove dataset-level FAIL with min/max alone.
+    return False
+
+
+# ---------- public API ----------
+
+def preplan_single_parquet(
+    path: str,
+    required_columns: List[str],
+    predicates: List[Predicate],
+    filesystem: pafs.FileSystem | None = None,  # <-- Updated
+) -> PrePlan:
+    """
+    Metadata-only pre-planner for a SINGLE Parquet file.
+
+    Inputs:
+      - path:             Parquet file path/URI
+      - required_columns: union of columns needed for *all* rules (from your CompiledPlan)
+      - predicates:       metadata-usable predicates -> List[(rule_id, column, op, value)]
+      - filesystem:       PyArrow filesystem object (e.g., for S3)
+
+    Outputs (PrePlan):
+      - manifest_row_groups: RG indices that STILL MATTER for remaining rules
+      - manifest_columns:    columns still needed (you can pass through required_columns)
+      - rule_decisions:      rule_id -> "pass_meta" | "fail_meta" | "unknown"
+      - stats:               {"rg_total": N, "rg_kept": K}
+    """
+    pf = pq.ParquetFile(path, filesystem=filesystem)  # <-- Updated
+    md = pf.metadata
+    schema_names = _schema_names(md.schema)
+
+    # Pre-extract per-RG per-column stats into a simple map:
+    # rg_stats[i][col_name] -> {"min":..., "max":..., "null_count":...}
+    rg_stats: List[Dict[str, Dict[str, Any]]] = []
+    for i in range(md.num_row_groups):
+        rg = md.row_group(i)
+        per_col: Dict[str, Dict[str, Any]] = {}
+        for j in range(rg.num_columns):
+            name = _name_for_rg_col(rg, j, schema_names[j] if j < len(schema_names) else f"col_{j}")
+            s = _rg_col_stats(rg, j)
+            if s is not None:
+                per_col[name] = s
+        rg_stats.append(per_col)
+
+    # Decide each rule at dataset-level (PASS/FAIL/UNKNOWN by metadata)
+    rule_decisions: Dict[str, Decision] = {}
+    for rule_id, col, op, val in predicates:
+        stats_iter = (rgc.get(col) for rgc in rg_stats)
+        if _decide_fail(op, val, stats_iter):
+            rule_decisions[rule_id] = "fail_meta"
+            continue
+        # need a fresh iterator
+        stats_iter = (rgc.get(col) for rgc in rg_stats)
+        if _decide_pass(op, val, stats_iter):
+            rule_decisions[rule_id] = "pass_meta"
+        else:
+            rule_decisions[rule_id] = "unknown"
+
+    # Determine which RGs we still need to scan (conservative):
+    # - If no predicates at all -> keep ALL RGs.
+    # - Else keep any RG that *might* be relevant for at least one UNKNOWN rule.
+    keep_rg: List[int] = list(range(md.num_row_groups))
+    unknown_preds = [(rid, col, op, val) for (rid, col, op, val) in predicates if rule_decisions.get(rid) == "unknown"]
+
+    if unknown_preds:
+        keep_rg = []
+        for i, per_col in enumerate(rg_stats):
+            # Keep if ANY unknown predicate "may overlap"
+            keep = False
+            for _, col, op, val in unknown_preds:
+                verdict = _verdict_overlaps(op, val, per_col.get(col))
+                # Verdict True  -> overlaps; Verdict None -> unknown -> keep to be safe
+                if verdict is True or verdict is None:
+                    keep = True
+                    break
+            if keep:
+                keep_rg.append(i)
+        if not keep_rg:
+            # Safety: if overlap logic ended up too strict, default to ALL
+            keep_rg = list(range(md.num_row_groups))
+
+    preplan = PrePlan(
+        manifest_columns=list(required_columns) if required_columns else [],
+        manifest_row_groups=keep_rg,
+        rule_decisions=rule_decisions,
+        stats={
+            "rg_total": md.num_row_groups,
+            "rg_kept": len(keep_rg),
+            "total_rows": md.num_rows,
+        },
+    )
+    return preplan

kontra/preplan/postgres.py

@@ -0,0 +1,179 @@
+# src/kontra/preplan/postgres.py
+"""
+PostgreSQL preplan - use pg_stats for metadata-only rule decisions.
+
+Similar to Parquet metadata preplan, but uses PostgreSQL's statistics catalog.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+
+from kontra.connectors.handle import DatasetHandle
+from kontra.connectors.postgres import PostgresConnectionParams, get_connection
+
+from .types import PrePlan, Decision
+
+
+# Predicate format: (rule_id, column, op, value)
+Predicate = Tuple[str, str, str, Any]
+
+
+def fetch_pg_stats_for_preplan(
+    params: PostgresConnectionParams,
+) -> Dict[str, Dict[str, Any]]:
+    """
+    Fetch PostgreSQL statistics for preplan decisions.
+
+    Returns:
+        {
+            "__table__": {"row_estimate": int, "page_count": int},
+            "column_name": {
+                "null_frac": float,      # 0.0-1.0 fraction of nulls
+                "n_distinct": float,     # -1 = unique, >0 = count, <0 = fraction
+                "most_common_vals": str, # Array literal
+            },
+            ...
+        }
+    """
+    with get_connection(params) as conn:
+        with conn.cursor() as cur:
+            # Table-level stats from pg_class
+            cur.execute(
+                """
+                SELECT reltuples::bigint AS row_estimate,
+                       relpages AS page_count
+                FROM pg_class
+                WHERE relname = %s
+                  AND relnamespace = %s::regnamespace
+                """,
+                (params.table, params.schema),
+            )
+            row = cur.fetchone()
+            table_stats = {
+                "row_estimate": row[0] if row else 0,
+                "page_count": row[1] if row else 0,
+            }
+
+            # Column-level stats from pg_stats
+            cur.execute(
+                """
+                SELECT attname AS column_name,
+                       null_frac,
+                       n_distinct,
+                       most_common_vals::text
+                FROM pg_stats
+                WHERE schemaname = %s AND tablename = %s
+                """,
+                (params.schema, params.table),
+            )
+
+            result: Dict[str, Dict[str, Any]] = {"__table__": table_stats}
+            for col_row in cur.fetchall():
+                col_name, null_frac, n_distinct, mcv = col_row
+                result[col_name] = {
+                    "null_frac": null_frac,
+                    "n_distinct": n_distinct,
+                    "most_common_vals": mcv,
+                }
+
+            return result
+
+
+def preplan_postgres(
+    handle: DatasetHandle,
+    required_columns: List[str],
+    predicates: List[Predicate],
+) -> PrePlan:
+    """
+    Metadata-only pre-planner for PostgreSQL tables using pg_stats.
+
+    Supports decisions for:
+      - not_null: if null_frac == 0 -> pass_meta
+      - unique: if n_distinct == -1 -> pass_meta (PostgreSQL's way of saying "all unique")
+
+    Args:
+        handle: DatasetHandle with db_params
+        required_columns: Columns needed for validation
+        predicates: List of (rule_id, column, op, value) tuples
+
+    Returns:
+        PrePlan with rule decisions based on pg_stats
+    """
+    if not handle.db_params:
+        raise ValueError("PostgreSQL handle missing db_params")
+
+    params: PostgresConnectionParams = handle.db_params
+    pg_stats = fetch_pg_stats_for_preplan(params)
+
+    table_stats = pg_stats.get("__table__", {})
+    row_estimate = table_stats.get("row_estimate", 0)
+
+    rule_decisions: Dict[str, Decision] = {}
+
+    for rule_id, column, op, value in predicates:
+        col_stats = pg_stats.get(column)
+
+        if col_stats is None:
+            # No stats for this column (ANALYZE not run or column doesn't exist)
+            rule_decisions[rule_id] = "unknown"
+            continue
+
+        null_frac = col_stats.get("null_frac")
+        n_distinct = col_stats.get("n_distinct")
+
+        if op == "not_null":
+            # If null_frac is exactly 0, the column has no nulls
+            if null_frac is not None and null_frac == 0:
+                rule_decisions[rule_id] = "pass_meta"
+            elif null_frac is not None and null_frac > 0:
+                # We know there ARE nulls, but pg_stats doesn't give exact count
+                # So we can't prove pass, but we know the rule will likely fail
+                # Be conservative: mark as unknown (actual execution will determine)
+                rule_decisions[rule_id] = "unknown"
+            else:
+                rule_decisions[rule_id] = "unknown"
+
+        elif op == "unique":
+            # n_distinct == -1 means PostgreSQL detected all values are unique
+            # n_distinct < 0 (other than -1) means n_distinct is a fraction of rows
+            if n_distinct is not None:
+                if n_distinct == -1:
+                    # All values are unique AND no nulls (unique constraint behavior)
+                    if null_frac == 0:
+                        rule_decisions[rule_id] = "pass_meta"
+                    else:
+                        # Unique values but has nulls - need to check if nulls cause dups
+                        rule_decisions[rule_id] = "unknown"
+                elif n_distinct < 0:
+                    # Fraction - close to -1 means high uniqueness but not guaranteed
+                    rule_decisions[rule_id] = "unknown"
+                else:
+                    # n_distinct > 0: exact count or estimate
+                    # If n_distinct equals row_estimate, likely unique
+                    if row_estimate > 0 and n_distinct >= row_estimate:
+                        rule_decisions[rule_id] = "pass_meta"
+                    else:
+                        rule_decisions[rule_id] = "unknown"
+            else:
+                rule_decisions[rule_id] = "unknown"
+
+        else:
+            # Other ops (>=, <=, ==, etc.) - pg_stats doesn't have min/max for general use
+            # (histogram_bounds could be used but it's complex)
+            rule_decisions[rule_id] = "unknown"
+
+    return PrePlan(
+        manifest_columns=list(required_columns) if required_columns else [],
+        manifest_row_groups=[],  # Not applicable for PostgreSQL
+        rule_decisions=rule_decisions,
+        stats={
+            "row_estimate": row_estimate,
+            "columns_with_stats": len([k for k in pg_stats if k != "__table__"]),
+        },
+    )
+
+
+def can_preplan_postgres(handle: DatasetHandle) -> bool:
+    """Check if PostgreSQL preplan is applicable for this handle."""
+    return handle.scheme in ("postgres", "postgresql") and handle.db_params is not None

kontra/preplan/sqlserver.py

@@ -0,0 +1,191 @@
+# src/kontra/preplan/sqlserver.py
+"""
+SQL Server preplan - use metadata for rule decisions.
+
+Uses SQL Server system views:
+- sys.dm_db_partition_stats for row count estimates
+- sys.columns for nullability constraints
+- sys.indexes + sys.index_columns for uniqueness constraints
+
+Note: Unlike PostgreSQL's pg_stats, SQL Server doesn't expose null_frac directly.
+We use constraint metadata instead (NOT NULL, UNIQUE indexes).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Tuple
+
+from kontra.connectors.handle import DatasetHandle
+from kontra.connectors.sqlserver import SqlServerConnectionParams, get_connection
+
+from .types import PrePlan, Decision
+
+
+# Predicate format: (rule_id, column, op, value)
+Predicate = Tuple[str, str, str, Any]
+
+
+def fetch_sqlserver_metadata(
+    params: SqlServerConnectionParams,
+) -> Dict[str, Dict[str, Any]]:
+    """
+    Fetch SQL Server metadata for preplan decisions.
+
+    Returns:
+        {
+            "__table__": {"row_estimate": int, "page_count": int},
+            "column_name": {
+                "is_nullable": bool,     # From column definition
+                "is_identity": bool,     # Identity column
+                "has_unique_constraint": bool,  # Unique index or constraint
+            },
+            ...
+        }
+    """
+    with get_connection(params) as conn:
+        cursor = conn.cursor()
+
+        # Table-level stats from sys.dm_db_partition_stats
+        cursor.execute(
+            """
+            SELECT SUM(row_count) AS row_estimate,
+                   SUM(used_page_count) AS page_count
+            FROM sys.dm_db_partition_stats ps
+            JOIN sys.objects o ON ps.object_id = o.object_id
+            JOIN sys.schemas s ON o.schema_id = s.schema_id
+            WHERE s.name = %s AND o.name = %s AND ps.index_id IN (0, 1)
+            """,
+            (params.schema, params.table),
+        )
+        row = cursor.fetchone()
+        table_stats = {
+            "row_estimate": int(row[0]) if row and row[0] else 0,
+            "page_count": int(row[1]) if row and row[1] else 0,
+        }
+
+        # Column-level metadata
+        cursor.execute(
+            """
+            SELECT
+                c.name AS column_name,
+                c.is_nullable,
+                c.is_identity
+            FROM sys.columns c
+            JOIN sys.objects o ON c.object_id = o.object_id
+            JOIN sys.schemas s ON o.schema_id = s.schema_id
+            WHERE s.name = %s AND o.name = %s
+            """,
+            (params.schema, params.table),
+        )
+
+        result: Dict[str, Dict[str, Any]] = {"__table__": table_stats}
+        for col_row in cursor.fetchall():
+            col_name, is_nullable, is_identity = col_row
+            result[col_name] = {
+                "is_nullable": bool(is_nullable),
+                "is_identity": bool(is_identity),
+                "has_unique_constraint": False,  # Will be updated below
+            }
+
+        # Check for unique constraints/indexes
+        cursor.execute(
+            """
+            SELECT c.name AS column_name
+            FROM sys.indexes i
+            JOIN sys.index_columns ic ON i.object_id = ic.object_id AND i.index_id = ic.index_id
+            JOIN sys.columns c ON ic.object_id = c.object_id AND ic.column_id = c.column_id
+            JOIN sys.objects o ON i.object_id = o.object_id
+            JOIN sys.schemas s ON o.schema_id = s.schema_id
+            WHERE s.name = %s AND o.name = %s
+              AND (i.is_unique = 1 OR i.is_primary_key = 1)
+              AND ic.is_included_column = 0
+            GROUP BY c.name
+            HAVING COUNT(*) = 1  -- Single-column unique constraint
+            """,
+            (params.schema, params.table),
+        )
+
+        for row in cursor.fetchall():
+            col_name = row[0]
+            if col_name in result:
+                result[col_name]["has_unique_constraint"] = True
+
+        return result
+
+
+def preplan_sqlserver(
+    handle: DatasetHandle,
+    required_columns: List[str],
+    predicates: List[Predicate],
+) -> PrePlan:
+    """
+    Metadata-only pre-planner for SQL Server tables.
+
+    Supports decisions for:
+      - not_null: if column is NOT NULL (is_nullable = 0) -> pass_meta
+      - unique: if column has unique index/constraint -> pass_meta
+
+    Args:
+        handle: DatasetHandle with db_params
+        required_columns: Columns needed for validation
+        predicates: List of (rule_id, column, op, value) tuples
+
+    Returns:
+        PrePlan with rule decisions based on SQL Server metadata
+    """
+    if not handle.db_params:
+        raise ValueError("SQL Server handle missing db_params")
+
+    params: SqlServerConnectionParams = handle.db_params
+    metadata = fetch_sqlserver_metadata(params)
+
+    table_stats = metadata.get("__table__", {})
+    row_estimate = table_stats.get("row_estimate", 0)
+
+    rule_decisions: Dict[str, Decision] = {}
+
+    for rule_id, column, op, value in predicates:
+        col_meta = metadata.get(column)
+
+        if col_meta is None:
+            # Column not found in metadata
+            rule_decisions[rule_id] = "unknown"
+            continue
+
+        is_nullable = col_meta.get("is_nullable", True)
+        is_identity = col_meta.get("is_identity", False)
+        has_unique = col_meta.get("has_unique_constraint", False)
+
+        if op == "not_null":
+            # If column is defined as NOT NULL, it definitely has no nulls
+            if not is_nullable:
+                rule_decisions[rule_id] = "pass_meta"
+            else:
+                # Column allows nulls - may or may not have any
+                rule_decisions[rule_id] = "unknown"
+
+        elif op == "unique":
+            # If column has unique constraint or is identity, it's unique
+            if has_unique or is_identity:
+                rule_decisions[rule_id] = "pass_meta"
+            else:
+                rule_decisions[rule_id] = "unknown"
+
+        else:
+            # Other ops - would need actual data statistics
+            rule_decisions[rule_id] = "unknown"
+
+    return PrePlan(
+        manifest_columns=list(required_columns) if required_columns else [],
+        manifest_row_groups=[],  # Not applicable for SQL Server
+        rule_decisions=rule_decisions,
+        stats={
+            "row_estimate": row_estimate,
+            "columns_with_metadata": len([k for k in metadata if k != "__table__"]),
+        },
+    )
+
+
+def can_preplan_sqlserver(handle: DatasetHandle) -> bool:
+    """Check if SQL Server preplan is applicable for this handle."""
+    return handle.scheme in ("mssql", "sqlserver") and handle.db_params is not None

kontra/preplan/types.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Literal
+
+# A rule can be proven by metadata to pass, proven to fail, or remain unknown.
+Decision = Literal["pass_meta", "fail_meta", "unknown"]
+
+
+@dataclass
+class PrePlan:
+    """
+    Result of the metadata-only pre-planning stage.
+
+    - manifest_columns: union of columns still needed for SQL/Polars after metadata decisions.
+    - manifest_row_groups: Parquet row-group indices that *may* affect remaining rules.
+                           (Single-file MVP; can evolve to a file list later.)
+    - rule_decisions: rule_id -> Decision ("pass_meta" | "fail_meta" | "unknown").
+    - stats: small numbers for observability (e.g., {"rg_total": 19, "rg_kept": 7}).
+    """
+    manifest_columns: List[str]
+    manifest_row_groups: List[int]
+    rule_decisions: Dict[str, Decision]
+    stats: Dict[str, int]

kontra/probes/__init__.py

@@ -0,0 +1,20 @@
+# src/kontra/probes/__init__.py
+"""
+Transformation probes for Kontra.
+
+Probes measure the structural effects of data transformations without
+assigning meaning or judgment. They provide deterministic, structured,
+token-efficient measurements for agents to reason about.
+
+Available probes:
+- compare: Measure differences between before/after transformation
+- profile_relationship: Measure JOIN viability between datasets
+"""
+
+from kontra.probes.compare import compare
+from kontra.probes.relationship import profile_relationship
+
+__all__ = [
+    "compare",
+    "profile_relationship",
+]