Fokko commented on code in PR #4346:
URL: https://github.com/apache/iceberg-python/pull/4346#discussion_r4542670789
##########
pyiceberg/io/pyarrow.py:
##########
@@ -1721,6 +1733,29 @@ def _get_column_projection_values(
return projected_missing_fields
+def _has_nan_sensitive_predicate(expr: BooleanExpression, negated: bool =
False) -> bool:
+ """Return whether a predicate in the expression can match NaN values that
pyarrow may prune away.
+
+ Parquet statistics exclude NaN, so a row group whose non-NaN values are
all equal reports min == max.
+ pyarrow then prunes it for predicates that the NaN values satisfy: IsNaN,
NotEqualTo and NotIn on a
+ float or double column, and any predicate on such a column under a NOT.
+ """
+ if isinstance(expr, (And, Or)):
+ return _has_nan_sensitive_predicate(expr.left, negated) or
_has_nan_sensitive_predicate(expr.right, negated)
+ if isinstance(expr, Not):
Review Comment:
I'm not sure NOT belongs here. The metrics evaluator applies `rewrite_not`,
so `not x < 6.0` is evaluated as `x >= 6.0` at the manifest level and a file
with upper bound 5.0 is skipped before we get here. With this change, the NaN
row comes back for `not x < 6.0` only when the file has another row group that
survives, and `x >= 6.0` never returns it. Could we instead apply `rewrite_not`
to the bound filter before `expression_to_pyarrow`, in a separate PR? Then NOT
is never NaN-sensitive and the `negated` flag can go.
##########
tests/io/test_pyarrow.py:
##########
@@ -4198,6 +4203,61 @@ def
test_task_to_record_batches_filter_after_positional_deletes_empty_result(tmp
assert result_batches == []
+def
test_task_to_record_batches_keeps_nan_rows_for_nan_sensitive_filters(tmpdir:
str) -> None:
+ """Regression test: pyarrow must not prune a row group that holds NaN rows
matching the filter.
+
+ Parquet statistics exclude NaN, so the row group [NaN, 5.0, 5.0] reports
min == max == 5.0.
+ pyarrow prunes such a row group for IsNaN, NotEqualTo, NotIn and a NOT
over a float comparison,
+ which drops the NaN row. NotNaN and a NOT over an integer column are
controls.
+ """
+ from pyiceberg.expressions.visitors import bind
+
+ arrow_schema = pa.schema(
+ (
+ pa.field("id", pa.int32(), nullable=True,
metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"}),
+ pa.field("x", pa.float64(), nullable=True,
metadata={PYARROW_PARQUET_FIELD_ID_KEY: "2"}),
+ )
+ )
+ arrow_table = pa.table(
+ {"id": pa.array([0, 1, 2], type=pa.int32()), "x":
pa.array([float("nan"), 5.0, 5.0], type=pa.float64())},
+ schema=arrow_schema,
+ )
+ data_file =
_write_table_to_data_file(f"{tmpdir}/nan_row_group_equal_bounds.parquet",
arrow_schema, arrow_table)
+
+ table_schema = Schema(
+ NestedField(1, "id", IntegerType(), required=False),
+ NestedField(2, "x", DoubleType(), required=False),
+ )
+ projected_schema = Schema(NestedField(1, "id", IntegerType(),
required=False))
+
+ cases: list[tuple[BooleanExpression, list[int]]] = [
+ (IsNaN("x"), [0]),
+ (NotEqualTo("x", 5.0), [0]),
+ (NotIn("x", {5.0, 6.0}), [0]),
+ (And(IsNaN("x"), EqualTo("id", 0)), [0]),
+ (Or(IsNaN("x"), EqualTo("id", 2)), [0, 2]),
+ (Not(LessThan("x", 6.0)), [0]),
Review Comment:
This locks in that NaN matches `NOT(x < 6.0)`, which disagrees with how the
manifest evaluator treats it (`x >= 6.0`, where NaN does not match). I'd remove
this case along with the NOT handling.
##########
pyiceberg/io/pyarrow.py:
##########
@@ -1761,21 +1796,35 @@ def _task_to_record_batches(
)
pyarrow_filter = None
+ scanner_filter = None
+ # Set when the scanner does not apply the whole filter, so each batch
is filtered after the read
+ filter_batches = False
if bound_row_filter is not AlwaysTrue():
translated_row_filter = translate_column_names(
bound_row_filter, file_schema, case_sensitive=case_sensitive,
projected_field_values=projected_missing_fields
)
bound_file_filter = bind(file_schema, translated_row_filter,
case_sensitive=case_sensitive)
pyarrow_filter = expression_to_pyarrow(bound_file_filter,
file_schema)
+ if positional_deletes:
+ # Positional deletes must be applied before any row is filtered
+ filter_batches = True
+ elif _has_nan_sensitive_predicate(bound_file_filter):
+ # Parquet statistics exclude NaN, so pyarrow can prune a row
group that holds NaN rows matching
Review Comment:
Is there an upstream Arrow issue for pruning on statistics that ignore NaN?
It would be nice to link it here, so we know when this workaround can be
removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
