Forward-merge main into pandas3

cpp/benchmarks/CMakeLists.txt

@@ -285,6 +285,7 @@ ConfigureNVBench(BITMASK_NVBENCH bitmask/bitmask_and.cpp bitmask/set_null_mask.c
 # * parquet writer benchmark ----------------------------------------------------------------------
 ConfigureNVBench(
   PARQUET_WRITER_NVBENCH io/parquet/parquet_writer.cpp io/parquet/parquet_writer_chunks.cpp
+  io/parquet/parquet_writer_dict.cpp
 )
 
 # ##################################################################################################

cpp/benchmarks/io/parquet/parquet_writer_dict.cpp

@@ -0,0 +1,239 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "io/parquet/compact_protocol_reader.hpp"
+
+#include <benchmarks/common/memory_stats.hpp>
+#include <benchmarks/io/cuio_common.hpp>
+
+#include <cudf_test/column_wrapper.hpp>
+
+#include <cudf/column/column.hpp>
+#include <cudf/io/datasource.hpp>
+#include <cudf/io/experimental/hybrid_scan.hpp>
+#include <cudf/io/parquet.hpp>
+#include <cudf/io/parquet_io_utils.hpp>
+#include <cudf/io/parquet_schema.hpp>
+#include <cudf/table/table.hpp>
+#include <cudf/types.hpp>
+#include <cudf/utilities/default_stream.hpp>
+#include <cudf/utilities/error.hpp>
+
+#include <nvbench/nvbench.cuh>
+
+#include <algorithm>
+#include <cmath>
+#include <numeric>
+#include <random>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace {
+
+constexpr auto frequent_pages_ratio =
+  0.8;  ///< 80% of the pages will only contain elements from the frequent set
+
+/**
+ * @brief Build a string column such that certain pages only contain elements from the frequent set
+ and others only contain elements from the rare set
+ *
+ * Each generated value is of the form `"k_<v>"`, where `v` is a `cudf::size_type` index in `[0,
+ cardinality)`, so cardinality maps 1:1 to distinct strings.
+ *
+ * @param num_rows Total number of rows
+ * @param page_size_rows Number of rows per page
+ * @param cardinality Total number of distinct values
+ * @param frequent_set_ratio Fraction of `cardinality` assigned to the frequent set
+ * @return Constructed column values
+ */
+std::vector<std::string> build_string_column(cudf::size_type num_rows,
+                                             cudf::size_type page_size_rows,
+                                             cudf::size_type cardinality,
+                                             double frequent_set_ratio)
+{
+  static constexpr auto dict_rng_seed = 0xC0DEFACE;
+
+  CUDF_EXPECTS(frequent_set_ratio > 0.0 and frequent_set_ratio < 1.0,
+               "frequent_set_ratio must be between 0.0 and 1.0");
+  CUDF_EXPECTS(num_rows % page_size_rows == 0, "num_rows must be a multiple of page_size_rows");
+  static_assert(frequent_pages_ratio > 0.0 and frequent_pages_ratio < 1.0,
+                "frequent_pages_ratio must be between 0.0 and 1.0");
+
+  auto const total_pages = num_rows / page_size_rows;
+  auto const frequent_set_threshold =
+    static_cast<cudf::size_type>(total_pages * frequent_pages_ratio) * page_size_rows;
+
+  auto const frequent_set_size =
+    static_cast<cudf::size_type>(static_cast<double>(cardinality) * frequent_set_ratio);
+
+  std::mt19937 rng{dict_rng_seed};
+  std::uniform_int_distribution<cudf::size_type> freq_dist(0, frequent_set_size - 1);
+  std::uniform_int_distribution<cudf::size_type> rare_dist(frequent_set_size, cardinality - 1);
+
+  cudf::size_type row_idx = 0;
+  std::vector<std::string> values(num_rows);
+  std::generate_n(values.begin(), num_rows, [&]() {
+    auto const v = row_idx++ < frequent_set_threshold ? freq_dist(rng) : rare_dist(rng);
+    return "k_" + std::to_string(v);
+  });
+
+  return values;
+}
+
+/**
+ * @brief Build a table with a single STRING column
+ *
+ * @tparam reverse_order Whether to reverse the order of the values
+ * @param num_rows Number of rows
+ * @param page_size_rows Number of rows per page
+ * @param cardinality Total number of distinct values
+ * @param frequent_set_ratio Fraction of `cardinality` assigned to the frequent set
+ * @return std::unique_ptr<cudf::table>
+ */
+template <bool reverse_order = false>
+[[nodiscard]] std::unique_ptr<cudf::table> build_table(cudf::size_type num_rows,
+                                                       cudf::size_type page_size_rows,
+                                                       cudf::size_type cardinality,
+                                                       double frequent_set_ratio)
+{
+  constexpr cudf::size_type num_cols = 1;
+
+  auto values = build_string_column(num_rows, page_size_rows, cardinality, frequent_set_ratio);
+  if constexpr (reverse_order) { std::reverse(values.begin(), values.end()); }
+  std::vector<std::unique_ptr<cudf::column>> cols;
+  cols.reserve(num_cols);
+  cols.emplace_back(cudf::test::strings_column_wrapper(values.begin(), values.end()).release());
+  return std::make_unique<cudf::table>(std::move(cols));
+}
+
+/**
+ * @brief Compute per-page RLE bit widths for dictionary-encoded pages from the parquet page index
+ *
+ * Assumption: All parquet pages are dictionary-encoded, no nulls, no rep/def levels
+ *
+ * @param source Datasource
+ * @param footer File metadata
+ * @return Vector of number of bits per page for dictionary-encoded pages
+ */
+[[nodiscard]] std::vector<int> compute_page_dict_bits(cudf::io::datasource& source,
+                                                      cudf::io::parquet::FileMetaData const& footer)
+{
+  using namespace cudf::io::parquet;
+
+  std::vector<int> bits;
+
+  for (auto const& rg : footer.row_groups) {
+    for (auto const& chunk : rg.columns) {
+      if (not chunk.offset_index.has_value()) { continue; }
+      for (auto const& page_loc : chunk.offset_index->page_locations) {
+        if (page_loc.offset <= 0 or page_loc.compressed_page_size <= 0) { continue; }
+        auto const buffer = source.host_read(page_loc.offset, page_loc.compressed_page_size);
+        detail::CompactProtocolReader cp(buffer->data(), buffer->size());
+        PageHeader page_header;
+        cp.read(&page_header);
+        // Check if the page is dictionary-encoded.
+        auto const is_dict_encoded =
+          (page_header.type == PageType::DATA_PAGE and
+           (page_header.data_page_header.encoding == Encoding::PLAIN_DICTIONARY or
+            page_header.data_page_header.encoding == Encoding::RLE_DICTIONARY)) or
+          (page_header.type == PageType::DATA_PAGE_V2 and
+           (page_header.data_page_header_v2.encoding == Encoding::PLAIN_DICTIONARY or
+            page_header.data_page_header_v2.encoding == Encoding::RLE_DICTIONARY));
+        if (not is_dict_encoded) { continue; }
+        // `cp` is positioned at the first byte of the page payload after the
+        // header thrift; that byte is the RLE bit width for dict-indexed
+        // pages (valid only with no rep/def levels).
+        bits.push_back(cp.getb());
+      }
+    }
+  }
+  return bits;
+}
+
+}  // namespace
+
+void BM_parq_write_dict_encoding(nvbench::state& state)
+{
+  auto const num_rows           = static_cast<cudf::size_type>(state.get_int64("num_rows"));
+  auto const reverse_order      = static_cast<bool>(state.get_int64("reverse_order"));
+  auto const cardinality        = static_cast<cudf::size_type>(state.get_int64("cardinality"));
+  auto const frequent_set_ratio = static_cast<double>(state.get_float64("freq_set_ratio"));
+  auto const page_size_rows     = static_cast<cudf::size_type>(state.get_int64("page_size_rows"));
+
+  CUDF_EXPECTS(page_size_rows <= num_rows and num_rows % page_size_rows == 0,
+               "num_rows must be a multiple of page_size_rows");
+
+  cuio_source_sink_pair source_sink(io_type::FILEPATH);
+
+  auto const table = [&]() {
+    if (reverse_order) {
+      return build_table<true>(num_rows, page_size_rows, cardinality, frequent_set_ratio);
+    } else {
+      return build_table<false>(num_rows, page_size_rows, cardinality, frequent_set_ratio);
+    }
+  }();
+
+  auto const mem_stats_logger = cudf::memory_stats_logger();
+  state.set_cuda_stream(nvbench::make_cuda_stream_view(cudf::get_default_stream().value()));
+  state.exec(
+    nvbench::exec_tag::timer | nvbench::exec_tag::sync, [&](nvbench::launch&, auto& timer) {
+      timer.start();
+      auto const write_opts =
+        cudf::io::parquet_writer_options::builder(source_sink.make_sink_info(), table->view())
+          .compression(cudf::io::compression_type::NONE)
+          .dictionary_policy(cudf::io::dictionary_policy::ALWAYS)
+          .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN)
+          .row_group_size_rows(num_rows)
+          .max_page_size_rows(page_size_rows)
+          .max_page_size_bytes(std::size_t{64} << 20)
+          .build();
+      cudf::io::write_parquet(write_opts);
+      timer.stop();
+    });
+
+  state.add_element_count(static_cast<double>(table->num_rows()), "rows");
+  state.add_buffer_size(
+    mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage");
+  state.add_buffer_size(source_sink.size(), "encoded_file_size", "encoded_file_size");
+
+  // Use hybrid scan reader to get footer with page index.
+  {
+    auto const datasource =
+      std::move(cudf::io::make_datasources(source_sink.make_source_info()).front());
+    auto const datasource_ref = std::ref(*datasource);
+    auto const footer_buf     = cudf::io::parquet::fetch_footer_to_host(datasource_ref);
+    cudf::io::parquet::experimental::hybrid_scan_reader reader(*footer_buf,
+                                                               cudf::io::parquet_reader_options{});
+    auto const page_index_bytes = reader.page_index_byte_range();
+    CUDF_EXPECTS(not page_index_bytes.is_empty(), "Page index is required");
+    auto const page_index_buffer =
+      cudf::io::parquet::fetch_page_index_to_host(datasource_ref, page_index_bytes);
+    reader.setup_page_index(*page_index_buffer);
+
+    auto const metadata       = reader.parquet_metadata();
+    auto const page_dict_bits = compute_page_dict_bits(datasource_ref, metadata);
+
+    CUDF_EXPECTS(not page_dict_bits.empty(), "No dictionary-encoded pages found");
+
+    auto const [min_it, max_it] = std::minmax_element(page_dict_bits.begin(), page_dict_bits.end());
+    auto const sum =
+      std::accumulate(page_dict_bits.begin(), page_dict_bits.end(), std::uint64_t{0});
+    auto const mean =
+      std::round(static_cast<double>(sum) / static_cast<double>(page_dict_bits.size()));
+    state.add_element_count(static_cast<double>(*min_it), "dict_rle_bits_min");
+    state.add_element_count(static_cast<double>(*max_it), "dict_rle_bits_max");
+    state.add_element_count(mean, "dict_rle_bits_mean");
+  }
+}
+
+NVBENCH_BENCH(BM_parq_write_dict_encoding)
+  .set_name("parquet_write_dict_encoding")
+  .set_min_samples(4)
+  .add_int64_axis("reverse_order", {false, true})
+  .add_int64_axis("num_rows", {1'000'000})
+  .add_int64_axis("page_size_rows", {10'000, 100'000})
+  .add_int64_axis("cardinality", {64'000, 100'000})
+  .add_float64_axis("freq_set_ratio", {0.001, 0.01});

