Fix Parquet statistics pruning for predicates satisfied by NaN

cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp

@@ -400,8 +400,8 @@ class dictionary_literals_collector : public equality_literals_collector {
 };
 
 /**
- * @brief Converts named columns to index reference columns and pushes logical negations down to
- * expression leaves
+ * @brief Converts named columns to index reference columns and rewrites the expression into
+ * negation normal form, pushing logical negations down to the leaves
  */
 class parquet_filter_normalizer : public parquet::detail::parquet_filter_normalizer {
  public:

cpp/src/io/parquet/experimental/page_index_filter.cu

@@ -562,9 +562,8 @@ struct page_stats_to_row_mask_converter : public page_stats_caster {
 
       auto page_stats_table = cudf::table(std::move(columns));
       // Converts AST to StatsAST with reference to min, max columns in above `stats_table`.
-      auto constexpr num_columns = 1;
       parquet::detail::stats_expression_converter const stats_expr{
-        filter.get(), num_columns, has_is_null_operator, stream};
+        filter.get(), std::span<cudf::data_type const>{&dtype, 1}, has_is_null_operator, stream};
 
       // Filter the input table using AST expression and return the (BOOL8) predicate column.
       auto const page_mask = cudf::detail::compute_column(page_stats_table,
@@ -866,9 +865,7 @@ std::unique_ptr<cudf::column> aggregate_reader_metadata::build_row_mask_with_pag
 
   // Get a boolean mask indicating which columns will participate in stats based filtering
   auto const [stats_columns_mask, has_is_null_operator] =
-    parquet::detail::stats_columns_collector{filter.get(),
-                                             static_cast<size_type>(output_dtypes.size())}
-      .get_stats_columns_mask();
+    parquet::detail::stats_columns_collector{filter.get(), output_dtypes}.get_stats_columns_mask();
 
   // Return early if no columns will participate in stats based page filtering
   if (stats_columns_mask.empty()) { return build_all_true_row_mask(row_group_indices, stream, mr); }
@@ -971,7 +968,7 @@ std::unique_ptr<cudf::column> aggregate_reader_metadata::build_row_mask_with_pag
 
   // Converts AST to StatsAST with reference to min, max columns in above `stats_table`.
   parquet::detail::stats_expression_converter const stats_expr{
-    filter.get(), static_cast<size_type>(output_dtypes.size()), has_is_null_operator, stream};
+    filter.get(), output_dtypes, has_is_null_operator, stream};
 
   // Filter the input table using AST expression and return the (BOOL8) predicate column.
   return cudf::detail::compute_column(

cpp/src/io/parquet/expression_transform_helpers.cpp

@@ -50,6 +50,9 @@ namespace {
 template <operator_transform mode>
 std::optional<ast::ast_operator> transform_operator(ast::ast_operator op)
 {
+  static_assert(mode == operator_transform::INVERT or mode == operator_transform::NEGATE,
+                "Unhandled operator transform");
+
   if constexpr (mode == operator_transform::INVERT) {
     switch (op) {
       case ast::ast_operator::LESS: return ast::ast_operator::GREATER;

cpp/src/io/parquet/expression_transform_helpers.hpp

@@ -157,8 +157,8 @@ class names_from_expression : public ast::detail::expression_transformer {
 };
 
 /**
- * @brief Converts named columns to index reference columns and pushes logical negations down to the
- * leaves of the expression.
+ * @brief Converts named columns to index reference columns and rewrites the expression into
+ * negation normal form, pushing logical negations down to the leaves.
  */
 class parquet_filter_normalizer : public ast::detail::expression_transformer {
  public:

cpp/src/io/parquet/predicate_pushdown.cpp

@@ -69,8 +69,7 @@ std::optional<std::vector<std::vector<size_type>>> aggregate_reader_metadata::ap
 
   // Get a boolean mask indicating which columns can participate in stats based filtering
   auto const [stats_columns_mask, has_is_null_operator] =
-    stats_columns_collector{filter.get(), static_cast<size_type>(output_dtypes.size())}
-      .get_stats_columns_mask();
+    stats_columns_collector{filter.get(), output_dtypes}.get_stats_columns_mask();
 
   // Return early if no columns will participate in stats based filtering
   if (stats_columns_mask.empty()) { return std::nullopt; }
@@ -150,7 +149,7 @@ std::optional<std::vector<std::vector<size_type>>> aggregate_reader_metadata::ap
 
   // Converts AST to StatsAST with reference to min, max columns in above `stats_table`.
   stats_expression_converter const stats_expr{
-    filter.get(), static_cast<size_type>(output_dtypes.size()), has_is_null_operator, stream};
+    filter.get(), output_dtypes, has_is_null_operator, stream};
 
   // Filter stats table with StatsAST expression and collect filtered row group indices
   return collect_filtered_row_group_indices(

cpp/src/io/parquet/stats_filter_helpers.cpp

@@ -14,11 +14,16 @@
 
 namespace cudf::io::parquet::detail {
 
+stats_columns_collector::stats_columns_collector(std::span<cudf::data_type const> output_dtypes)
+  : _num_columns(static_cast<size_type>(output_dtypes.size())), _output_dtypes(output_dtypes)
+{
+  _columns_mask.resize(_num_columns, false);
+}
+
 stats_columns_collector::stats_columns_collector(ast::expression const& expr,
-                                                 cudf::size_type num_columns)
-  : _num_columns(num_columns)
+                                                 std::span<cudf::data_type const> output_dtypes)
+  : stats_columns_collector(output_dtypes)
 {
-  _columns_mask.resize(num_columns, false);
   expr.accept(*this);
 }
 
@@ -75,7 +80,12 @@ std::reference_wrapper<ast::expression const> stats_columns_collector::visit(
     if (op == ast_operator::EQUAL or op == ast_operator::NOT_EQUAL or op == ast_operator::LESS or
         op == ast_operator::LESS_EQUAL or op == ast_operator::GREATER or
         op == ast_operator::GREATER_EQUAL) {
-      _columns_mask[col_ref->get_column_index()] = true;
+      // NOT_EQUAL leaf for floating points relaxes to always true as Parquet statistics do not
+      // record NaNs.
+      if (op != ast_operator::NOT_EQUAL or
+          not cudf::is_floating_point(_output_dtypes[col_ref->get_column_index()])) {
+        _columns_mask[col_ref->get_column_index()] = true;
+      }
     }
   } else {
     // Visit the operands and ignore any output as we only want to build the column mask
@@ -89,15 +99,16 @@ std::pair<thrust::host_vector<bool>, bool> stats_columns_collector::get_stats_co
   return {std::move(_columns_mask), _has_is_null_operator};
 }
 
-stats_expression_converter::stats_expression_converter(ast::expression const& expr,
-                                                       size_type num_columns,
-                                                       bool has_is_null_operator,
-                                                       cuda::stream_ref stream)
-  : _always_true_scalar{std::make_unique<cudf::numeric_scalar<bool>>(true, true, stream)},
+stats_expression_converter::stats_expression_converter(
+  ast::expression const& expr,
+  std::span<cudf::data_type const> output_dtypes,
+  bool has_is_null_operator,
+  cuda::stream_ref stream)
+  : stats_columns_collector{output_dtypes},
+    _always_true_scalar{std::make_unique<cudf::numeric_scalar<bool>>(true, true, stream)},
     _always_true{std::make_unique<ast::literal>(*_always_true_scalar)}
 {
   _stats_cols_per_column = has_is_null_operator ? 3 : 2;
-  _num_columns           = num_columns;
   expr.accept(*this);
 }
 
@@ -153,18 +164,32 @@ std::reference_wrapper<ast::expression const> stats_expression_converter::visit(
             }
           }  // Binary operation wrapped
           else if (cudf::ast::detail::ast_operator_arity(child_op) == 2) {
+            // For NOT(col op lit) or NOT(lit op col), negate the operator if negatable and visit
+            // the negated operation directly.
             auto const binary_operands = extract_binary_operands(*child_operation);
             auto const lhs_kind        = binary_operands.lhs_type;
             auto const rhs_kind        = binary_operands.rhs_type;
 
-            // For NOT(col op lit) negate the operator if negatable and visit the negated operation
-            // directly
+            // `col_ref` is only non-null for the `col op lit` form, so both checks below must
+            // stay inside this branch
             if (lhs_kind == operand_kind::COLUMN_REF and rhs_kind == operand_kind::LITERAL) {
-              auto const negated_op = transform_operator<operator_transform::NEGATE>(child_op);
-              if (negated_op.has_value()) {
-                auto const& child_operands = child_operation->get_operands();
-                return visit(
-                  ast::operation{*negated_op, child_operands.front(), child_operands.back()});
+              // Equality is always exact
+              auto const is_equality =
+                child_op == ast_operator::EQUAL or child_op == ast_operator::NOT_EQUAL;
+
+              // An ordering comparison is only exact when the column cannot hold a `NaN`. i.e., not
+              // a floating point type
+              auto const can_negate_ordering = not cudf::is_floating_point(
+                _output_dtypes[binary_operands.col_ref->get_column_index()]);
+
+              if (is_equality or can_negate_ordering) {
+                auto const negated_op =
+                  transform_operator<operator_transform::NEGATE>(child_operation->get_operator());
+                if (negated_op.has_value()) {
+                  auto const& child_operands = child_operation->get_operands();
+                  return visit(
+                    ast::operation{*negated_op, child_operands.front(), child_operands.back()});
+                }
               }
             }
           }
@@ -210,6 +235,12 @@ std::reference_wrapper<ast::expression const> stats_expression_converter::visit(
         break;
       }
       case ast_operator::NOT_EQUAL: {
+        // Some Parquet writers exclude `NaN`s from stats so we can't reliably prune row groups for
+        // columns that may contain them.
+        if (cudf::is_floating_point(_output_dtypes[col_index])) {
+          _stats_expr.push(ast::operation{ast_operator::IDENTITY, *_always_true});
+          return *_always_true;
+        }
         auto const& vmin =
           _stats_expr.push(ast::column_reference{col_index * _stats_cols_per_column});
         auto const& vmax =

cpp/src/io/parquet/stats_filter_helpers.hpp

@@ -308,7 +308,8 @@ class stats_columns_collector : public ast::detail::expression_transformer {
  public:
   stats_columns_collector() = default;
 
-  stats_columns_collector(ast::expression const& expr, cudf::size_type num_columns);
+  stats_columns_collector(ast::expression const& expr,
+                          std::span<cudf::data_type const> output_dtypes);
 
   /**
    * @copydoc ast::detail::expression_transformer::visit(ast::literal const& )
@@ -340,7 +341,10 @@ class stats_columns_collector : public ast::detail::expression_transformer {
   std::pair<thrust::host_vector<bool>, bool> get_stats_columns_mask() &&;
 
  protected:
+  stats_columns_collector(std::span<cudf::data_type const> output_dtypes);
+
   size_type _num_columns;
+  std::span<cudf::data_type const> _output_dtypes;
 
  private:
   thrust::host_vector<bool> _columns_mask;
@@ -358,7 +362,7 @@ class stats_columns_collector : public ast::detail::expression_transformer {
 class stats_expression_converter : public stats_columns_collector {
  public:
   stats_expression_converter(ast::expression const& expr,
-                             size_type num_columns,
+                             std::span<cudf::data_type const> output_dtypes,
                              bool has_is_null_operator,
                              cuda::stream_ref stream);
 

cpp/tests/io/parquet_reader_test.cpp

@@ -2422,6 +2422,28 @@ TEST_F(ParquetReaderTest, FilterNegationPushdown)
     expect_matches_unrewritten(cudf::ast::operation(cudf::ast::ast_operator::NOT, not_lt), 1);
   }
 
+  // NOT(50 op col_a) - literal-left ordering comparisons preserve operand order when complemented.
+  for (auto const op : {cudf::ast::ast_operator::LESS,
+                        cudf::ast::ast_operator::LESS_EQUAL,
+                        cudf::ast::ast_operator::GREATER,
+                        cudf::ast::ast_operator::GREATER_EQUAL}) {
+    auto literal_left = cudf::ast::operation(op, lit_50, col_ref_a);
+    auto const expected_row_groups =
+      op == cudf::ast::ast_operator::LESS or op == cudf::ast::ast_operator::LESS_EQUAL ? 1 : 4;
+    expect_matches_unrewritten(cudf::ast::operation(cudf::ast::ast_operator::NOT, literal_left),
+                               expected_row_groups);
+  }
+
+  // NOT(col_a + 10 > 50) - operand is not `col op lit`, so it must NOT be complemented. The
+  // `col_a < 150` conjunct keeps the filter stats-usable.
+  {
+    auto sum       = cudf::ast::operation(cudf::ast::ast_operator::ADD, col_ref_a, lit_10);
+    auto sum_gt_50 = cudf::ast::operation(cudf::ast::ast_operator::GREATER, sum, lit_50);
+    auto not_sum   = cudf::ast::operation(cudf::ast::ast_operator::NOT, sum_gt_50);
+    expect_matches_unrewritten(
+      cudf::ast::operation(cudf::ast::ast_operator::LOGICAL_AND, a_lt_150, not_sum), 1);
+  }
+
   // Double negation over a non-boolean operand must NOT be eliminated.
   {
     auto not_a     = cudf::ast::operation(cudf::ast::ast_operator::NOT, col_ref_a);
@@ -4414,9 +4436,10 @@ void filter_unary_operation_typed_test()
     filter_expression = cudf::ast::operation(cudf::ast::ast_operator::LOGICAL_OR, not_expr1, expr2);
     ref_filter =
       cudf::ast::operation(cudf::ast::ast_operator::LOGICAL_OR, ref_not_expr1, ref_expr2);
-    // For signed numeric types, RGs 1,2,3 pass. Otherwise, RGs 2,3 pass
+    // Signed numeric types pass RGs 1,2,3, others pass RGs 2,3. Floats keep all 4: NaN makes
+    // every ordered comparison false, so `NOT(col0 < 100)` is not `col0 >= 100` and gets relaxed.
     auto constexpr expected_filtered_row_groups_with_unary_or =
-      (cudf::is_numeric<T>() and cudf::is_signed<T>()) ? 3 : 2;
+      cudf::is_floating_point<T>() ? 4 : ((cudf::is_numeric<T>() and cudf::is_signed<T>()) ? 3 : 2);
     test_predicate_pushdown(filter_expression,
                             ref_filter,
                             expected_total_row_groups,

python/cudf/cudf/tests/input_output/test_parquet.py

@@ -4721,6 +4721,88 @@ def test_parquet_reader_mismatched_nullability_structs(tmp_path):
     )
 
 
+def test_parquet_not_equal_with_nan_stats(tmp_path):
+    """`col != v` must not prune matching `NaN` rows."""
+    import pylibcudf as plc
+    from pylibcudf.expressions import (
+        ASTOperator,
+        ColumnNameReference,
+        Literal,
+        Operation,
+    )
+
+    path = tmp_path / "nan_not_equal.parquet"
+    pq.write_table(
+        pa.table({"x": [float("nan"), 5.0, 7.0, 8.0]}), path, row_group_size=2
+    )
+
+    # Sanity check the fixture: NaN is excluded, so row group 0 looks constant
+    stats = pq.ParquetFile(path).metadata.row_group(0).column(0).statistics
+    assert_eq(stats.min, 5.0)
+    assert_eq(stats.max, 5.0)
+
+    scalar = plc.Scalar.from_arrow(pa.scalar(5.0))
+    filter_expr = Operation(
+        ASTOperator.NOT_EQUAL, ColumnNameReference("x"), Literal(scalar)
+    )
+
+    source = plc.io.SourceInfo([str(path)])
+    options = plc.io.parquet.ParquetReaderOptions.builder(source).build()
+    options.set_filter(filter_expr)
+    result = plc.io.parquet.read_parquet(options)
+
+    # Neither row group may be pruned: rg0 holds a NaN, rg1 holds 7.0 and 8.0
+    assert_eq(result.num_row_groups_after_stats_filter, 2)
+    got = result.tbl.to_arrow().column(0).to_pylist()
+    assert_eq(len(got), 3)
+    assert_eq(math.isnan(got[0]), True)
+    assert_eq(got[1:], [7.0, 8.0])
+
+
+def test_parquet_negated_ordering_with_nan_stats(tmp_path):
+    """`NOT(col < v)` must not prune matching `NaN` rows."""
+    import pylibcudf as plc
+    from pylibcudf.expressions import (
+        ASTOperator,
+        ColumnNameReference,
+        Literal,
+        Operation,
+    )
+
+    # One row group per 3 rows. The first holds NaN alongside small values, so its
+    # statistics are min=1.0/max=2.0 and `vmax >= 50` is false for it.
+    values = [float("nan"), 1.0, 2.0, 100.0, 200.0, 300.0]
+    path = tmp_path / "nan_ordering.parquet"
+    pq.write_table(pa.table({"x": values}), path, row_group_size=3)
+
+    # Sanity check the fixture actually reproduces the Arrow statistics behaviour
+    stats = pq.ParquetFile(path).metadata.row_group(0).column(0).statistics
+    assert_eq(stats.has_min_max, True)
+    assert_eq(stats.min, 1.0)
+    assert_eq(stats.max, 2.0)
+
+    col = ColumnNameReference("x")
+    lit = Literal(plc.Scalar.from_arrow(pa.scalar(50.0)))
+    filter_expr = Operation(
+        ASTOperator.NOT, Operation(ASTOperator.LESS, col, lit)
+    )
+
+    source = plc.io.SourceInfo([str(path)])
+    options = plc.io.parquet.ParquetReaderOptions.builder(source).build()
+    options.set_filter(filter_expr)
+    got = (
+        plc.io.parquet.read_parquet(options)
+        .tbl.to_arrow()
+        .column(0)
+        .to_pylist()
+    )
+
+    # NOT(x < 50) is true for NaN and for 100/200/300, and false for 1.0/2.0
+    assert_eq(len(got), 4)
+    assert_eq(math.isnan(got[0]), True)
+    assert_eq(got[1:], [100.0, 200.0, 300.0])
+
+
 @pytest.mark.skipif(
     pa.__version__ == "19.0.0",
     reason="https://github.com/apache/arrow/issues/45283, https://github.com/NVIDIA/cudf/issues/17806",