pointblank 0.17.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

pointblank/column.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 import re
 from dataclasses import dataclass
+from typing import Any
 
 import narwhals as nw
 from narwhals.typing import IntoDataFrame
 
 __all__ = [
     "col",
+    "ref",
     "starts_with",
     "ends_with",
     "contains",
@@ -192,6 +194,22 @@ class ColumnLiteral(Column):
         return self.exprs
 
 
+@dataclass
+class ReferenceColumn:
+    """
+    A class to represent a column from the reference data.
+
+    This is used with aggregate validation methods (like `col_sum_eq`, `col_avg_gt`, etc.)
+    to compare the aggregate value of a column in the main data against the aggregate
+    value of a column in the reference data.
+    """
+
+    column_name: str
+
+    def __repr__(self):
+        return f"ref({self.column_name!r})"
+
+
 @dataclass
 class ColumnSelectorNarwhals(Column):
     """
@@ -211,11 +229,16 @@ class ColumnSelectorNarwhals(Column):
 
     exprs: nw.selectors.Selector
 
-    def resolve(self, table) -> list[str]:
+    def resolve(
+        self, columns: list[str] | None = None, table: IntoDataFrame | None = None
+    ) -> list[str]:
+        # Note: columns parameter is unused - Narwhals selectors need the actual table
+        if table is None:
+            raise ValueError("ColumnSelectorNarwhals requires a table for resolution")
         # Convert the native table to a Narwhals DataFrame
         dfn = nw.from_native(table)
         # Use the selector to select columns and return their names
-        selected_df = dfn.select(self.exprs.exprs)
+        selected_df = dfn.select(self.exprs.exprs)  # type: ignore[attr-defined]
         # Use `collect_schema()` for LazyFrame to avoid performance warnings
         if hasattr(selected_df, "collect_schema"):
             return list(selected_df.collect_schema().keys())
@@ -224,7 +247,7 @@ class ColumnSelectorNarwhals(Column):
 
 
 def col(
-    exprs: str | ColumnSelector | ColumnSelectorNarwhals,
+    exprs: str | ColumnSelector | ColumnSelectorNarwhals | nw.selectors.Selector,
 ) -> Column | ColumnLiteral | ColumnSelectorNarwhals:
     """
     Helper function for referencing a column in the input table.
@@ -522,6 +545,83 @@ def col(
     raise TypeError(f"Unsupported type: {type(exprs)}")  # pragma: no cover
 
 
+def ref(column_name: str) -> ReferenceColumn:
+    """
+    Reference a column from the reference data for aggregate comparisons.
+
+    This function is used with aggregate validation methods (like `col_sum_eq`, `col_avg_gt`, etc.)
+    to compare the aggregate value of a column in the main data against the aggregate value of
+    a column in the reference data.
+
+    To use this function, you must first set the reference data on the `Validate` object using
+    the `reference=` parameter in the constructor.
+
+    Parameters
+    ----------
+    column_name
+        The name of the column in the reference data to compute the aggregate from.
+
+    Returns
+    -------
+    ReferenceColumn
+        A reference column marker that indicates the value should be computed from the
+        reference data.
+
+    Examples
+    --------
+    ```{python}
+    #| echo: false
+    #| output: false
+    import pointblank as pb
+    pb.config(report_incl_header=False, report_incl_footer=False, preview_incl_header=False)
+    ```
+
+    Suppose we have two DataFrames: a current data table and a reference (historical) table.
+    We want to validate that the sum of a column in the current data matches the sum of the
+    same column in the reference data.
+
+    ```{python}
+    import pointblank as pb
+    import polars as pl
+
+    # Current data
+    current_data = pl.DataFrame({"sales": [100, 200, 300]})
+
+    # Reference (historical) data
+    reference_data = pl.DataFrame({"sales": [100, 200, 300]})
+
+    validation = (
+        pb.Validate(data=current_data, reference=reference_data)
+        .col_sum_eq("sales", pb.ref("sales"))
+        .interrogate()
+    )
+
+    validation
+    ```
+
+    You can also compare different columns or use tolerance:
+
+    ```{python}
+    current_data = pl.DataFrame({"revenue": [105, 205, 305]})
+    reference_data = pl.DataFrame({"sales": [100, 200, 300]})
+
+    # Check if revenue sum is within 10% of sales sum
+    validation = (
+        pb.Validate(data=current_data, reference=reference_data)
+        .col_sum_eq("revenue", pb.ref("sales"), tol=0.1)
+        .interrogate()
+    )
+
+    validation
+    ```
+
+    See Also
+    --------
+    The [`col()`](`pointblank.col`) function for referencing columns within the same table.
+    """
+    return ReferenceColumn(column_name=column_name)
+
+
 def starts_with(text: str, case_sensitive: bool = False) -> StartsWith:
     """
     Select columns that start with specified text.
@@ -1634,13 +1734,25 @@ class ColumnExpression:
     Supports operations like >, <, +, etc. for creating backend-agnostic validation expressions.
     """
 
-    def __init__(self, column_name=None, operation=None, left=None, right=None):
+    column_name: str | None
+    operation: str | None
+    left: ColumnExpression | None
+    right: ColumnExpression | str | int | float | None
+
+    def __init__(
+        self,
+        column_name: str | None = None,
+        operation: str | None = None,
+        left: ColumnExpression | None = None,
+        right: ColumnExpression | str | int | float | None = None,
+    ):
         self.column_name = column_name  # Name of the column (for leaf nodes)
         self.operation = operation  # Operation type (gt, lt, add, etc.)
         self.left = left  # Left operand (ColumnExpression or None for column reference)
         self.right = right  # Right operand (ColumnExpression, value, or None)
 
-    def to_polars_expr(self):
+    # TODO: This method would benefit from stronger typing
+    def to_polars_expr(self) -> Any:
         """Convert this expression to a Polars expression."""
         import polars as pl
 
@@ -1650,16 +1762,16 @@ class ColumnExpression:
 
         # Handle unary operations like is_null
         if self.operation == "is_null":
-            left_expr = self.left
+            left_expr: Any = self.left
             if isinstance(left_expr, ColumnExpression):
                 left_expr = left_expr.to_polars_expr()
-            return left_expr.is_null()
+            return left_expr.is_null()  # type: ignore[union-attr]
 
         if self.operation == "is_not_null":
             left_expr = self.left
             if isinstance(left_expr, ColumnExpression):
                 left_expr = left_expr.to_polars_expr()
-            return left_expr.is_not_null()
+            return left_expr.is_not_null()  # type: ignore[union-attr]
 
         # Handle nested expressions through recursive evaluation
         if self.operation is None:
@@ -1667,6 +1779,7 @@ class ColumnExpression:
             raise ValueError("Invalid expression state: No operation or column name")
 
         # Get the left operand
+        left_expr: Any
         if self.left is None and self.column_name is not None:
             # Column name as left operand
             left_expr = pl.col(self.column_name)  # pragma: no cover
@@ -1678,6 +1791,7 @@ class ColumnExpression:
             left_expr = self.left  # pragma: no cover
 
         # Get the right operand
+        right_expr: Any
         if isinstance(self.right, ColumnExpression):
             # Nested expression as right operand
             right_expr = self.right.to_polars_expr()  # pragma: no cover
@@ -1688,35 +1802,35 @@ class ColumnExpression:
             # Literal value as right operand
             right_expr = self.right  # pragma: no cover
 
-        # Apply the operation
+        # Apply the operation (type ignore needed due to dynamic expression types)
         if self.operation == "gt":
-            return left_expr > right_expr
+            return left_expr > right_expr  # type: ignore[operator]
         elif self.operation == "lt":
-            return left_expr < right_expr
+            return left_expr < right_expr  # type: ignore[operator]
         elif self.operation == "eq":
             return left_expr == right_expr
         elif self.operation == "ne":
             return left_expr != right_expr
         elif self.operation == "ge":
-            return left_expr >= right_expr
+            return left_expr >= right_expr  # type: ignore[operator]
         elif self.operation == "le":
-            return left_expr <= right_expr
+            return left_expr <= right_expr  # type: ignore[operator]
         elif self.operation == "add":
-            return left_expr + right_expr
+            return left_expr + right_expr  # type: ignore[operator]
         elif self.operation == "sub":
-            return left_expr - right_expr
+            return left_expr - right_expr  # type: ignore[operator]
         elif self.operation == "mul":
-            return left_expr * right_expr
+            return left_expr * right_expr  # type: ignore[operator]
         elif self.operation == "div":
-            return left_expr / right_expr
+            return left_expr / right_expr  # type: ignore[operator]
         elif self.operation == "and":
-            return left_expr & right_expr
+            return left_expr & right_expr  # type: ignore[operator]
         elif self.operation == "or":
-            return left_expr | right_expr
+            return left_expr | right_expr  # type: ignore[operator]
         else:
             raise ValueError(f"Unsupported operation: {self.operation}")
 
-    def to_pandas_expr(self, df):
+    def to_pandas_expr(self, df: Any) -> Any:
         """Convert this expression to a Pandas Series of booleans."""
 
         # Handle is_null as a special case - but raise an error
@@ -1739,43 +1853,43 @@ class ColumnExpression:
             return df[self.column_name]
 
         # For other operations, recursively process operands
-        left_expr = self.left
+        left_expr: Any = self.left
         if isinstance(left_expr, ColumnExpression):
             left_expr = left_expr.to_pandas_expr(df)
         elif isinstance(left_expr, str) and left_expr in df.columns:  # pragma: no cover
             left_expr = df[left_expr]
 
-        right_expr = self.right
+        right_expr: Any = self.right
         if isinstance(right_expr, ColumnExpression):
             right_expr = right_expr.to_pandas_expr(df)
         elif isinstance(right_expr, str) and right_expr in df.columns:  # pragma: no cover
             right_expr = df[right_expr]
 
-        # Apply the operation
+        # Apply the operation (type ignore needed due to dynamic expression types)
         if self.operation == "gt":
-            return left_expr > right_expr
+            return left_expr > right_expr  # type: ignore[operator]
         elif self.operation == "lt":
-            return left_expr < right_expr
+            return left_expr < right_expr  # type: ignore[operator]
         elif self.operation == "eq":
             return left_expr == right_expr
         elif self.operation == "ne":
             return left_expr != right_expr
         elif self.operation == "ge":
-            return left_expr >= right_expr
+            return left_expr >= right_expr  # type: ignore[operator]
         elif self.operation == "le":
-            return left_expr <= right_expr
+            return left_expr <= right_expr  # type: ignore[operator]
         elif self.operation == "add":
-            return left_expr + right_expr
+            return left_expr + right_expr  # type: ignore[operator]
         elif self.operation == "sub":
-            return left_expr - right_expr
+            return left_expr - right_expr  # type: ignore[operator]
         elif self.operation == "mul":
-            return left_expr * right_expr
+            return left_expr * right_expr  # type: ignore[operator]
         elif self.operation == "div":
-            return left_expr / right_expr
+            return left_expr / right_expr  # type: ignore[operator]
         else:
             raise ValueError(f"Unsupported operation: {self.operation}")
 
-    def to_ibis_expr(self, table):
+    def to_ibis_expr(self, table: Any) -> Any:
         """Convert this expression to an Ibis expression."""
 
         # Base case: simple column reference
@@ -1784,16 +1898,16 @@ class ColumnExpression:
 
         # Handle unary operations
         if self.operation == "is_null":
-            left_expr = self.left
+            left_expr: Any = self.left
             if isinstance(left_expr, ColumnExpression):
                 left_expr = left_expr.to_ibis_expr(table)
-            return left_expr.isnull()
+            return left_expr.isnull()  # type: ignore[union-attr]
 
         if self.operation == "is_not_null":
             left_expr = self.left
             if isinstance(left_expr, ColumnExpression):
                 left_expr = left_expr.to_ibis_expr(table)
-            return ~left_expr.isnull()
+            return ~left_expr.isnull()  # type: ignore[union-attr,operator]
 
         # Handle nested expressions through recursive evaluation
         if self.operation is None:
@@ -1801,6 +1915,7 @@ class ColumnExpression:
             raise ValueError("Invalid expression state: No operation or column name")
 
         # Get the left operand
+        left_expr: Any
         if self.left is None and self.column_name is not None:
             # Column name as left operand
             left_expr = table[self.column_name]  # pragma: no cover
@@ -1812,6 +1927,7 @@ class ColumnExpression:
             left_expr = self.left  # pragma: no cover
 
         # Get the right operand
+        right_expr: Any
         if isinstance(self.right, ColumnExpression):
             # Nested expression as right operand
             right_expr = self.right.to_ibis_expr(table)  # pragma: no cover
@@ -1822,31 +1938,31 @@ class ColumnExpression:
             # Literal value as right operand
             right_expr = self.right  # pragma: no cover
 
-        # Apply the operation
+        # Apply the operation (type ignore needed due to dynamic expression types)
         if self.operation == "gt":
-            return left_expr > right_expr
+            return left_expr > right_expr  # type: ignore[operator]
         elif self.operation == "lt":
-            return left_expr < right_expr
+            return left_expr < right_expr  # type: ignore[operator]
         elif self.operation == "eq":
             return left_expr == right_expr
         elif self.operation == "ne":
             return left_expr != right_expr
         elif self.operation == "ge":
-            return left_expr >= right_expr
+            return left_expr >= right_expr  # type: ignore[operator]
         elif self.operation == "le":
-            return left_expr <= right_expr
+            return left_expr <= right_expr  # type: ignore[operator]
         elif self.operation == "add":
-            return left_expr + right_expr
+            return left_expr + right_expr  # type: ignore[operator]
         elif self.operation == "sub":
-            return left_expr - right_expr
+            return left_expr - right_expr  # type: ignore[operator]
         elif self.operation == "mul":
-            return left_expr * right_expr
+            return left_expr * right_expr  # type: ignore[operator]
         elif self.operation == "div":
-            return left_expr / right_expr
+            return left_expr / right_expr  # type: ignore[operator]
         elif self.operation == "and":
-            return left_expr & right_expr
+            return left_expr & right_expr  # type: ignore[operator]
         elif self.operation == "or":
-            return left_expr | right_expr
+            return left_expr | right_expr  # type: ignore[operator]
         else:
             raise ValueError(f"Unsupported operation: {self.operation}")