cpp/examples/hybrid_scan_io/common_utils.cpp

@@ -6,8 +6,7 @@
 #include "common_utils.hpp"
 
 #include <cudf/ast/expressions.hpp>
-#include <cudf/io/parquet.hpp>
-#include <cudf/join/filtered_join.hpp>
+#include <cudf/table/equality.hpp>
 #include <cudf/table/table_view.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -69,30 +68,10 @@ void check_tables_equal(cudf::table_view const& lhs_table,
                         cudf::table_view const& rhs_table,
                         rmm::cuda_stream_view stream)
 {
-  try {
-    // Left anti-join the original and transcoded tables identical tables should not throw an
-    // exception and return an empty indices vector
-    cudf::filtered_join join_obj(lhs_table, cudf::null_equality::EQUAL, stream);
-    auto const indices = join_obj.anti_join(rhs_table, stream);
-    // No exception thrown, check indices
-    auto const tables_equal = indices->size() == 0;
-    if (tables_equal) {
-      std::cout << "Tables identical: " << std::boolalpha << tables_equal << "\n\n";
-    } else {
-      // Helper to write parquet data for inspection
-      auto const write_parquet =
-        [](cudf::table_view table, std::string filepath, rmm::cuda_stream_view stream) {
-          auto sink_info = cudf::io::sink_info(filepath);
-          auto opts      = cudf::io::parquet_writer_options::builder(sink_info, table).build();
-          cudf::io::write_parquet(opts, stream);
-        };
-      write_parquet(lhs_table, "lhs_table.parquet", stream);
-      write_parquet(rhs_table, "rhs_table.parquet", stream);
-      throw std::logic_error("Tables identical: false\n\n");
-    }
-  } catch (std::exception& e) {
-    std::cout << e.what() << std::endl;
-  }
+  auto const tables_equal =
+    cudf::tables_equal(lhs_table, rhs_table, cudf::null_equality::EQUAL, stream);
+  std::cout << "Tables identical: " << std::boolalpha << tables_equal << "\n\n";
+  if (not tables_equal) { throw std::logic_error("Table equality check failed"); }
 }
 
 std::vector<io_source> extract_input_sources(std::string const& paths,

cpp/examples/parquet_io/common_utils.cpp

@@ -7,7 +7,7 @@
 
 #include <cudf/concatenate.hpp>
 #include <cudf/io/types.hpp>
-#include <cudf/join/filtered_join.hpp>
+#include <cudf/table/equality.hpp>
 #include <cudf/table/table_view.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -87,19 +87,10 @@ void check_tables_equal(cudf::table_view const& lhs_table,
                         cudf::table_view const& rhs_table,
                         rmm::cuda_stream_view stream)
 {
-  try {
-    // Left anti-join the original and transcoded tables identical tables should not throw an
-    // exception and return an empty indices vector
-    cudf::filtered_join join_obj(lhs_table, cudf::null_equality::EQUAL, stream);
-    auto const indices = join_obj.anti_join(rhs_table, stream);
-
-    // No exception thrown, check indices
-    auto const valid = indices->size() == 0;
-    std::cout << "Tables identical: " << valid << "\n\n";
-  } catch (std::exception& e) {
-    std::cerr << e.what() << std::endl << std::endl;
-    throw std::runtime_error("Tables identical: false\n\n");
-  }
+  auto const tables_equal =
+    cudf::tables_equal(lhs_table, rhs_table, cudf::null_equality::EQUAL, stream);
+  std::cout << "Tables identical: " << std::boolalpha << tables_equal << "\n\n";
+  if (not tables_equal) { throw std::logic_error("Table equality check failed"); }
 }
 
 std::unique_ptr<cudf::table> concatenate_tables(std::vector<std::unique_ptr<cudf::table>> tables,

cpp/examples/parquet_io/parquet_io.cpp

@@ -4,7 +4,6 @@
  */
 
 #include "common_utils.hpp"
-#include "io_source.hpp"
 #include "timer.hpp"
 
 #include <cudf/io/parquet.hpp>