Expose Parquet column-chunk statistics in pylibcudf

cpp/include/cudf/io/parquet_metadata.hpp

@@ -10,12 +10,18 @@
 
 #pragma once
 
+#include <cudf/column/column.hpp>
 #include <cudf/io/datasource.hpp>
 #include <cudf/io/parquet_schema.hpp>
 #include <cudf/io/types.hpp>
+#include <cudf/table/table.hpp>
+#include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/export.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 
+#include <memory>
 #include <span>
+#include <string>
 #include <string_view>
 #include <vector>
 
@@ -294,6 +300,49 @@ parquet_metadata read_parquet_metadata(source_info const& src_info);
 std::vector<parquet::FileMetaData> read_parquet_footers(
   std::span<std::unique_ptr<cudf::io::datasource> const> sources);
 
+/**
+ * @brief Min/max bounds decoded from parquet column-chunk statistics.
+ *
+ * ``file_indices`` and ``row_group_indices`` identify the column chunks represented by each row of
+ * every table in ``bounds``. Each table in ``bounds`` corresponds positionally to a requested
+ * column and contains exactly two columns: decoded minimum values followed by decoded maximum
+ * values.
+ */
+struct column_chunk_bounds_result {
+  /// File index for each row in every bounds table
+  std::unique_ptr<column> file_indices;
+  /// File-local row-group index for each row in every bounds table
+  std::unique_ptr<column> row_group_indices;
+  /// One two-column table per requested column, where column 0 is min and column 1 is max
+  std::vector<std::unique_ptr<table>> bounds;
+};
+
+/**
+ * @brief Decode parquet column-chunk min/max statistics for selected leaf columns.
+ *
+ * Missing min/max statistics are represented as nulls in the corresponding output column. Parquet
+ * min/max exactness flags are not interpreted by this function. The requested column names are
+ * resolved against each file's schema.
+ *
+ * @ingroup io_readers
+ *
+ * @param parquet_metadatas Parquet file metadata, one per source
+ * @param column_names Dotted leaf-column paths to decode statistics for
+ * @param stream CUDA stream used for device memory operations
+ * @param mr Device memory resource to use for device memory allocation
+ * @return Decoded min/max bounds and row-group identifiers
+ *
+ * @throw std::invalid_argument If a requested leaf-column path is missing or ambiguous.
+ * @throw std::invalid_argument If a requested column has unsupported or compound statistics dtype.
+ * @throw std::invalid_argument If a requested column has mismatching statistics dtype across
+ * sources.
+ */
+column_chunk_bounds_result column_chunk_bounds(
+  std::vector<parquet::FileMetaData> parquet_metadatas,
+  std::span<std::string const> column_names,
+  rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
+
 /** @} */  // end of group
 }  // namespace io
 }  // namespace CUDF_EXPORT cudf

cpp/src/io/parquet/predicate_pushdown.cpp

@@ -5,17 +5,16 @@
 
 #include "expression_transform_helpers.hpp"
 #include "reader_impl_helpers.hpp"
+#include "row_group_stats_helpers.hpp"
 #include "stats_filter_helpers.hpp"
-#include "timestamp_utils.cuh"
 
 #include <cudf/column/column_factories.hpp>
 #include <cudf/detail/iterator.cuh>
 #include <cudf/detail/transform.hpp>
-#include <cudf/detail/utilities/vector_factories.hpp>
+#include <cudf/table/table.hpp>
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 #include <cudf/utilities/span.hpp>
-#include <cudf/utilities/traits.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
 #include <cuda/iterator>
