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/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 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/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 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; 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/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] 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", ]