From 2a5d2f70af4edb0343fbc4d375b505407bc41398 Mon Sep 17 00:00:00 2001 From: Matthew Murray <41342305+Matt711@users.noreply.github.com> Date: Fri, 15 May 2026 15:41:00 -0400 Subject: [PATCH 1/5] Guarantee `insert_finished()` on bulk AllGather (#22516) Need to make sure `insert_finished` is called for collective operations. Authors: - Matthew Murray (https://github.com/Matt711) Approvers: - Matthew Roeschke (https://github.com/mroeschke) - Mads R. B. Kristensen (https://github.com/madsbk) - Lawrence Mitchell (https://github.com/wence-) URL: https://github.com/rapidsai/cudf/pull/22516 --- .../experimental/rapidsmpf/frontend/core.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py index 7d362270ffdc..ef7fbf244758 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py @@ -560,8 +560,14 @@ def all_gather_host_data( br=br, statistics=Statistics(enable=False), ) - allgather.insert(0, PackedData.from_host_bytes(data, br)) - allgather.insert_finished() + # TODO: Make AllGather (bulk) a context manager so this becomes + # with AllGather(...) as ag: + # ag.insert(0, PackedData.from_host_bytes(data, br)) + # results = ag.wait_and_extract(ordered=True) + try: + allgather.insert(0, PackedData.from_host_bytes(data, br)) + finally: + allgather.insert_finished() results = allgather.wait_and_extract(ordered=True) return [r.to_host_bytes() for r in results] From fc0598ff4016a9c106da65737d3852fbcb75f3b5 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Fri, 15 May 2026 13:53:59 -0700 Subject: [PATCH 2/5] Use `cudf::tables_equal` API in libcudf examples (#22428) Follow up #22319 Use `cudf::tables_equal` API in libcudf examples instead of anti join to check if two tables are equal Authors: - Muhammad Haseeb (https://github.com/mhaseeb123) - Lawrence Mitchell (https://github.com/wence-) - Vukasin Milovanovic (https://github.com/vuule) Approvers: - Shruti Shivakumar (https://github.com/shrshi) - David Wendt (https://github.com/davidwendt) URL: https://github.com/rapidsai/cudf/pull/22428 --- cpp/examples/hybrid_scan_io/common_utils.cpp | 31 ++++---------------- cpp/examples/parquet_io/common_utils.cpp | 19 ++++-------- cpp/examples/parquet_io/parquet_io.cpp | 1 - 3 files changed, 10 insertions(+), 41 deletions(-) diff --git a/cpp/examples/hybrid_scan_io/common_utils.cpp b/cpp/examples/hybrid_scan_io/common_utils.cpp index 635686171866..6f08f8e1965d 100644 --- a/cpp/examples/hybrid_scan_io/common_utils.cpp +++ b/cpp/examples/hybrid_scan_io/common_utils.cpp @@ -6,8 +6,7 @@ #include "common_utils.hpp" #include -#include -#include +#include #include #include @@ -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 extract_input_sources(std::string const& paths, diff --git a/cpp/examples/parquet_io/common_utils.cpp b/cpp/examples/parquet_io/common_utils.cpp index 26ae8ea5c0f0..2adc1b1007fb 100644 --- a/cpp/examples/parquet_io/common_utils.cpp +++ b/cpp/examples/parquet_io/common_utils.cpp @@ -7,7 +7,7 @@ #include #include -#include +#include #include #include @@ -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 concatenate_tables(std::vector> tables, diff --git a/cpp/examples/parquet_io/parquet_io.cpp b/cpp/examples/parquet_io/parquet_io.cpp index a80ec5d44d54..33f78f060a03 100644 --- a/cpp/examples/parquet_io/parquet_io.cpp +++ b/cpp/examples/parquet_io/parquet_io.cpp @@ -4,7 +4,6 @@ */ #include "common_utils.hpp" -#include "io_source.hpp" #include "timer.hpp" #include From 3ae4319021539ead8d2b46e2395442d6f4585160 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Fri, 15 May 2026 15:32:19 -0700 Subject: [PATCH 3/5] Fix potential malformed headers in parquet delta decoder (#22275) This PR checks for potential malformed headers in Parquet delta decoder and propagates it in decoders. Authors: - Muhammad Haseeb (https://github.com/mhaseeb123) Approvers: - Paul Mattione (https://github.com/pmattione-nvidia) - Bradley Dice (https://github.com/bdice) - Nghia Truong (https://github.com/ttnghia) URL: https://github.com/rapidsai/cudf/pull/22275 --- cpp/src/io/parquet/delta_binary.cuh | 17 +++++++++++- cpp/src/io/parquet/page_delta_decode.cu | 35 ++++++++++++++++++------- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/cpp/src/io/parquet/delta_binary.cuh b/cpp/src/io/parquet/delta_binary.cuh index 513c65515bf7..6dac50ce80bd 100644 --- a/cpp/src/io/parquet/delta_binary.cuh +++ b/cpp/src/io/parquet/delta_binary.cuh @@ -91,6 +91,7 @@ struct delta_binary_decoder { uint32_t cur_mb; // index of the current mini-block within the block uint8_t const* cur_mb_start; // pointer to the start of the current mini-block data uint8_t const* cur_bitwidths; // pointer to the bitwidth array in the block + bool error; // flag to catch malformed headers zigzag128_t value[delta_rolling_buf_size]; // circular buffer of delta values @@ -148,7 +149,21 @@ struct delta_binary_decoder { last_value = first_value; current_value_idx = 0; - values_per_mb = block_size / mini_block_count; + error = false; + + // Validate header against the DELTA_BINARY_PACKED spec invariants + if (mini_block_count == 0 or block_size == 0 or (block_size % mini_block_count) != 0) { + error = true; + value_count = 0; + values_per_mb = 1; + block_start = d_end; + cur_mb = 0; + cur_mb_start = d_end; + cur_bitwidths = d_end; + return; + } + + values_per_mb = block_size / mini_block_count; // init the first mini-block block_start = d_start; diff --git a/cpp/src/io/parquet/page_delta_decode.cu b/cpp/src/io/parquet/page_delta_decode.cu index 0900e73cd6dd..82ac47391811 100644 --- a/cpp/src/io/parquet/page_delta_decode.cu +++ b/cpp/src/io/parquet/page_delta_decode.cu @@ -366,9 +366,11 @@ CUDF_KERNEL void __launch_bounds__(decode_delta_binary_block_size) block.sync(); auto const batch_size = db->values_per_mb; - if (batch_size > max_delta_mini_block_size) { - set_error(static_cast(decode_error::DELTA_PARAMS_UNSUPPORTED), - error_code); + if (db->error or batch_size > max_delta_mini_block_size) { + if (block.thread_rank() == 0) { + set_error(static_cast(decode_error::DELTA_PARAMS_UNSUPPORTED), + error_code); + } return; } @@ -546,6 +548,15 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) } block.sync(); + // Propagate malformed-header errors from either underlying DELTA_BINARY_PACKED decoder. + if (prefix_db->error or suffix_db->error) { + if (block.thread_rank() == 0) { + set_error(static_cast(decode_error::DELTA_PARAMS_UNSUPPORTED), + error_code); + } + return; + } + // assert that prefix and suffix have same mini-block size if (prefix_db->values_per_mb != suffix_db->values_per_mb or prefix_db->block_size != suffix_db->block_size or @@ -562,8 +573,10 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) // sanity check to make sure we can process this page auto const batch_size = prefix_db->values_per_mb; if (batch_size > max_delta_mini_block_size) { - set_error(static_cast(decode_error::DELTA_PARAMS_UNSUPPORTED), - error_code); + if (block.thread_rank() == 0) { + set_error(static_cast(decode_error::DELTA_PARAMS_UNSUPPORTED), + error_code); + } return; } @@ -759,14 +772,18 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) } block.sync(); - int const leaf_level_index = s->col.max_nesting_depth - 1; - // sanity check to make sure we can process this page auto const batch_size = db->values_per_mb; - if (batch_size > max_delta_mini_block_size) { - set_error(static_cast(decode_error::DELTA_PARAMS_UNSUPPORTED), error_code); + if (db->error or batch_size > max_delta_mini_block_size) { + if (block.thread_rank() == 0) { + set_error(static_cast(decode_error::DELTA_PARAMS_UNSUPPORTED), + error_code); + } return; } + + int const leaf_level_index = s->col.max_nesting_depth - 1; + // db->init_binary_block below resets db->values_per_mb block.sync(); // if this is a bounds page, then we need to decode up to the first mini-block From 2c52ba1ee74100492c3dea7fb5fa517fc9c15448 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Fri, 15 May 2026 15:59:16 -0700 Subject: [PATCH 4/5] Write variable bit-width keys for Parquet dictionary encoded pages (#22279) Contributes to #13995 This PR enables the Parquet writer to assign lower dictionary indices to elements appearing earlier in each column chunk thereby writing variable (reduced) number of dictionary key bits for earlier pages. New algorithm: 1. `map_insert_fn` now inserts this pair in the static map: `{row_idx, frag_idx}` instead of `{row_idx, row_idx}`. `frag_idx` is just `blockIdx.x` and indicates which page fragment actually inserted this entry in the static map (CAS race dependent but earlier thread blocks generally win this race) 2. `map_insert_fn` also writes the `PageFragment::num_dict_vals` field for each page fragment to keep a record of number of unique values inserted by it. 3. Then in `collect_map_entries_kernel` (where we assign keys to dict values), we optionally assign spatially-local keys to dictionary values inserted by the same page fragment. See algorithm details [here](https://github.com/rapidsai/cudf/pull/22279#discussion_r3156710776) 4. Finally, we launch `compute_page_dict_bits_kernel` to compute the max bit width required to encode dict keys for each parquet page. See algorithm [here](https://github.com/rapidsai/cudf/pull/22279#discussion_r3156694592) See follow up PR #22323 that deterministically writes the index of the first page fragment that sees each dictionary key Authors: - Muhammad Haseeb (https://github.com/mhaseeb123) Approvers: - Vukasin Milovanovic (https://github.com/vuule) - Paul Mattione (https://github.com/pmattione-nvidia) - Nghia Truong (https://github.com/ttnghia) URL: https://github.com/rapidsai/cudf/pull/22279 --- cpp/benchmarks/CMakeLists.txt | 1 + .../io/parquet/parquet_writer_dict.cpp | 239 ++++++++++++ cpp/src/io/parquet/chunk_dict.cu | 350 ++++++++++++++---- cpp/src/io/parquet/page_enc.cu | 11 +- cpp/src/io/parquet/parquet_gpu.cuh | 17 +- cpp/src/io/parquet/parquet_gpu.hpp | 6 +- cpp/src/io/parquet/writer_impl.cu | 10 +- cpp/tests/io/parquet_misc_test.cpp | 17 +- cpp/tests/io/parquet_writer_test.cpp | 102 +++++ 9 files changed, 665 insertions(+), 88 deletions(-) create mode 100644 cpp/benchmarks/io/parquet/parquet_writer_dict.cpp diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt index 066716e85f25..e8bc5ae97a00 100644 --- a/cpp/benchmarks/CMakeLists.txt +++ b/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 ) # ################################################################################################## diff --git a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp new file mode 100644 index 000000000000..e4a571c60b09 --- /dev/null +++ b/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 +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +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_"`, 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 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(total_pages * frequent_pages_ratio) * page_size_rows; + + auto const frequent_set_size = + static_cast(static_cast(cardinality) * frequent_set_ratio); + + std::mt19937 rng{dict_rng_seed}; + std::uniform_int_distribution freq_dist(0, frequent_set_size - 1); + std::uniform_int_distribution rare_dist(frequent_set_size, cardinality - 1); + + cudf::size_type row_idx = 0; + std::vector 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 + */ +template +[[nodiscard]] std::unique_ptr 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> cols; + cols.reserve(num_cols); + cols.emplace_back(cudf::test::strings_column_wrapper(values.begin(), values.end()).release()); + return std::make_unique(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 compute_page_dict_bits(cudf::io::datasource& source, + cudf::io::parquet::FileMetaData const& footer) +{ + using namespace cudf::io::parquet; + + std::vector 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(state.get_int64("num_rows")); + auto const reverse_order = static_cast(state.get_int64("reverse_order")); + auto const cardinality = static_cast(state.get_int64("cardinality")); + auto const frequent_set_ratio = static_cast(state.get_float64("freq_set_ratio")); + auto const page_size_rows = static_cast(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(num_rows, page_size_rows, cardinality, frequent_set_ratio); + } else { + return build_table(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(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(sum) / static_cast(page_dict_bits.size())); + state.add_element_count(static_cast(*min_it), "dict_rle_bits_min"); + state.add_element_count(static_cast(*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}); diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index d24d82d518f1..e7800da5e5ca 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -8,19 +8,40 @@ #include #include #include +#include +#include #include #include +#include +#include +#include #include #include +#include +#include namespace cudf::io::parquet::detail { namespace { + +/// Maximum number of chunk fragments for which spatially-local dictionary indices across fragment +/// keys are claimed using shared memory +constexpr size_type MAX_FRAGMENTS_PER_CHUNK = 1024; + +/// Default block size for kernel launches constexpr int DEFAULT_BLOCK_SIZE = 256; -} +/** + * @brief Functor for checking equality of two keys. + * + * @tparam T Type of the keys + * @param col Column view + * @param lhs_idx Index of the left-hand side key + * @param rhs_idx Index of the right-hand side key + * @return Whether the keys are equal + */ template struct equality_functor { column_device_view const& col; @@ -32,6 +53,14 @@ struct equality_functor { } }; +/** + * @brief Functor for hashing a key + * + * @tparam T Type of the key + * @param col Column view + * @param idx Key index of the key + * @return Hash value of the key + */ template struct hash_functor { column_device_view const& col; @@ -42,21 +71,36 @@ struct hash_functor { } }; +/** + * @brief Functor for inserting a key into a hash map + * + * @tparam block_size Thread block size + * @param storage_ref Hash map storage reference + * @param chunk Column chunk pointer + * @param frag Page fragment pointer + * @param frag_idx Page fragment index + */ template struct map_insert_fn { storage_ref_type const& storage_ref; - EncColumnChunk* const& chunk; + EncColumnChunk* const chunk; + PageFragment* const frag; + mapped_type const frag_idx; template - __device__ void operator()(size_type const s_start_value_idx, size_type const end_value_idx) + __device__ void operator()(size_type const start_value_idx, size_type const end_value_idx) { if constexpr (column_device_view::has_element_accessor()) { using block_reduce = cub::BlockReduce; __shared__ typename block_reduce::TempStorage reduce_storage; + auto const t = threadIdx.x; auto const col = chunk->col_desc; column_device_view const& data_col = *col->leaf_column; __shared__ size_type total_num_dict_entries; + __shared__ size_type num_dict_vals; + if (t == 0) { num_dict_vals = 0; }; + __syncthreads(); using equality_fn_type = equality_functor; using hash_fn_type = hash_functor; @@ -74,14 +118,13 @@ struct map_insert_fn { // Create a map ref with `cuco::insert` operator auto map_insert_ref = hash_map_ref.rebind_operators(cuco::insert); - auto const t = threadIdx.x; // Create atomic refs to the current chunk's num_dict_entries and uniq_data_size cuda::atomic_ref const chunk_num_dict_entries{chunk->num_dict_entries}; cuda::atomic_ref const chunk_uniq_data_size{chunk->uniq_data_size}; // Note: Adjust the following loop to use `cg::tile` if needed in the future. - for (thread_index_type val_idx = s_start_value_idx + t; val_idx - t < end_value_idx; + for (size_type val_idx = start_value_idx + t; val_idx - t < end_value_idx; val_idx += block_size) { size_type is_unique = 0; size_type uniq_elem_size = 0; @@ -90,26 +133,29 @@ struct map_insert_fn { auto const is_valid = val_idx < end_value_idx and val_idx < data_col.size() and data_col.is_valid(val_idx); - // Insert tile_val_idx to hash map and count successful insertions. + // Insert fragment index to hash map using a single thread (for best performance for now) + // and count successful insertions. if (is_valid) { - // Insert the keys using a single thread for best performance for now. - is_unique = map_insert_ref.insert(cuco::pair{val_idx, val_idx}); + // TODO(mh): Here we insert the fragment index of the CAS winner, which may not be the + // smallest one (relies on monotonic block scheduling). Switch to static_map's + // `insert_or_apply` with `cuco::op::min` for deterministic first-fragment semantics + is_unique = map_insert_ref.insert(slot_type{static_cast(val_idx), frag_idx}); uniq_elem_size = [&]() -> size_type { if (not is_unique) { return 0; } switch (col->physical_type) { - case Type::INT32: return 4; - case Type::INT64: return 8; - case Type::INT96: return 12; - case Type::FLOAT: return 4; - case Type::DOUBLE: return 8; + case Type::INT32: return sizeof(int32_t); + case Type::INT64: return sizeof(int64_t); + case Type::INT96: return sizeof(int32_t) + sizeof(int64_t); + case Type::FLOAT: return sizeof(float); + case Type::DOUBLE: return sizeof(double); case Type::BYTE_ARRAY: { auto const col_type = data_col.type().id(); if (col_type == type_id::STRING) { - // Strings are stored as 4 byte length + string bytes - return 4 + data_col.element(val_idx).size_bytes(); + // Strings are stored as int32_t length + string bytes + return sizeof(int32_t) + data_col.element(val_idx).size_bytes(); } else if (col_type == type_id::LIST) { - // Binary is stored as 4 byte length + bytes - return 4 + + // Binary is stored as int32_t length + bytes + return sizeof(int32_t) + get_element(data_col, val_idx).size_bytes(); } CUDF_UNREACHABLE( @@ -129,36 +175,47 @@ struct map_insert_fn { auto num_unique = block_reduce(reduce_storage).Sum(is_unique); __syncthreads(); auto uniq_data_size = block_reduce(reduce_storage).Sum(uniq_elem_size); - // The first thread in the block atomically updates total num_unique and uniq_data_size + // One thread atomically updates the number and data size of total unique values as well as + // the number of unique values inserted by this fragment. if (t == 0) { total_num_dict_entries = chunk_num_dict_entries.fetch_add(num_unique, cuda::std::memory_order_relaxed); total_num_dict_entries += num_unique; + num_dict_vals += num_unique; chunk_uniq_data_size.fetch_add(uniq_data_size, cuda::std::memory_order_relaxed); } __syncthreads(); // Check if the num unique values in chunk has already exceeded max dict size and early exit - if (total_num_dict_entries > MAX_DICT_SIZE) { return; } + if (total_num_dict_entries > MAX_DICT_SIZE) { break; } } // for loop + // Flush the number of unique values inserted by this fragment + if (t == 0) { frag->num_dict_vals = num_dict_vals; }; } else { CUDF_UNREACHABLE("Unsupported type to insert in map"); } } }; +/** + * @brief Functor for finding a key in a hash map + * + * @tparam block_size Thread block size + * @param storage_ref Hash map storage reference + * @param chunk Column chunk pointerk + */ template struct map_find_fn { storage_ref_type const& storage_ref; - EncColumnChunk* const& chunk; + EncColumnChunk* const chunk; template - __device__ void operator()(size_type const s_start_value_idx, + __device__ void operator()(size_type const start_value_idx, size_type const end_value_idx, - size_type const s_ck_start_val_idx) + size_type const ck_start_val_idx) { if constexpr (column_device_view::has_element_accessor()) { - auto const col = chunk->col_desc; - column_device_view const& data_col = *col->leaf_column; + auto const col = chunk->col_desc; + auto const& data_col = *col->leaf_column; using equality_fn_type = equality_functor; using hash_fn_type = hash_functor; @@ -179,8 +236,7 @@ struct map_find_fn { auto const t = threadIdx.x; // Note: Adjust the following loop to use `cg::tiles` if needed in the future. - for (thread_index_type val_idx = s_start_value_idx + t; val_idx < end_value_idx; - val_idx += block_size) { + for (key_type val_idx = start_value_idx + t; val_idx < end_value_idx; val_idx += block_size) { // Find the key using a single thread for best performance for now. if (data_col.is_valid(val_idx)) { auto const found_slot = map_find_ref.find(val_idx); @@ -188,7 +244,7 @@ struct map_find_fn { cudf_assert(found_slot != map_find_ref.end() && "Unable to find value in map in dictionary index construction"); // No need for atomic as this is not going to be modified by any other thread. - chunk->dict_index[val_idx - s_ck_start_val_idx] = found_slot->second; + chunk->dict_index[val_idx - ck_start_val_idx] = found_slot->second; } } } else { @@ -197,85 +253,154 @@ struct map_find_fn { } }; +/** + * @brief Populate the hash maps for all chunks (thread block per page fragment) + * + * @tparam block_size Thread block size + * @param map_storage Hash map storage span + * @param frags 2D span of page fragments + */ template CUDF_KERNEL void __launch_bounds__(block_size) populate_chunk_hash_maps_kernel(device_span const map_storage, - cudf::detail::device_2dspan frags) + cudf::detail::device_2dspan frags) { - auto const col_idx = blockIdx.y; - auto const block_x = blockIdx.x; - auto const frag = frags[col_idx][block_x]; - auto chunk = frag.chunk; - auto col = chunk->col_desc; + auto const col_idx = blockIdx.y; + auto const frag_idx = blockIdx.x; + auto& frag = frags[col_idx][frag_idx]; + auto const chunk = frag.chunk; + auto col = chunk->col_desc; if (not chunk->use_dictionary) { return; } - size_type start_row = frag.start_row; - size_type end_row = frag.start_row + frag.num_rows; + auto const start_row = frag.start_row; + auto const end_row = frag.start_row + frag.num_rows; // Find the bounds of values in leaf column to be inserted into the map for current chunk. - size_type const s_start_value_idx = row_to_value_idx(start_row, *col); - size_type const end_value_idx = row_to_value_idx(end_row, *col); + auto const start_value_idx = row_to_value_idx(start_row, *col); + auto const end_value_idx = row_to_value_idx(end_row, *col); column_device_view const& data_col = *col->leaf_column; storage_ref_type const storage_ref{chunk->dict_map_size, map_storage.data() + chunk->dict_map_offset}; - type_dispatcher(data_col.type(), - map_insert_fn{storage_ref, chunk}, - s_start_value_idx, - end_value_idx); + type_dispatcher( + data_col.type(), + map_insert_fn{storage_ref, chunk, &frag, static_cast(frag_idx)}, + start_value_idx, + end_value_idx); } +/** + * @brief Collect the dictionary indices for all chunks (thread block per column chunk) + * + * @tparam block_size Thread block size + * @param map_storage Hash map storage span + * @param chunks Column chunks span + * @param frags 2D span of page fragments + */ template CUDF_KERNEL void __launch_bounds__(block_size) collect_map_entries_kernel(device_span const map_storage, - device_span chunks) + device_span chunks, + cudf::detail::device_2dspan frags) { auto& chunk = chunks[blockIdx.x]; if (not chunk.use_dictionary) { return; } - auto t = threadIdx.x; - __shared__ cuda::atomic counter; - using cuda::std::memory_order_relaxed; - if (t == 0) { new (&counter) cuda::atomic{0}; } - __syncthreads(); - - // Iterate over all slots in the map. - for (; t < chunk.dict_map_size; t += block_size) { - auto* slot = map_storage.data() + chunk.dict_map_offset + t; - auto const key = slot->first; - if (key != KEY_SENTINEL) { - auto const loc = counter.fetch_add(1, memory_order_relaxed); - cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); - chunk.dict_data[loc] = key; - // If sorting dict page ever becomes a hard requirement, enable the following statement - // and add a dict sorting step before storing into the slot's second field. - // chunk.dict_data_idx[loc] = idx; - slot->second = loc; + auto t = threadIdx.x; + auto const num_frags = chunk.num_fragments; + + // If a chunk has less fragments than the maximum fragments per chunk, resolve its + // column-relative fragment range. + if (num_frags <= MAX_FRAGMENTS_PER_CHUNK) { + auto const col_idx = chunk.col_desc_id; + auto const& col_frags = frags[col_idx]; + auto const frag_start = static_cast(chunk.fragments - col_frags.data()); + + // Initialize fragment_offsets with prefix sum (exclusive) of number of dictionary values + // inserted by each fragment in this chunk. + __shared__ size_type fragment_offsets[MAX_FRAGMENTS_PER_CHUNK]; + { + using block_scan = cub::BlockScan; + __shared__ typename block_scan::TempStorage scan_storage; + + auto base_idx = 0; + while (base_idx < num_frags) { + auto const idx = base_idx + t; + auto const per_thread_count = + (idx < num_frags) ? col_frags[frag_start + idx].num_dict_vals : 0; + auto per_thread_offset = 0; + block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); + if (idx < num_frags) { fragment_offsets[idx] = per_thread_offset; } + base_idx += block_size; + __syncthreads(); + } + } + + // Iterate over slots and claim a dictionary index from the fragment offsets (spatial-locality + // for keys in the same fragment). + for (; t < chunk.dict_map_size; t += block_size) { + auto* slot = map_storage.data() + chunk.dict_map_offset + t; + auto const key = slot->first; + if (key != KEY_SENTINEL) { + auto const frag_loc = static_cast(slot->second) - frag_start; + cudf_assert(frag_loc >= 0 && frag_loc < num_frags && + "populate stamped a fragment hint outside this chunk's fragment range"); + auto const loc = atomicAdd(&fragment_offsets[frag_loc], 1); + cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); + chunk.dict_data[loc] = key; + slot->second = loc; + } + } + } + // Otherwise, iterate over slots and claim monotonically increasing dictionary indices for each + // key + else { + __shared__ cuda::atomic counter; + using cuda::std::memory_order_relaxed; + if (t == 0) { new (&counter) cuda::atomic{0}; } + __syncthreads(); + + for (; t < chunk.dict_map_size; t += block_size) { + auto* slot = map_storage.data() + chunk.dict_map_offset + t; + auto const key = slot->first; + if (key != KEY_SENTINEL) { + auto const loc = counter.fetch_add(1, memory_order_relaxed); + cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); + chunk.dict_data[loc] = key; + slot->second = loc; + } } } } +/** + * @brief Get the dictionary indices for all chunks (thread block per page fragment) + * + * @tparam block_size Thread block size + * @param map_storage Hash map storage span + * @param frags 2D span of page fragments + */ template CUDF_KERNEL void __launch_bounds__(block_size) get_dictionary_indices_kernel(device_span const map_storage, cudf::detail::device_2dspan frags) { - auto const col_idx = blockIdx.y; - auto const block_x = blockIdx.x; - auto const frag = frags[col_idx][block_x]; - auto chunk = frag.chunk; + auto const col_idx = blockIdx.y; + auto const frag_idx = blockIdx.x; + auto const& frag = frags[col_idx][frag_idx]; + auto const chunk = frag.chunk; if (not chunk->use_dictionary) { return; } - size_type start_row = frag.start_row; - size_type end_row = frag.start_row + frag.num_rows; + auto const start_row = frag.start_row; + auto const end_row = frag.start_row + frag.num_rows; auto const col = chunk->col_desc; // Find the bounds of values in leaf column to be searched in the map for current chunk - auto const s_start_value_idx = row_to_value_idx(start_row, *col); - auto const s_ck_start_val_idx = row_to_value_idx(chunk->start_row, *col); - auto const end_value_idx = row_to_value_idx(end_row, *col); + auto const start_value_idx = row_to_value_idx(start_row, *col); + auto const end_value_idx = row_to_value_idx(end_row, *col); + auto const ck_start_val_idx = row_to_value_idx(chunk->start_row, *col); column_device_view const& data_col = *col->leaf_column; storage_ref_type const storage_ref{chunk->dict_map_size, @@ -283,13 +408,76 @@ CUDF_KERNEL void __launch_bounds__(block_size) type_dispatcher(data_col.type(), map_find_fn{storage_ref, chunk}, - s_start_value_idx, + start_value_idx, end_value_idx, - s_ck_start_val_idx); + ck_start_val_idx); +} + +/** + * @brief Compute the RLE bits for the dictionary indices for all pages (warp per page) + * + * @tparam block_size Thread block size + * @param pages Pages span + */ +CUDF_KERNEL void __launch_bounds__(DEFAULT_BLOCK_SIZE) + compute_page_dict_bits_kernel(device_span pages) +{ + namespace cg = cooperative_groups; + using cudf::detail::warp_size; + + auto const page_idx = cudf::detail::grid_1d::global_thread_id() / warp_size; + if (page_idx >= static_cast(pages.size())) { return; } + + auto const warp = cg::tiled_partition(cg::this_thread_block()); + auto& page = pages[page_idx]; + auto const chunk = page.chunk; + + // Return if non-dict chunk: `dict_rle_bits` is unused by the encoder + if (not chunk->use_dictionary) { return; } + // Return if dictionary page itself does not encode dict_indices + if (page.page_type == PageType::DICTIONARY_PAGE) { return; } + + auto const col = chunk->col_desc; + + // Return if BOOLEAN column: dict_bits=1 through a separate code path in + // `gpuEncodeDictPages`, independent of `dict_rle_bits` + if (col->physical_type == Type::BOOLEAN) { return; } + + auto const chunk_start_val = row_to_value_idx(chunk->start_row, *col); + auto const page_start_val = row_to_value_idx(page.start_row, *col); + auto const page_num_leaf_values = static_cast(page.num_leaf_values); + auto const begin = page_start_val - chunk_start_val; + auto const end = begin + page_num_leaf_values; + + auto const* dict_index = chunk->dict_index; + auto const& leaf_col = *col->leaf_column; + auto const leaf_size = leaf_col.size(); + + // Accumulate per-lane max dict index for this page + size_type lane_max = 0; + for (size_type i = begin + warp.thread_rank(); i < end; i += warp_size) { + auto const val_idx = chunk_start_val + i; + // Null rows leave `dict_index` undefined; gate the read with the column's validity bitmap to + // avoid pulling garbage bits into the max. + if (val_idx < leaf_size && leaf_col.is_valid(val_idx)) { + lane_max = cuda::std::max(lane_max, dict_index[i]); + } + } + + // Write this page's RLE bits + auto const page_max = cg::reduce(warp, lane_max, cg::greater{}); + cg::invoke_one(warp, [&] { + // Floor at 1 to match the chunk-wide convention (all-null pages still emit a 1-bit RLE + // preamble) + auto const nbits = cuda::std::max(cuda::std::bit_width(static_cast(page_max)), 1); + page.dict_rle_bits = static_cast(nbits); + }); } +} // namespace + void populate_chunk_hash_maps(device_span const map_storage, - cudf::detail::device_2dspan frags, + cudf::detail::device_2dspan frags, rmm::cuda_stream_view stream) { dim3 const dim_grid(frags.size().second, frags.size().first); @@ -299,11 +487,15 @@ void populate_chunk_hash_maps(device_span const map_storage, void collect_map_entries(device_span const map_storage, device_span chunks, + cudf::detail::device_2dspan frags, rmm::cuda_stream_view stream) { constexpr int block_size = 1024; + static_assert(block_size >= MAX_FRAGMENTS_PER_CHUNK, + "block_size must be >= MAX_FRAGMENTS_PER_CHUNK so one BlockScan thread backs " + "each histogram bucket."); collect_map_entries_kernel - <<>>(map_storage, chunks); + <<>>(map_storage, chunks, frags); } void get_dictionary_indices(device_span const map_storage, @@ -314,4 +506,14 @@ void get_dictionary_indices(device_span const map_storage, get_dictionary_indices_kernel <<>>(map_storage, frags); } + +void compute_per_page_dict_bits(device_span pages, rmm::cuda_stream_view stream) +{ + if (pages.empty()) { return; } + auto constexpr warps_per_block = DEFAULT_BLOCK_SIZE / cudf::detail::warp_size; + auto const num_blocks = + cudf::util::div_rounding_up_safe(static_cast(pages.size()), warps_per_block); + compute_page_dict_bits_kernel<<>>(pages); +} + } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/page_enc.cu b/cpp/src/io/parquet/page_enc.cu index fcd90f6f710e..bf0edd14dc10 100644 --- a/cpp/src/io/parquet/page_enc.cu +++ b/cpp/src/io/parquet/page_enc.cu @@ -152,7 +152,7 @@ void __device__ init_frag_state(frag_init_state_s* const s, // frag.num_rows = fragment_size except for the last fragment in partition which can be // smaller. num_rows is fixed but fragment size could be larger if the data is strings or // nested. - s->frag.num_rows = min(fragment_size, part_end_row - s->frag.start_row); + s->frag.num_rows = cuda::std::min(fragment_size, part_end_row - s->frag.start_row); s->frag.num_dict_vals = 0; s->frag.fragment_data_size = 0; s->frag.dict_data_size = 0; @@ -642,7 +642,6 @@ CUDF_KERNEL void __launch_bounds__(128) uint32_t num_rows = 0; uint32_t page_start = 0; uint32_t page_offset = ck_g.ck_stat_size; - uint32_t num_dict_entries = 0; uint32_t comp_page_offset = ck_g.ck_stat_size; uint32_t page_headers_size = 0; uint32_t max_page_data_size = 0; @@ -677,6 +676,7 @@ CUDF_KERNEL void __launch_bounds__(128) page_g.num_rows = ck_g.num_dict_entries; page_g.num_leaf_values = ck_g.num_dict_entries; page_g.num_values = ck_g.num_dict_entries; // TODO: shouldn't matter for dict page + page_g.dict_rle_bits = ck_g.dict_rle_bits; // TODO: shouldn't matter for dict page page_offset += util::round_up_unsafe(page_g.max_hdr_size + page_g.max_data_size, page_align); if (not comp_page_sizes.empty()) { @@ -778,7 +778,9 @@ CUDF_KERNEL void __launch_bounds__(128) page_g.data_size = 0; page_g.comp_data_size = 0; page_g.is_compressed = false; - page_g.max_hdr_size = max_data_page_hdr_size; // Max size excluding statistics + page_g.dict_rle_bits = + ck_g.dict_rle_bits; // Conservatively set to the chunk-wide bit width + page_g.max_hdr_size = max_data_page_hdr_size; // Max size excluding statistics if (ck_g.stats) { uint32_t stats_hdr_len = 16; if (col_g.stats_dtype == dtype_string || col_g.stats_dtype == dtype_byte_array) { @@ -890,7 +892,6 @@ CUDF_KERNEL void __launch_bounds__(128) max_stats_len = 0; } max_stats_len = max(max_stats_len, minmax_len); - num_dict_entries += frag_g.num_dict_vals; page_size += fragment_data_size; // fragment_data_size includes the length indicator...remove it var_bytes_size += frag_g.fragment_data_size - frag_g.num_valid * sizeof(size_type); @@ -1911,7 +1912,7 @@ CUDF_KERNEL void __launch_bounds__(block_size, 8) // TODO assert dict_bits >= 0 auto const dict_bits = (physical_type == Type::BOOLEAN) ? 1 : (s->ck.use_dictionary and s->page.page_type != PageType::DICTIONARY_PAGE) - ? s->ck.dict_rle_bits + ? s->page.dict_rle_bits // Use `page.dict_rle_bits` for data pages : -1; if (t == 0) { uint8_t* dst = s->cur; diff --git a/cpp/src/io/parquet/parquet_gpu.cuh b/cpp/src/io/parquet/parquet_gpu.cuh index ed79cb1ff06d..a57e405ec6c0 100644 --- a/cpp/src/io/parquet/parquet_gpu.cuh +++ b/cpp/src/io/parquet/parquet_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -91,18 +91,23 @@ inline size_type __device__ row_to_value_idx(size_type idx, * @param stream CUDA stream to use */ void populate_chunk_hash_maps(device_span const map_storage, - cudf::detail::device_2dspan frags, + cudf::detail::device_2dspan frags, rmm::cuda_stream_view stream); /** * @brief Compact dictionary hash map entries into chunk.dict_data * + * `dict_id`s are assigned in fragment-first-insert order, so pages over earlier + * fragments see smaller ids and can use narrower RLE bit widths. + * * @param map_storage Bulk hashmap storage * @param chunks Flat span of chunks to compact hash maps for + * @param frags 2D span of per-column page fragments * @param stream CUDA stream to use */ void collect_map_entries(device_span const map_storage, device_span chunks, + cudf::detail::device_2dspan frags, rmm::cuda_stream_view stream); /** @@ -122,4 +127,12 @@ void get_dictionary_indices(device_span const map_storage, cudf::detail::device_2dspan frags, rmm::cuda_stream_view stream); +/** + * @brief Compute the minimum width required for the dictionary indices for each data page + * + * @param pages Device span of encoder pages + * @param stream CUDA stream to use + */ +void compute_per_page_dict_bits(device_span pages, rmm::cuda_stream_view stream); + } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index 7d07f39aa388..5892c5086179 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -530,8 +530,8 @@ struct PageFragment { //!< non-leaf level uint32_t num_valid; //>, std::vector>> build_chunk_dictionaries(hostdevice_2dvector& chunks, host_span col_desc, - device_2dspan frags, + device_2dspan frags, compression_type compression, dictionary_policy dict_policy, size_t max_dict_size, @@ -1410,7 +1410,7 @@ build_chunk_dictionaries(hostdevice_2dvector& chunks, chunk.dict_index = inserted_dict_index.data(); } chunks.host_to_device_async(stream); - collect_map_entries(map_storage_data, chunks.device_view().flat_view(), stream); + collect_map_entries(map_storage_data, chunks.device_view().flat_view(), frags, stream); get_dictionary_indices(map_storage_data, frags, stream); return std::pair(std::move(dict_data), std::move(dict_index)); @@ -1899,6 +1899,7 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta, ck.start_row = start_row; ck.num_rows = (uint32_t)row_group.num_rows; ck.first_fragment = c * num_fragments + f; + ck.num_fragments = fragments_in_chunk; ck.encodings = 0; auto chunk_fragments = row_group_fragments[c].subspan(f, fragments_in_chunk); // In fragment struct, add a pointer to the chunk it belongs to @@ -1963,6 +1964,7 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta, EncColumnChunk& ck = chunks[r + first_rg_in_part[p]][c]; ck.fragments = page_fragments.device_ptr(frag_offset); ck.first_fragment = frag_offset; + ck.num_fragments = fragments_in_chunk; // update the chunk pointer here for each fragment in chunk.fragments for (uint32_t i = 0; i < fragments_in_chunk; i++) { @@ -2119,6 +2121,10 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta, max_page_size_rows, write_v2_headers, stream); + + // Now that page boundaries are finalized and dictionary indices have been materialized, compute + // minimum required RLE bit width for each data page + compute_per_page_dict_bits({pages.data(), pages.size()}, stream); } // Check device write support for all chunks and initialize bounce_buffer. diff --git a/cpp/tests/io/parquet_misc_test.cpp b/cpp/tests/io/parquet_misc_test.cpp index 4bcf31146a33..f37cb8ccfc6e 100644 --- a/cpp/tests/io/parquet_misc_test.cpp +++ b/cpp/tests/io/parquet_misc_test.cpp @@ -12,7 +12,9 @@ #include +#include #include +#include //////////////////////////////// // delta encoding writer tests @@ -170,9 +172,18 @@ TEST_P(ParquetSizedTest, DictionaryTest) EXPECT_TRUE(used_dict); // and check that the correct number of bits was used - auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); - auto const nbits = read_dict_bits(source, oi.page_locations.front()); - EXPECT_EQ(nbits, GetParam()); + auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); + std::vector page_dict_bits; + page_dict_bits.reserve(oi.page_locations.size()); + std::transform( + oi.page_locations.begin(), + oi.page_locations.end(), + std::back_inserter(page_dict_bits), + [&source](auto const& page_location) { return read_dict_bits(source, page_location); }); + + auto const max_bits = std::max_element(page_dict_bits.begin(), page_dict_bits.end()); + ASSERT_NE(max_bits, page_dict_bits.end()); + EXPECT_EQ(*max_bits, GetParam()); } /////////////////////// diff --git a/cpp/tests/io/parquet_writer_test.cpp b/cpp/tests/io/parquet_writer_test.cpp index f54a36d4b870..fd285daabe60 100644 --- a/cpp/tests/io/parquet_writer_test.cpp +++ b/cpp/tests/io/parquet_writer_test.cpp @@ -22,13 +22,19 @@ #include #include +#include + #include #include +#include #include +#include #include #include +#include +#include using cudf::test::iterators::no_nulls; @@ -1106,6 +1112,102 @@ TEST_F(ParquetWriterTest, SingleValueDictionaryTest) EXPECT_EQ(nbits, expected_bits); } +TEST_F(ParquetWriterTest, VariableBitWidthDictEncoding) +{ + constexpr auto num_rows = 100'000; + constexpr auto num_pages = 10; + constexpr auto page_size = num_rows / num_pages; + constexpr auto freq_pages = num_pages - 2; + constexpr auto cardinality = 64'000; + constexpr auto frequent_set_size = 64; + + auto const filepath = temp_env->get_temp_filepath("VariableBitWidthDictEncoding.parquet"); + { + std::mt19937 rng{0xACAD1A}; + using ColumnType = cudf::test::fixed_width_column_wrapper; + + // Hot pages contain values in [0, frequent_set_size - 1], (first `freq_pages` pages in the + // chunk). Rare pages contain values in [frequent_set_size, cardinality - 1], (last `num_pages - + // freq_pages` pages in the chunk) + std::uniform_int_distribution freq_dist(0, frequent_set_size - 1); + std::uniform_int_distribution rare_dist(frequent_set_size, cardinality - 1); + auto constexpr threshold = freq_pages * page_size; + auto values = std::vector{}; + values.reserve(num_rows); + std::transform( + cuda::counting_iterator(0), + cuda::counting_iterator(num_rows), + std::back_inserter(values), + [&](auto row_idx) { return row_idx < threshold ? freq_dist(rng) : rare_dist(rng); }); + auto const col = ColumnType(values.begin(), values.end()); + + auto writer_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, table_view{{col}}) + .compression(cudf::io::compression_type::NONE) // No compression + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) // Write page index + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) // Always dictionary encode + .row_group_size_rows(num_rows) // Single row group only + .max_page_size_rows(page_size) // Max page size is set to the page size + .max_page_size_bytes(std::size_t{64} << 20) // Unlimited page size + .build(); + cudf::io::write_parquet(writer_opts); + } + + auto datasource = cudf::io::datasource::create(filepath); + + // Extract dictionary bit width for each data page from page index + std::vector page_dict_bits; + { + // Read file metadata + cudf::io::parquet::FileMetaData file_metadata; + read_footer(datasource, &file_metadata); + + // Check dictionary encoded pages + auto const& colchunk = file_metadata.row_groups.front().columns.front(); + EXPECT_TRUE(std::any_of(colchunk.meta_data.encodings.begin(), + colchunk.meta_data.encodings.end(), + [](auto const encoding) { + return encoding == cudf::io::parquet::Encoding::PLAIN_DICTIONARY or + encoding == cudf::io::parquet::Encoding::RLE_DICTIONARY; + })); + + auto const oi = read_offset_index(datasource, colchunk); + page_dict_bits.reserve(oi.page_locations.size()); + std::transform(oi.page_locations.begin(), + oi.page_locations.end(), + std::back_inserter(page_dict_bits), + [&datasource](auto const& page_location) { + return read_dict_bits(datasource, page_location); + }); + + // Check page count + EXPECT_EQ(oi.page_locations.size(), static_cast(num_pages)); + } + + // Checks + { + // Check min and max bit widths + auto const [min_bits_iter, max_bits_iter] = std::ranges::minmax_element(page_dict_bits); + auto const chunk_wide_max_bits = std::bit_width(cardinality - 1); + auto const frequent_max_bits = std::bit_width(frequent_set_size - 1); + + ASSERT_FALSE(page_dict_bits.empty()); + EXPECT_GT(*min_bits_iter, 1); + EXPECT_GT(*max_bits_iter, frequent_max_bits); + EXPECT_LE(*max_bits_iter, chunk_wide_max_bits); + + // Check expected number of freq and rare pages + // TODO(mh): Race-dependent checks. Enable these along with cuDF PR #22323. + + // auto const total_page_count = static_cast(page_dict_bits.size()); + // auto const freq_page_count = static_cast( + // std::ranges::count_if(page_dict_bits, [&](int nbits) { return nbits <= frequent_max_bits; + // })); + // EXPECT_EQ(freq_page_count, freq_pages); + // EXPECT_EQ(total_page_count - freq_page_count, num_pages - freq_pages); + } +} + TEST_F(ParquetWriterTest, DictionaryNeverTest) { constexpr unsigned int nrows = 1'000U; From 8acc070d980568f8f29ac4035455913d514c4ae4 Mon Sep 17 00:00:00 2001 From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com> Date: Fri, 15 May 2026 17:32:33 -0700 Subject: [PATCH 5/5] Add additional AI review guidelines for pylibcudf/cudf_polars (#22340) closes https://github.com/rapidsai/cudf/issues/21919 Authors: - Matthew Roeschke (https://github.com/mroeschke) Approvers: - Matthew Murray (https://github.com/Matt711) URL: https://github.com/rapidsai/cudf/pull/22340 --- python/REVIEW_GUIDELINES.md | 7 +++- python/pylibcudf/pylibcudf/interop.pyi | 51 +------------------------- python/pylibcudf/pylibcudf/interop.pyx | 2 - 3 files changed, 7 insertions(+), 53 deletions(-) diff --git a/python/REVIEW_GUIDELINES.md b/python/REVIEW_GUIDELINES.md index dceb696e06c0..dd251be9472f 100644 --- a/python/REVIEW_GUIDELINES.md +++ b/python/REVIEW_GUIDELINES.md @@ -58,14 +58,19 @@ ### pylibcudf (Cython Bindings) - Incorrect Cython object lifetime management -- Exceptions not handled correctly across Python/C++ boundary +- Exceptions not handled correctly across Python/C++ boundary (missing `+libcudf_exception_handler` if not `noexcept`) +- Cython binding of a C++ function declaring `noexcept` when the C++ function can raise exceptions - Incorrect GIL handling for CUDA operations - Cython bindings not matching the C++ API +- Using pylibcudf or Polars APIs that require pyarrow (like `to_arrow`) when cudf_polars containers should be used instead ### cudf_polars (Polars GPU Executor) - Missing coverage of Polars expression types (silent fallback to CPU without warning) - Incorrect GPU executor fallback logic - IR nodes not properly translated +- Stream argument not explicitly passed to a pylibcudf API +- `asyncio.Task`s not explicitly canceled in a finally block upon failure +- rapidsmpf `Channel`s not eventually entering the `shutdown_on_error` context manager ### dask_cudf - Dask DataFrame API compatibility issues diff --git a/python/pylibcudf/pylibcudf/interop.pyi b/python/pylibcudf/pylibcudf/interop.pyi index 34fe9394f7dd..3e3666219f44 100644 --- a/python/pylibcudf/pylibcudf/interop.pyi +++ b/python/pylibcudf/pylibcudf/interop.pyi @@ -1,18 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 -from collections.abc import Iterable, Mapping from dataclasses import dataclass -from typing import Any, overload - -import pyarrow as pa +from typing import Any from rmm.pylibrmm.memory_resource import DeviceMemoryResource -from pylibcudf.column import Column -from pylibcudf.scalar import Scalar from pylibcudf.table import Table -from pylibcudf.types import DataType from pylibcudf.utils import CudaStreamLike @dataclass @@ -22,49 +16,6 @@ class ColumnMetadata: precision: int | None = ... children_meta: list[ColumnMetadata] = ... -@overload -def from_arrow(obj: pa.DataType) -> DataType: ... -@overload -def from_arrow( - obj: pa.Scalar[Any], *, data_type: DataType | None = None -) -> Scalar: ... -@overload -def from_arrow( - obj: pa.Array[Any], - *, - data_type: DataType | None = None, - stream: CudaStreamLike | None = None, - mr: DeviceMemoryResource | None = None, -) -> Column: ... -@overload -def from_arrow( - obj: pa.Table, - *, - stream: CudaStreamLike | None = None, - mr: DeviceMemoryResource | None = None, -) -> Table: ... -@overload -def to_arrow( - obj: DataType, - *, - precision: int | None = None, - fields: Iterable[pa.Field[pa.DataType] | tuple[str, pa.DataType]] - | Mapping[str, pa.DataType] - | None = None, - value_type: pa.DataType | None = None, -) -> pa.DataType: ... -@overload -def to_arrow( - obj: Table, metadata: list[ColumnMetadata | str] | None = None -) -> pa.Table: ... -@overload -def to_arrow( - obj: Column, metadata: ColumnMetadata | str | None = None -) -> pa.Array[Any]: ... -@overload -def to_arrow( - obj: Scalar, metadata: ColumnMetadata | str | None = None -) -> pa.Scalar[Any]: ... def from_dlpack( managed_tensor: Any, stream: CudaStreamLike | None = None, diff --git a/python/pylibcudf/pylibcudf/interop.pyx b/python/pylibcudf/pylibcudf/interop.pyx index 23c47bb090fe..b43233ef5492 100644 --- a/python/pylibcudf/pylibcudf/interop.pyx +++ b/python/pylibcudf/pylibcudf/interop.pyx @@ -28,9 +28,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t __all__ = [ "ColumnMetadata", - "from_arrow", "from_dlpack", - "to_arrow", "to_dlpack", ]