@@ -28,104 +27,6 @@
 
 namespace cudf::io::parquet::detail {
 
-namespace {
-
-/**
- * @brief Converts column chunk statistics to 2 device columns - min, max values.
- *
- * Each column's number of rows equals the total number of row groups.
- *
- */
-struct row_group_stats_caster : public stats_caster_base {
-  size_type total_row_groups;
-  std::vector<metadata> const& per_file_metadata;
-  host_span<std::vector<size_type> const> row_group_indices;
-  bool has_is_null_operator;
-
-  // Creates device columns from column statistics (min, max)
-  template <typename T>
-  std::
-    tuple<std::unique_ptr<column>, std::unique_ptr<column>, std::optional<std::unique_ptr<column>>>
-    operator()(host_span<int const> per_source_schema_indices,
-               cudf::data_type dtype,
-               cuda::stream_ref stream,
-               rmm::device_async_resource_ref mr) const
-  {
-    // List, Struct, Dictionary types are not supported
-    if constexpr (cudf::is_compound<T>() && !std::is_same_v<T, string_view>) {
-      CUDF_FAIL("Compound types do not have statistics");
-    } else {
-      host_column<T> min(total_row_groups, stream);
-      host_column<T> max(total_row_groups, stream);
-      std::optional<host_column<bool>> is_null;
-      if (has_is_null_operator) { is_null = host_column<bool>(total_row_groups, stream); }
-
-      size_type stats_idx = 0;
-      for (size_t src_idx = 0; src_idx < row_group_indices.size(); ++src_idx) {
-        auto const mapped_schema_idx = per_source_schema_indices[src_idx];
-        // Compute timestamp scale factor for precision conversion from the mapped source schema.
-        auto const ts_scale = [&] {
-          if constexpr (cudf::is_timestamp<T>()) {
-            auto const& schema = per_file_metadata[src_idx].schema[mapped_schema_idx];
-            return calc_timestamp_scale(schema.logical_type, static_cast<int32_t>(T::period::den));
-          }
-          return 0;
-        }();
-
-        for (auto const rg_idx : row_group_indices[src_idx]) {
-          auto const& row_group = per_file_metadata[src_idx].row_groups[rg_idx];
-          auto col              = std::find_if(row_group.columns.begin(),
-                                  row_group.columns.end(),
-                                  [mapped_schema_idx](ColumnChunk const& col) {
-                                    return col.schema_idx == mapped_schema_idx;
-                                  });
-          if (col != std::end(row_group.columns)) {
-            auto const& colchunk = *col;
-            // To support deprecated min, max fields.
-            auto const& min_value = colchunk.meta_data.statistics.min_value.has_value()
-                                      ? colchunk.meta_data.statistics.min_value
-                                      : colchunk.meta_data.statistics.min;
-            auto const& max_value = colchunk.meta_data.statistics.max_value.has_value()
-                                      ? colchunk.meta_data.statistics.max_value
-                                      : colchunk.meta_data.statistics.max;
-            // translate binary data to Type then to <T>
-            min.set_index(stats_idx, min_value, colchunk.meta_data.type, ts_scale);
-            max.set_index(stats_idx, max_value, colchunk.meta_data.type, ts_scale);
-            // Check the nullability of this column chunk
-            if (has_is_null_operator) {
-              if (colchunk.meta_data.statistics.null_count.has_value()) {
-                auto const& null_count = colchunk.meta_data.statistics.null_count.value();
-                if (null_count == 0) {
-                  is_null->val[stats_idx] = false;
-                } else if (null_count < colchunk.meta_data.num_values) {
-                  is_null->set_index(stats_idx, std::nullopt, {});
-                } else if (null_count == colchunk.meta_data.num_values) {
-                  is_null->val[stats_idx] = true;
-                } else {
-                  CUDF_FAIL("Invalid null count");
-                }
-              }
-            }
-          } else {
-            // Marking it null, if column present in row group
-            min.set_index(stats_idx, std::nullopt, {});
-            max.set_index(stats_idx, std::nullopt, {});
-            if (has_is_null_operator) { is_null->set_index(stats_idx, std::nullopt, {}); }
-          }
-          stats_idx++;
-        }
-      };
-      return {min.to_device(dtype, stream, mr),
-              max.to_device(dtype, stream, mr),
-              has_is_null_operator ? std::make_optional(is_null->to_device(
-                                       data_type{cudf::type_id::BOOL8}, stream, mr))
-                                   : std::nullopt};
-    }
-  }
-};
-
-}  // namespace
-
 bool aggregate_reader_metadata::any_row_group_stats_available(
   host_span<std::vector<size_type> const> input_row_group_indices,
   host_span<int const> filter_column_schemas) const
@@ -288,7 +189,8 @@ aggregate_reader_metadata::filter_row_groups(
   // Span of row groups to apply bloom filtering on.
   auto const bloom_filter_input_row_groups =
     stats_filtered_row_groups.has_value()
-      ? host_span<std::vector<size_type> const>(stats_filtered_row_groups.value())
+      ? host_span<std::vector<size_type> const>{stats_filtered_row_groups.value().data(),
+                                                stats_filtered_row_groups.value().size()}
       : input_row_group_indices;
 
   // Collect equality literals for each input table column for bloom filtering