From f949272073b78743cad9fe5748f291389da352ee Mon Sep 17 00:00:00 2001 From: Konstantinos Stefanidis Vozikis Date: Thu, 13 Aug 2026 14:30:55 +0200 Subject: [PATCH] fix: build delete filters with PyIceberg expressions --- INSTALL-AND-REFERENCE.md | 19 +++- src/tower/_tables.py | 96 +++++++++++------ src/tower/exceptions.py | 11 ++ src/tower/utils/pyarrow.py | 150 --------------------------- tests/tower/test_table_filters.py | 167 ++++++++++++++++++++++++++++++ tests/tower/test_tables.py | 2 +- 6 files changed, 260 insertions(+), 185 deletions(-) create mode 100644 tests/tower/test_table_filters.py diff --git a/INSTALL-AND-REFERENCE.md b/INSTALL-AND-REFERENCE.md index 3dd698de..a586dd07 100644 --- a/INSTALL-AND-REFERENCE.md +++ b/INSTALL-AND-REFERENCE.md @@ -128,8 +128,23 @@ pip install "tower[ai]" pip install "tower[iceberg]" ``` -- `tower.create_table`: create Iceberg tables -- `tower.load_table`: load data from Iceberg tables +- `tower.tables(...)`: load, create, update, and delete Iceberg table data + +Delete filters are SQL-like strings or native PyIceberg boolean expressions. The +`Table.column()` builder creates composable PyIceberg predicates: + +```python +table = tower.tables("events").load() +table.delete( + (table.column("age") >= 18) + & ~(table.column("status") == "inactive") +) +``` + +PyArrow compute expressions and lists of expressions are no longer accepted as delete +filters. Replace `pc.field("age") >= 18` with `table.column("age") >= 18`, and replace +`[a, b]` with `a & b`. Code using PyArrow for Arrow-side filtering can continue to use +`pyarrow.compute.field()` outside the Tower table API. ### dbt Core support diff --git a/src/tower/_tables.py b/src/tower/_tables.py index b477fe79..a099085e 100644 --- a/src/tower/_tables.py +++ b/src/tower/_tables.py @@ -5,8 +5,17 @@ import random import time from dataclasses import dataclass -from typing import Callable, List, Optional, TypeVar, Union - +from typing import Any, Callable, Optional, TypeVar, Union + +from pyiceberg.expressions import ( + BooleanExpression, + EqualTo, + GreaterThan, + GreaterThanOrEqual, + LessThan, + LessThanOrEqual, + NotEqualTo, +) from pyiceberg.exceptions import CommitFailedException, NoSuchTableError TTable = TypeVar("TTable", bound="Table") @@ -28,11 +37,9 @@ get_tower_catalog_credentials, load_vended_catalog, ) +from .exceptions import PyArrowFilterMigrationError from .tower_api_client.models import CatalogCredentials -from .utils.pyarrow import ( - convert_pyarrow_expressions, - convert_pyarrow_schema, -) +from .utils.pyarrow import convert_pyarrow_schema from .utils.tables import ( make_table_name, namespace_or_default, @@ -47,6 +54,29 @@ class RowsAffectedInformation: updates: int +@dataclass(frozen=True, eq=False) +class _TableColumn: + name: str + + def __eq__(self, value: Any) -> BooleanExpression: + return EqualTo(self.name, value) + + def __ne__(self, value: Any) -> BooleanExpression: + return NotEqualTo(self.name, value) + + def __gt__(self, value: Any) -> BooleanExpression: + return GreaterThan(self.name, value) + + def __ge__(self, value: Any) -> BooleanExpression: + return GreaterThanOrEqual(self.name, value) + + def __lt__(self, value: Any) -> BooleanExpression: + return LessThan(self.name, value) + + def __le__(self, value: Any) -> BooleanExpression: + return LessThanOrEqual(self.name, value) + + _VendedCatalogIdentity = tuple[str, str, str] @@ -410,7 +440,7 @@ def upsert( def delete( self, - filters: Union[str, List[pc.Expression]], + filters: str | BooleanExpression, max_retries: int = 5, retry_delay_seconds: float = 0.5, ) -> TTable: @@ -423,11 +453,8 @@ def delete( cannot be tracked due to limitations in the underlying Iceberg implementation. Args: - filters (Union[str, List[pc.Expression]]): The filter conditions to apply. - Can be either: - - A single PyArrow compute expression - - A list of PyArrow compute expressions (combined with AND) - - A string expression + filters (str | BooleanExpression): A SQL-like string or a PyIceberg + boolean expression. Use ``Table.column()`` to construct expressions. max_retries (int): Maximum number of retry attempts on commit conflicts. Defaults to 5. retry_delay_seconds (float): Maximum randomized wait before the first retry, @@ -450,21 +477,17 @@ def delete( >>> # Delete rows where age is greater than 30 >>> table.delete(table.column("age") > 30) >>> # Delete rows matching multiple conditions - >>> table.delete([ - ... table.column("age") > 30, - ... table.column("department") == "IT" - ... ]) + >>> table.delete( + ... (table.column("age") > 30) + ... & (table.column("department") == "IT") + ... ) >>> # Delete rows using a string expression >>> table.delete("age > 30 AND department = 'IT'") """ self._validate_retry_args(max_retries, retry_delay_seconds) + filters = self._normalize_delete_filter(filters) self._ensure_read_write_table() - if isinstance(filters, list): - # We need to convert the pc.Expression into PyIceberg - next_filters = convert_pyarrow_expressions(filters) - filters = next_filters - self._commit_with_retry( lambda: self._table.delete( delete_filter=filters, @@ -480,6 +503,16 @@ def delete( return self + @staticmethod + def _normalize_delete_filter(filters: object) -> str | BooleanExpression: + if isinstance(filters, (pc.Expression, list)): + raise PyArrowFilterMigrationError() + if isinstance(filters, (str, BooleanExpression)): + return filters + raise TypeError( + "filters must be a SQL-like string or a PyIceberg BooleanExpression" + ) + def schema(self) -> pa.Schema: """ Returns the schema of the table as a PyArrow schema. @@ -496,19 +529,19 @@ def schema(self) -> pa.Schema: iceberg_schema = self._table.schema() return iceberg_schema.as_arrow() - def column(self, name: str) -> pa.compute.Expression: + def column(self, name: str) -> _TableColumn: """ - Returns a column from the table as a PyArrow compute expression. + Returns a structural builder for PyIceberg filter expressions. This method is useful for creating column-based expressions that can be used in - operations like filtering, sorting, or aggregating data. The returned expression - can be used with PyArrow's compute functions. + comparison operators build PyIceberg boolean expressions that can be passed to + ``delete()`` and composed with ``&``, ``|``, and ``~``. Args: name (str): The name of the column to retrieve from the table schema. Returns: - pa.compute.Expression: A PyArrow compute expression representing the column. + _TableColumn: A builder for PyIceberg comparison expressions. Raises: ValueError: If the specified column name is not found in the table schema. @@ -520,13 +553,12 @@ def column(self, name: str) -> pa.compute.Expression: >>> # Use the expression in a delete operation >>> table.delete(age_expr) """ - field = self.schema().field(name) - - if field is None: - raise ValueError(f"Column {name} not found in table schema") + try: + self._table.schema().find_field(name, case_sensitive=True) + except ValueError: + raise ValueError(f"Column {name} not found in table schema") from None - # We need to convert the PyArrow field into pa.compute.Expression - return pa.compute.field(name) + return _TableColumn(name) class TableReference: diff --git a/src/tower/exceptions.py b/src/tower/exceptions.py index b37c2a80..0035f952 100644 --- a/src/tower/exceptions.py +++ b/src/tower/exceptions.py @@ -32,3 +32,14 @@ def __init__(self, app_name: str, number: int, state: str): class AppNotFoundError(RuntimeError): def __init__(self, app_name: str): super().__init__(f"App '{app_name}' not found in the Tower.") + + +class PyArrowFilterMigrationError(TypeError): + def __init__(self): + super().__init__( + "PyArrow compute expressions are no longer accepted as table delete " + 'filters. Replace pc.field("age") >= 18 with ' + 'table.column("age") >= 18. Combine predicates with &, |, and ~; ' + "replace [a, b] with a & b. You can also pass a PyIceberg " + "BooleanExpression or a SQL-like filter string." + ) diff --git a/src/tower/utils/pyarrow.py b/src/tower/utils/pyarrow.py index 9b8c1111..d3b46186 100644 --- a/src/tower/utils/pyarrow.py +++ b/src/tower/utils/pyarrow.py @@ -1,23 +1,7 @@ -from typing import Any, Optional, List - import pyarrow as pa -import pyarrow.compute as pc from pyiceberg import types as iceberg_types from pyiceberg.schema import Schema as IcebergSchema -from pyiceberg.expressions import ( - BooleanExpression, - And, - Or, - Not, - EqualTo, - NotEqualTo, - GreaterThan, - GreaterThanOrEqual, - LessThan, - LessThanOrEqual, - Reference, -) class FieldIdManager: @@ -192,137 +176,3 @@ def convert_pyarrow_schema( ) ) return IcebergSchema(*iceberg_fields, schema_id=schema_id) - - -def extract_field_and_literal(expr: pc.Expression) -> tuple[str, Any]: - """Extract field name and literal value from a comparison expression.""" - # First, convert the expression to a string and parse it - expr_str = str(expr) - - # PyArrow expression strings look like: "(field_name == literal)" or similar - # Need to determine the operator and then split accordingly - operators = ["==", "!=", ">", ">=", "<", "<="] - op_used = None - for op in operators: - if op in expr_str: - op_used = op - break - - if not op_used: - raise ValueError( - f"Could not find comparison operator in expression: {expr_str}" - ) - - # Remove parentheses and split by operator - expr_clean = expr_str.strip("()") - parts = expr_clean.split(op_used) - if len(parts) != 2: - raise ValueError(f"Expected binary comparison in expression: {expr_str}") - - # Determine which part is the field and which is the literal - field_name = None - literal_value = None - - # Clean up the parts - left = parts[0].strip() - right = parts[1].strip() - - # Typically field name doesn't have quotes, literals (strings) do - if left.startswith('"') or left.startswith("'"): - # Right side is the field - field_name = right - # Extract the literal value - this is a simplification - literal_value = left.strip("\"'") - else: - # Left side is the field - field_name = left - # Extract the literal value - this is a simplification - literal_value = right.strip("\"'") - - # Try to convert numeric literals - try: - if "." in literal_value: - literal_value = float(literal_value) - else: - literal_value = int(literal_value) - except ValueError: - # Keep as string if not numeric - pass - - return field_name, literal_value - - -def convert_pyarrow_expression(expr: pc.Expression) -> Optional[BooleanExpression]: - """Convert a PyArrow compute expression to a PyIceberg boolean expression.""" - if expr is None: - return None - - # Handle the expression based on its string representation - expr_str = str(expr) - - # Handle logical operations - if "and" in expr_str.lower() and isinstance(expr, pc.Expression): - # This is a simplification - in real code, you'd need to parse the expression - # to extract the sub-expressions properly - left_expr = None # You'd need to extract this - right_expr = None # You'd need to extract this - return And( - convert_pyarrow_expression(left_expr), - convert_pyarrow_expression(right_expr), - ) - elif "or" in expr_str.lower() and isinstance(expr, pc.Expression): - # Similar simplification - left_expr = None # You'd need to extract this - right_expr = None # You'd need to extract this - return Or( - convert_pyarrow_expression(left_expr), - convert_pyarrow_expression(right_expr), - ) - elif "not" in expr_str.lower() and isinstance(expr, pc.Expression): - # Similar simplification - inner_expr = None # You'd need to extract this - return Not(convert_pyarrow_expression(inner_expr)) - - # Handle comparison operations - try: - if "==" in expr_str: - field_name, value = extract_field_and_literal(expr) - return EqualTo(Reference(field_name), value) - elif "!=" in expr_str: - field_name, value = extract_field_and_literal(expr) - return NotEqualTo(Reference(field_name), value) - elif ">=" in expr_str: - field_name, value = extract_field_and_literal(expr) - return GreaterThanOrEqual(Reference(field_name), value) - elif ">" in expr_str: - field_name, value = extract_field_and_literal(expr) - return GreaterThan(Reference(field_name), value) - elif "<=" in expr_str: - field_name, value = extract_field_and_literal(expr) - return LessThanOrEqual(Reference(field_name), value) - elif "<" in expr_str: - field_name, value = extract_field_and_literal(expr) - return LessThan(Reference(field_name), value) - else: - raise ValueError(f"Unsupported expression: {expr_str}") - except Exception as e: - raise ValueError(f"Failed to convert expression '{expr_str}': {str(e)}") - - -def convert_pyarrow_expressions(exprs: List[pc.Expression]) -> BooleanExpression: - """ - Convert a list of PyArrow expressions to a single PyIceberg expression. - Multiple expressions are combined with AND. - """ - if not exprs: - raise ValueError("No expressions provided") - - if len(exprs) == 1: - return convert_pyarrow_expression(exprs[0]) - - # Combine multiple expressions with AND - result = convert_pyarrow_expression(exprs[0]) - for expr in exprs[1:]: - result = And(result, convert_pyarrow_expression(expr)) - - return result diff --git a/tests/tower/test_table_filters.py b/tests/tower/test_table_filters.py new file mode 100644 index 00000000..b75c3b1c --- /dev/null +++ b/tests/tower/test_table_filters.py @@ -0,0 +1,167 @@ +import operator + +import pyarrow.compute as pc +import pytest +from pyiceberg.expressions import ( + And, + EqualTo, + GreaterThan, + GreaterThanOrEqual, + LessThan, + LessThanOrEqual, + NotEqualTo, + Or, +) +from pyiceberg.schema import Schema +from pyiceberg.types import IntegerType, NestedField, StringType, StructType + +import tower._tables as tables_module +from tower._context import TowerContext +from tower.exceptions import PyArrowFilterMigrationError + + +class FakeFilterTable: + def __init__(self): + self.delete_calls = [] + self._schema = Schema( + NestedField(1, "age", IntegerType(), required=False), + NestedField(2, "brand", StringType(), required=False), + NestedField(3, "origin", StringType(), required=False), + NestedField(4, "notice", StringType(), required=False), + NestedField( + 5, + "profile", + StructType( + NestedField(6, "name", StringType(), required=False), + ), + required=False, + ), + ) + + def schema(self): + return self._schema + + def delete(self, **kwargs): + self.delete_calls.append(kwargs) + + def refresh(self): + raise AssertionError("a successful delete must not refresh") + + +def make_table(): + context = TowerContext( + tower_url="https://api.example.com", + environment="production", + ) + iceberg_table = FakeFilterTable() + return tables_module.Table(context, iceberg_table), iceberg_table + + +@pytest.mark.parametrize( + ("comparison", "expected"), + [ + (lambda column: operator.eq(column, 18), EqualTo("age", 18)), + (lambda column: operator.ne(column, 18), NotEqualTo("age", 18)), + (lambda column: operator.gt(column, 18), GreaterThan("age", 18)), + (lambda column: operator.ge(column, 18), GreaterThanOrEqual("age", 18)), + (lambda column: operator.lt(column, 18), LessThan("age", 18)), + (lambda column: operator.le(column, 18), LessThanOrEqual("age", 18)), + ], +) +def test_table_column_builds_all_pyiceberg_comparisons(comparison, expected): + table, _ = make_table() + + assert comparison(table.column("age")) == expected + + +def test_table_column_expressions_compose_structurally(): + table, _ = make_table() + + expression = ( + (table.column("brand") == "candy or not") + & (table.column("origin") != "north and west") + ) | ~(table.column("notice") >= "not available") + + assert expression == Or( + And( + EqualTo("brand", "candy or not"), + NotEqualTo("origin", "north and west"), + ), + LessThan("notice", "not available"), + ) + + +def test_table_column_validates_nested_names_case_sensitively(): + table, _ = make_table() + + assert table.column("profile.name").name == "profile.name" + + with pytest.raises(ValueError, match="Column Profile.name not found"): + table.column("Profile.name") + + with pytest.raises(ValueError, match="Column profile.missing not found"): + table.column("profile.missing") + + +@pytest.mark.parametrize( + "delete_filter", + [ + "age >= 18 AND brand = 'candy'", + GreaterThanOrEqual("age", 18), + ], +) +def test_delete_forwards_canonical_filters_unchanged(delete_filter): + table, iceberg_table = make_table() + + result = table.delete(filters=delete_filter, max_retries=0) + + assert result is table + assert len(iceberg_table.delete_calls) == 1 + assert iceberg_table.delete_calls[0]["delete_filter"] is delete_filter + assert iceberg_table.delete_calls[0]["case_sensitive"] is True + + +@pytest.mark.parametrize( + "legacy_filter", + [ + pc.field("age") >= 18, + [pc.field("age") >= 18, pc.field("brand") == "candy"], + ], +) +def test_pyarrow_filters_raise_migration_error_before_write_escalation( + monkeypatch, legacy_filter +): + table, iceberg_table = make_table() + + def unexpected_escalation(mode): + raise AssertionError("invalid filters must fail before credential vending") + + monkeypatch.setattr(table, "_ensure_read_write_table", unexpected_escalation) + + with pytest.raises( + PyArrowFilterMigrationError, + match=r'table\.column\("age"\) >= 18.*a & b', + ): + table.delete(legacy_filter) + + assert iceberg_table.delete_calls == [] + + +@pytest.mark.parametrize("invalid_filter", [None, True, 42, ("age = 18",)]) +def test_delete_rejects_other_filter_types_before_write_escalation( + monkeypatch, invalid_filter +): + table, iceberg_table = make_table() + + def unexpected_escalation(mode): + raise AssertionError("invalid filters must fail before credential vending") + + monkeypatch.setattr(table, "_ensure_read_write_table", unexpected_escalation) + + with pytest.raises( + TypeError, + match="filters must be a SQL-like string or a PyIceberg BooleanExpression", + ): + table.delete(invalid_filter) + + assert iceberg_table.delete_calls == [] diff --git a/tests/tower/test_tables.py b/tests/tower/test_tables.py index a5663172..4bec5997 100644 --- a/tests/tower/test_tables.py +++ b/tests/tower/test_tables.py @@ -943,7 +943,7 @@ def test_delete_from_tables(in_memory_catalog): assert table.rows_affected().inserts == 3 # Perform the underlying delete from the table... - table.delete(filters=[table.column("username") == "bobb"]) + table.delete(filters=table.column("username") == "bobb") # ...and let's make sure that record is actually gone. df = table.to_polars()