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

pyiceberg/table/__init__.py

@@ -34,15 +34,28 @@
 from pydantic import Field
 
 from pyiceberg.exceptions import CommitFailedException, ValidationException
-from pyiceberg.expressions import AlwaysFalse, AlwaysTrue, And, BooleanExpression, EqualTo, IsNull, Or, Reference
+from pyiceberg.expressions import (
+    AlwaysFalse,
+    AlwaysTrue,
+    And,
+    BooleanExpression,
+    BoundPredicate,
+    EqualTo,
+    IsNull,
+    NotNull,
+    Or,
+    Reference,
+)
 from pyiceberg.expressions.visitors import (
     ResidualEvaluator,
+    _ExpressionFieldIDs,
     _InclusiveMetricsEvaluator,
     bind,
     expression_evaluator,
     extract_field_ids,
     inclusive_projection,
     manifest_evaluator,
+    visit,
 )
 from pyiceberg.io import FileIO, load_file_io
 from pyiceberg.manifest import (
@@ -2429,6 +2442,14 @@ def _to_arrow_batch_reader_via_file_scan_tasks(
     return pa.RecordBatchReader.from_batches(target_schema, batches).cast(target_schema)
 
 
+class _FieldIDsOutsideNullChecks(_ExpressionFieldIDs):
+    # Projected null checks preserve the source's NULL status, and void does not project predicates
+    def visit_bound_predicate(self, predicate: BoundPredicate) -> set[int]:
+        if predicate.as_unbound in (IsNull, NotNull):
+            return set()
+        return super().visit_bound_predicate(predicate)
+
+
 class DataScan(TableScan):
     @cached_property
     def _manifest_planner(self) -> ManifestGroupPlanner:
@@ -2557,12 +2578,31 @@ def count(self) -> int:
         res = 0
         # every task is a FileScanTask
         tasks = self.plan_files()
+        bound_filter = bind(self.table_metadata.schema(), self.row_filter, case_sensitive=self.case_sensitive)
+        field_ids_outside_null_checks = visit(bound_filter, _FieldIDsOutsideNullChecks())
+        specs = self.table_metadata.specs() if field_ids_outside_null_checks else {}
+        referenced_positions = {
+            spec_id: [pos for pos, field in enumerate(spec.fields) if field.source_id in field_ids_outside_null_checks]
+            for spec_id, spec in specs.items()
+        }
 
         for task in tasks:
             # task.residual is a Boolean Expression if the filter condition is fully satisfied by the
             # partition value and task.delete_files represents that positional delete haven't been merged yet
             # hence those files have to read as a pyarrow table applying the filter and deletes
-            if task.residual == AlwaysTrue() and len(task.delete_files) == 0:
+            if (
+                task.residual == AlwaysTrue()
+                and len(task.delete_files) == 0
+                # A NULL partition value can leave an AlwaysTrue residual although the filter rejects its rows
+                and (
+                    not field_ids_outside_null_checks
+                    or (
+                        task.file.spec_id in specs
+                        and len(task.file.partition) == len(specs[task.file.spec_id].fields)
+                        and all(task.file.partition[pos] is not None for pos in referenced_positions[task.file.spec_id])
+                    )
+                )
+            ):
                 # Every File has a metadata stat that stores the file record count
                 res += task.file.record_count
             else:

tests/catalog/test_catalog_behaviors.py

@@ -42,11 +42,11 @@
 from pyiceberg.io.pyarrow import _dataframe_to_data_files, schema_to_pyarrow
 from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
-from pyiceberg.table import TableProperties
+from pyiceberg.table import Table, TableProperties
 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
 
@@ -1345,6 +1345,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 test_scan_integer_column_with_decimal_literal(catalog: Catalog) -> None:
     catalog.create_namespace("default")
     identifier = f"default.scan_integer_decimal_literal_{catalog.name}"

tests/catalog/test_scan_planning_models.py

@@ -14,8 +14,11 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from pathlib import Path
 from typing import Any
 
+import pyarrow as pa
+import pyarrow.parquet as pq
 import pytest
 from pydantic import TypeAdapter, ValidationError
 from requests_mock import Mocker
@@ -39,6 +42,7 @@
     ValueMap,
 )
 from pyiceberg.expressions import AlwaysTrue, EqualTo, Reference
+from pyiceberg.io.pyarrow import schema_to_pyarrow
 from pyiceberg.manifest import FileFormat
 from pyiceberg.table import FileScanTask
 
@@ -832,6 +836,36 @@ def test_scan_uses_server_side_planning_from_load_table_config(
     assert plan_mock.call_count == 1
 
 
+def test_count_reads_server_planned_file_without_partition_values(
+    rest_scan_catalog: RestCatalog, requests_mock: Mocker, example_table_metadata_v2: dict[str, Any], tmp_path: Path
+) -> None:
+    _mock_load_table(requests_mock, example_table_metadata_v2, {})
+    table = rest_scan_catalog.load_table(("db", "tbl"))
+    data_file = tmp_path / "file1.parquet"
+    pq.write_table(
+        pa.table({"x": [1, 1, 1], "y": [4, 5, 6], "z": [7, 8, 9]}, schema=schema_to_pyarrow(table.schema())), data_file
+    )
+    requests_mock.post(
+        f"{TEST_URI}v1/namespaces/db/tables/tbl/plan",
+        json={
+            "status": "completed",
+            "plan-id": "plan-123",
+            "delete-files": [],
+            "file-scan-tasks": [
+                {
+                    "data-file": _rest_data_file(
+                        file_path=str(data_file), file_size_in_bytes=data_file.stat().st_size, record_count=3
+                    )
+                }
+            ],
+            "plan-tasks": [],
+        },
+        status_code=200,
+    )
+
+    assert table.scan(row_filter="x != 1").count() == 0
+
+
 def test_scan_keeps_client_side_planning_without_load_table_config(
     rest_client_planning_catalog: RestCatalog, requests_mock: Mocker, example_table_metadata_v2: dict[str, Any]
 ) -> None: