Fix Parquet statistics pruning for predicates satisfied by NaN

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

@@ -395,8 +395,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{&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

@@ -48,6 +48,7 @@ namespace {
 }  // namespace
 
 template <operator_transform mode>
+  requires(mode == operator_transform::INVERT or mode == operator_transform::NEGATE)
 std::optional<ast::ast_operator> transform_operator(ast::ast_operator op)
 {
   if constexpr (mode == operator_transform::INVERT) {

cpp/src/io/parquet/expression_transform_helpers.hpp

@@ -88,6 +88,7 @@ enum class operator_transform : uint8_t {
  * untransformable operators are returned as is (no std::nullopt)
  */
 template <operator_transform mode>
+  requires(mode == operator_transform::INVERT or mode == operator_transform::NEGATE)
 [[nodiscard]] std::optional<ast::ast_operator> transform_operator(ast::ast_operator op);
 
 /**
@@ -157,8 +158,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,45 @@
 
 namespace cudf::io::parquet::detail {
 
+namespace {
+
+/**
+ * @brief Returns whether a comparison operator can prune row groups via statistics
+ *
+ * Some Parquet writers exclude `NaN`s from stats, so a floating-point chunk holding a NaN is
+ * indistinguishable from one that does not. `col != val` is the only comparison leaf a NaN
+ * satisfies, so it is the only one we cannot prune.
+ *
+ * @param op The comparison operator
+ * @param dtype The data type of the column being compared
+ * @return true if the comparison can be used to prune row groups
+ */
+[[nodiscard]] bool is_prunable_comparison(ast::ast_operator op, cudf::data_type dtype)
+{
+  using cudf::ast::ast_operator;
+  switch (op) {
+    case ast_operator::EQUAL: [[fallthrough]];
+    case ast_operator::LESS: [[fallthrough]];
+    case ast_operator::LESS_EQUAL: [[fallthrough]];
+    case ast_operator::GREATER: [[fallthrough]];
+    case ast_operator::GREATER_EQUAL: return true;
+    case ast_operator::NOT_EQUAL: return not cudf::is_floating_point(dtype);
+    default: return false;
+  }
+}
+
+}  // namespace
+
+stats_columns_collector::stats_columns_collector(std::span<cudf::data_type const> output_dtypes)
+  : _output_dtypes(output_dtypes)
+{
+  _columns_mask.resize(_output_dtypes.size(), 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);
 }
 
@@ -33,7 +67,7 @@ std::reference_wrapper<ast::expression const> stats_columns_collector::visit(
 {
   CUDF_EXPECTS(expr.get_table_source() == ast::table_reference::LEFT,
                "Statistics AST supports only left table");
-  CUDF_EXPECTS(expr.get_column_index() < _num_columns,
+  CUDF_EXPECTS(static_cast<size_t>(expr.get_column_index()) < _output_dtypes.size(),
                "Column index cannot be more than number of columns in the table");
   return expr;
 }
@@ -72,11 +106,8 @@ std::reference_wrapper<ast::expression const> stats_columns_collector::visit(
 
   if (lhs_kind == operand_kind::COLUMN_REF and rhs_kind == operand_kind::LITERAL) {
     col_ref->accept(*this);
-    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;
-    }
+    auto const col_index = col_ref->get_column_index();
+    if (is_prunable_comparison(op, _output_dtypes[col_index])) { _columns_mask[col_index] = true; }
   } else {
     // Visit the operands and ignore any output as we only want to build the column mask
     std::ignore = visit_operands(expr.get_operands());
@@ -89,15 +120,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);
 }
 
@@ -132,9 +164,11 @@ std::reference_wrapper<ast::expression const> stats_expression_converter::visit(
         return *_always_true;
       }
     } else {
-      // Special handling for the NOT operator since is necessary as stats transforms use different
-      // columns (vmin, vmax, is_null) for different operators such that NOT(col < val) is not
-      // equivalent to NOT(vmin < val) and instead is equivalent to vmax >= val.
+      // `parquet_filter_normalizer::push_down_negation` deliberately does not complement ordering
+      // comparisons (NaN makes `NOT(a < b)` differ from `a >= b`), so `NOT(col op lit)` forms
+      // reach here. Stats transforms use different columns (vmin, vmax, is_null) for different
+      // operators such that NOT(col < val) is not equivalent to NOT(vmin < val) and instead is
+      // equivalent to vmax >= val.
       if (input_op == ast_operator::NOT) {
         auto const* child_operation =
           dynamic_cast<ast::operation const*>(&expr.get_operands().front().get());
@@ -153,18 +187,27 @@ 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()});
+              binary_operands.col_ref->accept(*this);
+
+              // A comparison cannot be negated when the column may hold a `NaN` (floating points).
+              if (not cudf::is_floating_point(
+                    _output_dtypes[binary_operands.col_ref->get_column_index()])) {
+                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()});
+                }
               }
             }
           }
