Make Parquet statistics pruning null-aware

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

@@ -277,13 +277,13 @@ struct page_stats_caster : public stats_caster_base {
 
   /**
    * @brief Computes host side data including page row offsets, column chunk page offsets, and host
-   * columns containing page-level min, max and (optional) is_null statistics for a column
+   * columns containing page-level min, max and (optional) all-null statistics for a column
    *
    * @param schema_idx Column schema index
    * @param dtype Column data type
    * @param stream CUDA stream
    * @return A tuple of page row offsets, column chunk page offsets, and host columns containing
-   * page-level min, max and (optional) is_null statistics
+   * page-level min, max and (optional) all-null statistics
    */
   template <typename T>
   [[nodiscard]] auto compute_host_data(cudf::size_type schema_idx,
@@ -300,11 +300,13 @@ struct page_stats_caster : public stats_caster_base {
 
     auto const total_pages = col_chunk_page_offsets.back();
 
-    // Create host columns with page-level min, max and optionally is_null statistics
+    // Create host columns with page-level min, max and optionally all-null statistics. The
+    // all-null column is true only when every value in the page is null, false when none are, and
+    // null when only some are, which is what lets it answer both IS_NULL and IS NOT NULL.
     host_column<T> min(total_pages, stream);
     host_column<T> max(total_pages, stream);
-    std::optional<host_column<bool>> is_null;
-    if (has_is_null_operator) { is_null = host_column<bool>(total_pages, stream); }
+    std::optional<host_column<bool>> all_null;
+    if (has_is_null_operator) { all_null = host_column<bool>(total_pages, stream); }
 
     // Compute timestamp scale factor for precision conversion
     auto const ts_scale = [&] {
@@ -353,22 +355,24 @@ struct page_stats_caster : public stats_caster_base {
               if (has_is_null_operator) {
                 // Check if the page is completely null
                 if (column_index.null_pages[page_idx]) {
-                  is_null->val[column_page_idx] = true;
+                  all_null->val[column_page_idx] = true;
                   return;
                 }
                 // Check if the page doesn't have a null count
                 if (not column_index.null_counts.has_value()) {
-                  is_null->set_index(column_page_idx, std::nullopt, {});
+                  all_null->set_index(column_page_idx, std::nullopt, {});
                   return;
                 }
                 // Use the null count to determine if the page is completely null
                 auto const page_row_count =
                   page_row_offsets[column_page_idx + 1] - page_row_offsets[column_page_idx];
                 auto const& null_count = column_index.null_counts.value()[page_idx];
-                if (null_count == page_row_count) {
-                  is_null->val[column_page_idx] = false;
-                } else if (null_count > 0 and null_count < page_row_count) {
-                  is_null->set_index(column_page_idx, std::nullopt, {});
+                if (null_count == 0) {
+                  all_null->val[column_page_idx] = false;
+                } else if (null_count < page_row_count) {
+                  all_null->set_index(column_page_idx, std::nullopt, {});
+                } else if (null_count == page_row_count) {
+                  all_null->val[column_page_idx] = true;
                 } else {
                   CUDF_FAIL("Invalid null count");
                 }
@@ -381,7 +385,7 @@ struct page_stats_caster : public stats_caster_base {
                       std::move(col_chunk_page_offsets),
                       std::move(min),
                       std::move(max),
-                      std::move(is_null)};
+                      std::move(all_null)};
   }
 
   /**

cpp/src/io/parquet/row_group_stats_helpers.hpp

@@ -138,6 +138,12 @@ struct row_group_stats_caster : public stats_caster_base {
                 } else {
                   CUDF_FAIL("Invalid null count");
                 }
+              } else {
+                // Statistics without a null count say nothing about this chunk's nullability. The
+                // value array is allocated uninitialized and the null mask starts out all valid, so
+                // this entry has to be marked null; leaving it alone would let an uninitialized
+                // byte be read as an answer.
+                is_null->set_index(stats_idx, std::nullopt, {});
               }
             }
           } else {

cpp/src/io/parquet/stats_filter_helpers.cpp

@@ -14,6 +14,29 @@
 
 namespace cudf::io::parquet::detail {
 
+namespace {
+
+/**
+ * @brief Maps a logical connective to its null-aware equivalent, returning any other operator as is
+ *
+ * A null in a statistics column says the writer did not record the statistic, never that the data
+ * is null, so the statistics expression is a three-valued predicate in which null means "unknown,
+ * keep this chunk". Three-valued logic is what propagates that: `false AND unknown` is false,
+ * because a chunk holding no row that can satisfy one conjunct cannot satisfy the conjunction
+ * whatever the other side turns out to be. The plain connectives instead return null whenever
+ * either side is null, which lets one absent statistic switch off pruning for the whole expression.
+ */
+[[nodiscard]] ast::ast_operator null_aware_operator(ast::ast_operator op)
+{
+  switch (op) {
+    case ast::ast_operator::LOGICAL_AND: return ast::ast_operator::NULL_LOGICAL_AND;
+    case ast::ast_operator::LOGICAL_OR: return ast::ast_operator::NULL_LOGICAL_OR;
+    default: return op;
+  }
+}
+
+}  // namespace
+
 stats_columns_collector::stats_columns_collector(ast::expression const& expr,
                                                  cudf::size_type num_columns)
   : _num_columns(num_columns)
@@ -76,6 +99,9 @@ std::reference_wrapper<ast::expression const> stats_columns_collector::visit(
         op == ast_operator::LESS_EQUAL or op == ast_operator::GREATER or
         op == ast_operator::GREATER_EQUAL) {
       _columns_mask[col_ref->get_column_index()] = true;
+      // None of these can match a null, so their stats expressions consult the nullability column
+      // to rule out a chunk of nothing but nulls, which has no min or max to compare against.
+      _has_is_null_operator = true;
     }
   } else {
     // Visit the operands and ignore any output as we only want to build the column mask
@@ -101,6 +127,28 @@ stats_expression_converter::stats_expression_converter(ast::expression const& ex
   expr.accept(*this);
 }
 
+void stats_expression_converter::push_non_null_guard(size_type col_index,
+                                                     ast::expression const& stats_expr)
+{
+  using cudf::ast::ast_operator;
+
+  if (not std::cmp_equal(_stats_cols_per_column, 3)) { return; }
+
+  auto const& all_null =
+    _stats_expr.push(ast::column_reference{col_index * _stats_cols_per_column + 2});
+  // Answering "not entirely null" takes all three of the column's states, so a plain NOT will not
+  // do: its null state says the chunk holds both nulls and values, or that the writer recorded no
+  // null count, and both of those answer this question true. NOT alone answers it null and hands an
+  // unknown to a comparison that is in fact decisive.
+  auto const& not_all_null = _stats_expr.push(
+    ast::operation{ast_operator::NULL_LOGICAL_OR,
+                   _stats_expr.push(ast::operation{ast_operator::IS_NULL, all_null}),
+                   _stats_expr.push(ast::operation{ast_operator::NOT, all_null})});
+  // Null-aware so that the false this side pushes for an all-null chunk prunes it even though the
+  // min and max it lacks leave `stats_expr` unknown.
+  _stats_expr.push(ast::operation{ast_operator::NULL_LOGICAL_AND, not_all_null, stats_expr});
+}
+
 std::reference_wrapper<ast::expression const> stats_expression_converter::visit(
   ast::operation const& expr)
 {
@@ -203,35 +251,47 @@ std::reference_wrapper<ast::expression const> stats_expression_converter::visit(
           _stats_expr.push(ast::column_reference{col_index * _stats_cols_per_column});
         auto const& vmax =
           _stats_expr.push(ast::column_reference{col_index * _stats_cols_per_column + 1});
-        _stats_expr.push(ast::operation{
-          ast::ast_operator::LOGICAL_AND,
+        // The two halves are separately optional in the statistics, so they are combined null-aware
+        // to keep whichever one is present decisive.
+        auto const& in_range = _stats_expr.push(ast::operation{
+          ast::ast_operator::NULL_LOGICAL_AND,
           _stats_expr.push(ast::operation{ast_operator::GREATER_EQUAL, vmax, literal}),
           _stats_expr.push(ast::operation{ast_operator::LESS_EQUAL, vmin, literal})});
+        // An all-null chunk has no min or max, so this range test is unknown there and would keep
+        // the chunk. The guard makes it prune instead.
+        push_non_null_guard(col_index, in_range);
         break;
       }
       case ast_operator::NOT_EQUAL: {
         auto const& vmin =
           _stats_expr.push(ast::column_reference{col_index * _stats_cols_per_column});
         auto const& vmax =
           _stats_expr.push(ast::column_reference{col_index * _stats_cols_per_column + 1});
-        _stats_expr.push(
-          ast::operation{ast_operator::LOGICAL_OR,
+        // Null-aware for the same reason as the range test above: either half can be the one the
+        // statistics carry.
+        auto const& outside_range = _stats_expr.push(
+          ast::operation{ast_operator::NULL_LOGICAL_OR,
                          _stats_expr.push(ast::operation{ast_operator::NOT_EQUAL, vmin, vmax}),
                          _stats_expr.push(ast::operation{ast_operator::NOT_EQUAL, vmax, literal})});
+        // A null does not satisfy `!=` either, and an all-null chunk has no min or max to make this
+        // test decisive, so the guard prunes it.
+        push_non_null_guard(col_index, outside_range);
         break;
       }
       case ast_operator::LESS: [[fallthrough]];
       case ast_operator::LESS_EQUAL: {
         auto const& vmin =
           _stats_expr.push(ast::column_reference{col_index * _stats_cols_per_column});
-        _stats_expr.push(ast::operation{op, vmin, literal});
+        // An all-null chunk has no min, leaving this test unknown, so the guard prunes it.
+        push_non_null_guard(col_index, _stats_expr.push(ast::operation{op, vmin, literal}));
         break;
       }
       case ast_operator::GREATER: [[fallthrough]];
       case ast_operator::GREATER_EQUAL: {
         auto const& vmax =
           _stats_expr.push(ast::column_reference{col_index * _stats_cols_per_column + 1});
-        _stats_expr.push(ast::operation{op, vmax, literal});
+        // An all-null chunk has no max, leaving this test unknown, so the guard prunes it.
+        push_non_null_guard(col_index, _stats_expr.push(ast::operation{op, vmax, literal}));
         break;
       }
       default: {
@@ -242,7 +302,8 @@ std::reference_wrapper<ast::expression const> stats_expression_converter::visit(
   }  // Visit operands and push expression for `expr op expr` form
   else if (lhs_kind == operand_kind::EXPRESSION and rhs_kind == operand_kind::EXPRESSION) {
     auto new_operands = visit_operands(expr.get_operands());
-    _stats_expr.push(ast::operation{op, new_operands.front(), new_operands.back()});
+    _stats_expr.push(
+      ast::operation{null_aware_operator(op), new_operands.front(), new_operands.back()});
   }  // Push _always_true for `col op col`, `expr op col`, `expr op lit` forms
   else {
     _stats_expr.push(ast::operation{ast_operator::IDENTITY, *_always_true});

cpp/src/io/parquet/stats_filter_helpers.hpp

@@ -333,9 +333,13 @@ class stats_columns_collector : public ast::detail::expression_transformer {
 
   /**
    * @brief Return a boolean vector indicating input columns that can participate in stats based
-   * filtering
+   * filtering, and whether the stats table needs a per-column nullability column
    *
-   * @return Boolean vector indicating input columns that can participate in stats based filtering
+   * The nullability column is needed by an `IS_NULL` operator, which is answered from it alone, and
+   * by any comparison against a literal, which uses it to rule out a chunk of nothing but nulls.
+   *
+   * @return Boolean vector indicating input columns that can participate in stats based filtering,
+   * and whether the nullability column is needed
    */
   std::pair<thrust::host_vector<bool>, bool> get_stats_columns_mask() &&;
 
@@ -383,6 +387,24 @@ class stats_expression_converter : public stats_columns_collector {
   thrust::host_vector<bool> get_stats_columns_mask() && = delete;
 
  private:
+  /**
+   * @brief Push `not_all_null AND stats_expr` for a column, so that a chunk holding nothing but
+   * nulls fails a predicate that needs a non-null value to match
+   *
+   * A writer has no non-null value to compute min and max from for such a chunk, so it omits them
+   * and every min/max comparison evaluates to null, which keeps the chunk. The nullability
+   * statistic is decisive where min and max are absent: despite being built as `is_null`, it is
+   * true only when *every* value in the chunk is null, false when none are, and null when only some
+   * are or when the writer recorded no null count. Reading three states out of that column takes
+   * more than a `NOT`, since the null state answers "not entirely null" with a definite yes.
+   *
+   * Does nothing when the nullability column was not built, leaving `stats_expr` as the result.
+   *
+   * @param col_index Index of the column in the input table
+   * @param stats_expr Statistics expression to guard, already pushed onto the tree
+   */
+  void push_non_null_guard(size_type col_index, ast::expression const& stats_expr);
+
   ast::tree _stats_expr;
   cudf::size_type _stats_cols_per_column;
   std::unique_ptr<cudf::numeric_scalar<bool>> _always_true_scalar;