diff --git a/pyiceberg/expressions/visitors.py b/pyiceberg/expressions/visitors.py index a8de088669..7504548507 100644 --- a/pyiceberg/expressions/visitors.py +++ b/pyiceberg/expressions/visitors.py @@ -50,7 +50,9 @@ BoundStartsWith, BoundTerm, BoundUnaryPredicate, + IsNull, Not, + NotNull, Or, UnboundPredicate, ) @@ -1969,36 +1971,29 @@ def visit_bound_predicate(self, predicate: BoundPredicate) -> BooleanExpression: return predicate for part in parts: - strict_projection = part.transform.strict_project(part.name, predicate) - strict_result = None - - if strict_projection is not None: - bound = strict_projection.bind(self.partition_schema, case_sensitive=self.case_sensitive) - if isinstance(bound, BoundPredicate): - strict_result = super().visit_bound_predicate(bound) - else: - # if the result is not a predicate, then it must be a constant like alwaysTrue or alwaysFalse - strict_result = bound - - if isinstance(strict_result, AlwaysTrue): + if isinstance(self._eval_projection(part.transform.strict_project(part.name, predicate)), AlwaysTrue): return AlwaysTrue() - inclusive_projection = part.transform.project(part.name, predicate) - inclusive_result = None - if inclusive_projection is not None: - bound_inclusive = inclusive_projection.bind(self.partition_schema, case_sensitive=self.case_sensitive) - if isinstance(bound_inclusive, BoundPredicate): - # using predicate method specific to inclusive - inclusive_result = super().visit_bound_predicate(bound_inclusive) - else: - # if the result is not a predicate, then it must be a constant like alwaysTrue or - # alwaysFalse - inclusive_result = bound_inclusive - if isinstance(inclusive_result, AlwaysFalse): + if isinstance(self._eval_projection(part.transform.project(part.name, predicate)), AlwaysFalse): return AlwaysFalse() return predicate + def _eval_projection(self, projection: UnboundPredicate | None) -> BooleanExpression | None: + if projection is None: + return None + + bound = projection.bind(self.partition_schema, case_sensitive=self.case_sensitive) + if not isinstance(bound, BoundPredicate): + # if the result is not a predicate, then it must be a constant like alwaysTrue or alwaysFalse + return bound + + # A NULL partition value only decides null checks; leave other predicates to the reader's NULL semantics + if bound.as_unbound not in (IsNull, NotNull) and bound.term.eval(self.struct) is None: + return None + + return super().visit_bound_predicate(bound) + def visit_unbound_predicate(self, predicate: UnboundPredicate) -> BooleanExpression: bound = predicate.bind(self.schema, case_sensitive=self.case_sensitive) diff --git a/tests/catalog/test_catalog_behaviors.py b/tests/catalog/test_catalog_behaviors.py index 6d694a9448..8115a3bfc4 100644 --- a/tests/catalog/test_catalog_behaviors.py +++ b/tests/catalog/test_catalog_behaviors.py @@ -47,7 +47,7 @@ from pyiceberg.table.snapshots import Operation from pyiceberg.table.sorting import NullOrder, SortDirection, SortField, SortOrder from pyiceberg.table.update import AddSchemaUpdate, SetCurrentSchemaUpdate -from pyiceberg.transforms import IdentityTransform +from pyiceberg.transforms import BucketTransform, IdentityTransform from pyiceberg.typedef import Identifier from pyiceberg.types import BooleanType, DoubleType, IntegerType, LongType, NestedField, StringType @@ -1376,6 +1376,73 @@ def test_append_nan_to_identity_partitioned_table(catalog: Catalog) -> None: assert sorted(tbl.scan(row_filter="value is nan").to_arrow()["id"].to_pylist()) == [2, 4] +def _table_with_null_partition(catalog: Catalog, name: str, partition_spec: PartitionSpec) -> Table: + schema = Schema(NestedField(1, "p", StringType(), required=False), NestedField(2, "f", DoubleType(), required=False)) + catalog.create_namespace_if_not_exists("default") + tbl = catalog.create_table(identifier=f"default.{name}_{catalog.name}", schema=schema, partition_spec=partition_spec) + tbl.append( + pa.Table.from_pydict( + { + "p": ["a", "a", "b", "b", "c", "c", None, None, "b", "b"], + "f": [1.0, 1.0, 2.0, 2.0, float("nan"), float("nan"), 1.0, 1.0, None, None], + }, + schema=schema_to_pyarrow(schema), + ) + ) + return tbl + + +def test_count_matches_to_arrow_with_null_partition_values(catalog: Catalog) -> None: + identity = PartitionSpec( + PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="p"), + PartitionField(source_id=2, field_id=1001, transform=IdentityTransform(), name="f"), + ) + bucket = PartitionSpec( + PartitionField(source_id=1, field_id=1000, transform=BucketTransform(4), name="p_bucket"), + PartitionField(source_id=2, field_id=1001, transform=IdentityTransform(), name="f"), + ) + evolved = _table_with_null_partition( + catalog, + "count_null_partition_evolved", + PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="p")), + ) + # Files written under the old spec keep a NULL p partition value that the current spec no longer has + with evolved.update_spec() as update: + update.remove_field("p") + evolved.append(evolved.scan().to_arrow()) + tables = [ + _table_with_null_partition(catalog, "count_null_partition_identity", identity), + _table_with_null_partition(catalog, "count_null_partition_bucket", bucket), + evolved, + ] + row_filters = ["p != 'a'", "p not in ('a', 'c')", "not p = 'a'", "p not like 'a%'", "f is not nan", "p is null", "true"] + for tbl in tables: + for row_filter in row_filters: + scan = tbl.scan(row_filter=row_filter) + assert scan.count() == len(scan.to_arrow()), (tbl.name(), row_filter) + + +def test_count_reads_null_partition_only_when_filter_compares_it(catalog: Catalog) -> None: + partition_spec = PartitionSpec( + PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="p"), + PartitionField(source_id=2, field_id=1001, transform=IdentityTransform(), name="f"), + ) + # Each row filter maps to (filter selecting files to delete, expected count), and count() must not read those files + cases = { + "true": ("true", 10), + "p is null": ("p is null", 2), + "p is not null or f = 1.0": ("p is null", 10), + "p != 'a'": ("p is not null", 6), + "f = 1.0": ("f = 1.0", 4), + } + for i, (row_filter, (metadata_only_filter, expected)) in enumerate(cases.items()): + tbl = _table_with_null_partition(catalog, f"count_metadata_only_{i}", partition_spec) + # Deleting those files makes any read of them fail + for task in tbl.scan(row_filter=metadata_only_filter).plan_files(): + tbl.io.delete(task.file.file_path) + assert tbl.scan(row_filter=row_filter).count() == expected, row_filter + + def _nan_partitioned_table(catalog: Catalog, name: str) -> Table: iceberg_schema = Schema( NestedField(1, "id", IntegerType(), required=False), diff --git a/tests/expressions/test_residual_evaluator.py b/tests/expressions/test_residual_evaluator.py index 4c6295e9eb..abcf1fb8ce 100644 --- a/tests/expressions/test_residual_evaluator.py +++ b/tests/expressions/test_residual_evaluator.py @@ -29,6 +29,8 @@ IsNaN, IsNull, LessThan, + Not, + NotEqualTo, NotIn, NotNaN, NotNull, @@ -40,7 +42,7 @@ from pyiceberg.expressions.visitors import ResidualVisitor, residual_evaluator_of from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema -from pyiceberg.transforms import DayTransform, IdentityTransform +from pyiceberg.transforms import BucketTransform, DayTransform, IdentityTransform from pyiceberg.typedef import Record from pyiceberg.types import DoubleType, FloatType, IntegerType, NestedField, StringType, TimestampType @@ -233,7 +235,7 @@ def test_is_not_nan() -> None: res_eval = residual_evaluator_of(spec=spec, expr=predicate, case_sensitive=True, schema=schema) residual = res_eval.residual_for(Record(None)) - assert residual == AlwaysTrue() + assert residual == predicate residual = res_eval.residual_for(Record(float("nan"))) assert residual == AlwaysFalse() @@ -248,7 +250,7 @@ def test_is_not_nan() -> None: res_eval = residual_evaluator_of(spec=spec, expr=predicate, case_sensitive=True, schema=schema) residual = res_eval.residual_for(Record(None)) - assert residual == AlwaysTrue() + assert residual == predicate residual = res_eval.residual_for(Record(float("nan"))) assert residual == AlwaysFalse() @@ -303,3 +305,23 @@ def test_not_starts_with() -> None: assert res_eval.residual_for(Record("abc")) == AlwaysFalse() assert res_eval.residual_for(Record("a")) == AlwaysFalse() assert res_eval.residual_for(Record("zoo")) == AlwaysTrue() + + +def test_null_partition_value_only_decides_null_checks() -> None: + schema = Schema(NestedField(1, "x", StringType(), required=False)) + identity = PartitionSpec(PartitionField(1, 1001, IdentityTransform(), "x_part")) + bucket = PartitionSpec(PartitionField(1, 1001, BucketTransform(4), "x_bucket")) + + for spec in [identity, bucket]: + for predicate in [NotEqualTo("x", "a"), Not(EqualTo("x", "a")), NotIn("x", ["a", "c"]), NotStartsWith("x", "a")]: + res_eval = residual_evaluator_of(spec=spec, expr=predicate, case_sensitive=True, schema=schema) + assert res_eval.residual_for(Record(None)) != AlwaysTrue(), (spec, predicate) + + assert ( + residual_evaluator_of(spec=spec, expr=IsNull("x"), case_sensitive=True, schema=schema).residual_for(Record(None)) + == AlwaysTrue() + ) + assert ( + residual_evaluator_of(spec=spec, expr=NotNull("x"), case_sensitive=True, schema=schema).residual_for(Record(None)) + == AlwaysFalse() + )