Make count() agree with to_arrow() for NULL partitions

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)
 

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),

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()
+        )