rahulsmahadev commented on code in PR #4346:
URL: https://github.com/apache/iceberg-python/pull/4346#discussion_r4542944033
##########
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:
Agreed, removed the NOT handling and the `negated` flag in 191387d0, so NOT
is not treated as NaN-sensitive here. I'll put the `rewrite_not` change before
`expression_to_pyarrow` in a separate PR.
##########
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:
Removed in 191387d0, along with the NOT over an integer column case, which
only existed for 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:
Yes, apache/arrow#51491. It covers `!=`, negated comparisons and `is_nan`
when min == max. Linked it in the comment in 191387d0.
--
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]