@@ -186,13 +229,21 @@ std::reference_wrapper<ast::expression const> stats_expression_converter::visit(
     col_ref->accept(*this);
 
     auto const col_index = col_ref->get_column_index();
+
+    // Some Parquet writers exclude `NaN`s from stats, so we can't reliably prune row groups for
+    // columns that may contain them.
+    if (not is_prunable_comparison(op, _output_dtypes[col_index])) {
+      _stats_expr.push(ast::operation{ast_operator::IDENTITY, *_always_true});
+      return *_always_true;
+    }
+
     // Push literal into the ast::tree
     auto const& literal = _stats_expr.push(*literal_ptr);
 
     switch (op) {
       /* transform to stats conditions
       col == val --> vmin <= val && vmax >= val
-      col != val --> !(vmin == val && vmax == val)
+      col != val --> vmin != vmax || vmax != val
       col >  val --> vmax > val
       col <  val --> vmin < val
       col >= val --> vmax >= val
@@ -234,10 +285,7 @@ std::reference_wrapper<ast::expression const> stats_expression_converter::visit(
         _stats_expr.push(ast::operation{op, vmax, literal});
         break;
       }
-      default: {
-        _stats_expr.push(ast::operation{ast_operator::IDENTITY, *_always_true});
-        return *_always_true;
-      }
+      default: CUDF_UNREACHABLE("Non-prunable operator should not reach stats conversion");
     };
   }  // Visit operands and push expression for `expr op expr` form
   else if (lhs_kind == operand_kind::EXPRESSION and rhs_kind == operand_kind::EXPRESSION) {

cpp/src/io/parquet/stats_filter_helpers.hpp

@@ -306,9 +306,8 @@ class stats_caster_base {
  */
 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 +339,9 @@ class stats_columns_collector : public ast::detail::expression_transformer {
   std::pair<thrust::host_vector<bool>, bool> get_stats_columns_mask() &&;
 
  protected:
-  size_type _num_columns;
+  explicit stats_columns_collector(std::span<cudf::data_type const> output_dtypes);
+
+  std::span<cudf::data_type const> _output_dtypes;
 
  private:
   thrust::host_vector<bool> _columns_mask;
@@ -358,7 +359,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

@@ -2371,8 +2371,12 @@ TEST_F(ParquetReaderTest, FilterNegationPushdown)
   auto const floats = cudf::detail::make_counting_transform_iterator(
     0, [](auto i) { return i % 2 == 0 ? NAN : static_cast<float>(i); });
   auto col_c = cudf::test::fixed_width_column_wrapper<float>(floats, floats + num_rows);
+  // NaN-free float column so that stats are generated
+  auto const clean_floats =
+    cudf::detail::make_counting_transform_iterator(0, [](auto i) { return static_cast<float>(i); });
+  auto col_d = cudf::test::fixed_width_column_wrapper<float>(clean_floats, clean_floats + num_rows);
 
-  auto const written_table = cudf::table_view{{col_a, col_b, col_c}};
+  auto const written_table = cudf::table_view{{col_a, col_b, col_c, col_d}};
   auto const filepath      = temp_env->get_temp_filepath("FilterNegationPushdown.parquet");
   cudf::io::parquet_writer_options const out_opts =
     cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, written_table)
@@ -2398,6 +2402,7 @@ TEST_F(ParquetReaderTest, FilterNegationPushdown)
   auto col_ref_a = cudf::ast::column_reference(0);
   auto col_ref_b = cudf::ast::column_reference(1);
   auto col_ref_c = cudf::ast::column_reference(2);
+  auto col_ref_d = cudf::ast::column_reference(3);
 
   auto lit_10_value  = cudf::numeric_scalar<int32_t>(10);
   auto lit_10        = cudf::ast::literal(lit_10_value);
@@ -2423,6 +2428,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);
@@ -2480,6 +2507,12 @@ TEST_F(ParquetReaderTest, FilterNegationPushdown)
     expect_matches_unrewritten(cudf::ast::operation(cudf::ast::ast_operator::NOT, conjunction));
   }
 
+  // NOT(col_d < NaN), ordered comparisons against NaN are always false, even with valid stats.
+  {
+    auto d_lt_nan = cudf::ast::operation(cudf::ast::ast_operator::LESS, col_ref_d, lit_nan);
+    expect_matches_unrewritten(cudf::ast::operation(cudf::ast::ast_operator::NOT, d_lt_nan), 4);
+  }
+
   // Operators with no complement are left intact
   {
     auto a_is_null = cudf::ast::operation(cudf::ast::ast_operator::IS_NULL, col_ref_a);
@@ -4460,9 +4493,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 integral types pass RGs 1,2,3, others pass RGs 2,3. Floats keep all 4 as they may
+    // hold NaNs, making every ordered comparison false, and get relaxed instead.
     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,