diff --git a/cpp/src/io/parquet/reader_impl_dict_transcode.cu b/cpp/src/io/parquet/reader_impl_dict_transcode.cu index ef6425ed5e8..be01b711938 100644 --- a/cpp/src/io/parquet/reader_impl_dict_transcode.cu +++ b/cpp/src/io/parquet/reader_impl_dict_transcode.cu @@ -13,17 +13,22 @@ #include #include #include +#include #include #include #include #include #include +#include + #include +#include +#include +#include #include #include -#include #include #include @@ -151,6 +156,44 @@ void update_from_chunk(column_eligibility& e, ColumnChunkDesc const& chunk) return cudf::strings::detail::make_strings_column(begin, begin + entry_count, stream, mr); } +/** + * @brief Remap each row's dictionary index onto the deduplicated key space (in place). + * + * Each row's decoded index is local to its own row group's dictionary. This shifts that index into + * the row group's region of the stacked (non-deduplicated) key space, then translates it through + * `stacked_to_unique` -- the position-to-index map produced by encoding the stacked keys -- so it + * points at the correct entry in the compact, unique keys column. Done in place, in one pass over + * the rows, in lieu of `cudf::dictionary::detail::concatenate`. + * + * @param d_indices Device pointer to the INT32 index buffer, mutated in place + * @param num_rows Number of index values + * @param row_offsets Per-chunk row boundaries `[offsets[k], offsets[k+1])`, size num_chunks+1 + * @param key_counts_prefix Per-chunk key-prefix offsets into the stacked key space, size + * num_chunks+1 + * @param stacked_to_unique Map from stacked-key position to compact unique-key index + * @param stream CUDA stream used for the kernel launch + */ +void remap_dict_indices_by_chunk(int32_t* d_indices, + size_type num_rows, + cudf::device_span row_offsets, + cudf::device_span key_counts_prefix, + cudf::device_span stacked_to_unique, + rmm::cuda_stream_view stream) +{ + thrust::for_each(rmm::exec_policy_nosync(stream, get_current_device_resource_ref()), + cuda::counting_iterator{0}, + cuda::counting_iterator{num_rows}, + [row_offsets, key_counts_prefix, stacked_to_unique, d_indices] __device__( + size_type row) -> void { + // Chunk owning `row` is the last offset <= row. + auto const it = thrust::upper_bound( + thrust::seq, row_offsets.begin(), row_offsets.end(), row); + auto const k = static_cast(it - row_offsets.begin() - 1); + auto const stacked_pos = key_counts_prefix[k] + d_indices[row]; + d_indices[row] = stacked_to_unique[stacked_pos]; + }); +} + } // namespace void reader_impl::prepare_dict_transcode(read_mode mode) @@ -247,55 +290,79 @@ void reader_impl::assemble_dict_transcoded_columns( auto const& pass = *_pass_itm_data; - // For each eligible input column, collect its chunks in row-group order, build a per-chunk - // DICTIONARY32 segment (local 0-based indices + per-chunk keys column), and concatenate. + // Every string chunk's dictionary entries live contiguously in one buffer in + // `pass.str_dict_index` (each `chunk.str_dict_index` is a pointer into that one buffer). + // Materialize all keys into a single column using `make_strings_column` (contains duplicates). + // All keys stores keys of all columns, and not just column i. + std::unique_ptr all_keys; + auto ensure_all_keys = [&]() -> column_view { + if (all_keys == nullptr) { + all_keys = + make_keys_column_from_index_pairs(pass.str_dict_index.data(), + static_cast(pass.str_dict_index.size()), + _stream, + get_current_device_resource_ref()); + } + return all_keys->view(); + }; + + // Pre-pass 1: Map each chunk to its dictionary page's key count. + // Chunks without a dictionary page keep a count of 0. + std::vector chunk_dict_key_counts(pass.chunks.size(), 0); + for (auto const& page : pass.pages) { + if ((page.flags & PAGEINFO_FLAGS_DICTIONARY) == 0) { continue; } + auto const chunk_idx = page.chunk_idx; + if (chunk_idx < 0 or static_cast(chunk_idx) >= pass.chunks.size()) { continue; } + if (pass.chunks[chunk_idx].dict_page == nullptr) { continue; } + chunk_dict_key_counts[chunk_idx] = static_cast(page.num_input_values); + } + + // Pre-pass 2: Bucket chunk indices by their source input-column ordinal. Because + // `pass.chunks` is laid out row-group-major, appending in index order yields each column's chunks + // already in row-group order. + std::vector> chunks_by_input_col(_input_columns.size()); + for (size_t c = 0; c < pass.chunks.size(); ++c) { + auto const col = pass.chunks[c].src_col_index; + if (col >= 0 and static_cast(col) < _input_columns.size()) { + chunks_by_input_col[col].push_back(c); + } + } + + // For each eligible input column, collect its chunks in row-group order and assemble a + // DICTIONARY32 output. // - // IMPORTANT: Each segment carries row-group-local indices into its own keys column. We do NOT - // pre-shift indices into a global keyspace, because `cudf::dictionary::detail::concatenate` - // already re-maps the indices using `compute_children_offsets_fn`. Pre-shifting would cause - // double-offsetting and out-of-bounds reads in the `dispatch_compute_indices` kernel. + // A single-row-group column takes a zero-copy fast path (keys + decoded indices stapled + // together). A multi-row-group column stacks the per-chunk keys, deduplicates them, and remaps + // the decoded indices onto the compact key space in place -- avoiding + // `cudf::dictionary::detail::concatenate` and its redundant per-chunk index copy. std::for_each( cuda::counting_iterator{0}, cuda::counting_iterator{_input_columns.size()}, [&](size_t i) { if (not _dict_transcode_eligible[i]) { return; } - // Gather chunk indices for this input column in row-group order. - std::vector chunk_indices; - chunk_indices.reserve(pass.chunks.size() / std::max(_input_columns.size(), 1)); - std::copy_if(cuda::counting_iterator{0}, - cuda::counting_iterator{pass.chunks.size()}, - std::back_inserter(chunk_indices), - [&](size_t c) { return pass.chunks[c].src_col_index == static_cast(i); }); + // This column's chunks, in row-group order (bucketed in pre-pass 2 above). + auto const& chunk_indices = chunks_by_input_col[i]; if (chunk_indices.empty()) { return; } // `out_columns` is indexed by output-buffer (root column) ordinal, not input-column // ordinal: a nested struct/list column contributes one entry to `_output_buffers` but one // entry per leaf to `_input_columns`, so `i` and the corresponding root index can diverge // as soon as any nested column precedes this one. Eligibility requires a flat (depth-1) - // column, so `nesting[0]` is the correct, and only, output-buffer index to use here. + // column, so `nesting[0]` is the correct, and only, output-buffer index to use. auto const out_idx = static_cast(_input_columns[i].nesting[0]); - // Per-chunk key counts from the dictionary page's `num_input_values`, mirrored back to - // host when `pass.pages` was copied by `decode_page_headers`. - std::vector chunk_key_counts(chunk_indices.size(), 0); + // Per-chunk key counts, looked up from the pre-pass 1 map. + std::vector chunk_key_counts(chunk_indices.size()); std::transform(chunk_indices.begin(), chunk_indices.end(), chunk_key_counts.begin(), - [&](size_t chunk_idx) -> size_type { - if (pass.chunks[chunk_idx].dict_page == nullptr) { return 0; } - for (auto const& page : pass.pages) { - if (page.chunk_idx == static_cast(chunk_idx) and - (page.flags & PAGEINFO_FLAGS_DICTIONARY) != 0) { - return static_cast(page.num_input_values); - } - } - return size_type{0}; - }); + [&](size_t chunk_idx) { return chunk_dict_key_counts[chunk_idx]; }); auto& indices_col = out_columns[out_idx]; CUDF_EXPECTS(indices_col != nullptr and indices_col->type().id() == type_id::INT32, "Expected INT32 indices column for dict-transcoded flat string column"); + // Claim ownership of the indices column; the `out_idx` entry in `out_columns` is now empty. auto indices_owner = std::move(indices_col); // Single row group fast path: the Parquet dictionary page's entries become the keys as-is, @@ -321,13 +388,14 @@ void reader_impl::assemble_dict_transcoded_columns( // Keys were not distinct: fall through to the multi-row-group path, which deduplicates. } - // Multi-row-group path: the indices buffer is shared (aliased) by per-chunk DICTIONARY32 - // views below via the parent's offset/size, so it must stay alive until concatenate - // completes. - column_view const indices_view{indices_owner->view()}; + // Multi-row-group path (dedup-and-shift): stack every chunk's keys into a single column, + // deduplicate the key set once, then remap each row's index onto the compact key space in + // place. This avoids `cudf::dictionary::detail::concatenate`, which would re-copy the + // already-contiguous per-chunk indices (`indices_owner`) into a fresh buffer. + auto const num_row_vals = static_cast(indices_owner->size()); - // Per-chunk boundaries along the row axis: chunk k occupies rows - // [chunk_row_offsets[k], chunk_row_offsets[k+1]). + // Per-chunk row boundaries: chunk k occupies rows [chunk_row_offsets[k], + // chunk_row_offsets[k+1]). std::vector chunk_row_offsets(chunk_indices.size() + 1, 0); std::transform( chunk_indices.begin(), @@ -336,58 +404,90 @@ void reader_impl::assemble_dict_transcoded_columns( [&](size_t chunk_idx) { return static_cast(pass.chunks[chunk_idx].num_rows); }); std::inclusive_scan( chunk_row_offsets.begin() + 1, chunk_row_offsets.end(), chunk_row_offsets.begin() + 1); - CUDF_EXPECTS(chunk_row_offsets.back() == indices_view.size(), + CUDF_EXPECTS(chunk_row_offsets.back() == num_row_vals, "Row counts on pass chunks must sum to the indices column size"); - // Pre-compute null counts for all segments in a single kernel launch. Building the - // column_views below requires a per-segment null count, and calling null_count(begin, end) - // inside the loop would launch one kernel per chunk. Batch them here instead. - std::vector seg_null_counts(chunk_indices.size(), 0); - if (indices_view.nullable()) { - std::vector indices_pairs; - indices_pairs.reserve(chunk_indices.size() * 2); - for (size_t k = 0; k < chunk_indices.size(); ++k) { - indices_pairs.push_back(chunk_row_offsets[k]); - indices_pairs.push_back(chunk_row_offsets[k + 1]); + // Per-chunk key prefix offsets into the stacked key space: chunk k's keys occupy + // [key_counts_prefix[k], key_counts_prefix[k+1]). + std::vector key_counts_prefix(chunk_indices.size() + 1, 0); + std::inclusive_scan( + chunk_key_counts.begin(), chunk_key_counts.end(), key_counts_prefix.begin() + 1); + auto const total_keys = key_counts_prefix.back(); + + // Stack this column's per-chunk keys, sliced out of the batched keys view + // `all_string_column_keys` -- a view of the caller-scoped owning column `all_keys`, which + // outlives this block, so any view into it stays valid. Chunk `k`'s entries occupy + // `[key_offset, key_offset + chunk_key_counts[k])` in `pass.str_dict_index`, where + // `key_offset` is recovered from the chunk's stored pointer into that buffer. + // + // When those per-chunk ranges are already contiguous in `all_string_column_keys` -- e.g. a + // single string column, whose chunks are laid out consecutively -- the stacked keys are just + // one zero-copy sub-range of it, so the per-chunk gather (`concatenate`) is skipped entirely. + // Otherwise (multiple string columns interleaved row-group-major) the strided slices are + // concatenated into one contiguous column. + auto const all_string_column_keys = ensure_all_keys(); + auto const key_offset_of = [&](size_t k) { + return static_cast(pass.chunks[chunk_indices[k]].str_dict_index - + pass.str_dict_index.data()); + }; + bool contiguous = true; + for (size_t k = 0; k + 1 < chunk_indices.size(); ++k) { + if (key_offset_of(k + 1) != key_offset_of(k) + chunk_key_counts[k]) { + contiguous = false; + break; } - seg_null_counts = - cudf::detail::segmented_null_count(indices_view.null_mask(), indices_pairs, _stream); } - // Build a per-chunk DICTIONARY32 *view* that aliases the shared decoded INT32 buffer (no - // copy): keys = this chunk's STRING column, indices = `indices_view`. The row range, null - // mask, and null count must all live on the *parent* view (via offset/size), not the indices - // child, because `get_indices_annotated()` rebuilds the indices from the child's `head()` - // plus the parent's offset/size/null_mask -- anything set on the child is ignored. A wrong - // null count (e.g. a hardcoded 0) would silently turn nulls into a valid index once - // `cudf::detail::concatenate` remaps the indices against the unified keys. - std::vector> seg_keys_owners(chunk_indices.size()); - std::vector dict_segment_views(chunk_indices.size()); - std::transform( - cuda::counting_iterator{0}, - cuda::counting_iterator{chunk_indices.size()}, - dict_segment_views.begin(), - [&](size_t k) { - auto const chunk_idx = chunk_indices[k]; - auto const& chunk = pass.chunks[chunk_idx]; - - seg_keys_owners[k] = make_keys_column_from_index_pairs( - chunk.str_dict_index, chunk_key_counts[k], _stream, get_current_device_resource_ref()); - - auto const seg_begin = chunk_row_offsets[k]; - auto const seg_end = chunk_row_offsets[k + 1]; - auto const seg_rows = seg_end - seg_begin; - return column_view{data_type{type_id::DICTIONARY32}, - seg_rows, - nullptr, // dictionary parent holds no data - indices_view.null_mask(), // shared with indices_view - seg_null_counts[k], - seg_begin, // reslices shared indices child + null mask - {indices_view, seg_keys_owners[k]->view()}}; - }); - - // `cudf::detail::concatenate` deduplicates + sorts keys and recomputes indices. - out_columns[out_idx] = cudf::detail::concatenate(dict_segment_views, _stream, _mr); + std::unique_ptr stacked_keys_owner; // holds the gathered keys in the strided case + column_view const stacked_keys = [&] { + if (contiguous) { + auto const first = key_offset_of(0); + return cudf::detail::slice(all_string_column_keys, first, first + total_keys, _stream); + } + std::vector key_slices(chunk_indices.size()); + std::transform(cuda::counting_iterator{0}, + cuda::counting_iterator{chunk_indices.size()}, + key_slices.begin(), + [&](size_t k) { + return cudf::detail::slice(all_string_column_keys, + key_offset_of(k), + key_offset_of(k) + chunk_key_counts[k], + _stream); + }); + stacked_keys_owner = + cudf::detail::concatenate(key_slices, _stream, get_current_device_resource_ref()); + return stacked_keys_owner->view(); + }(); + + // Deduplicate the stacked keys. `encode` yields the compact unique keys (on `_mr`, the output + // keys child) plus an INT32 map from each stacked-key position to its compact index. + auto encoded = + cudf::dictionary::detail::encode(stacked_keys, data_type{type_id::INT32}, _stream, _mr); + auto encoded_contents = encoded->release(); + auto stacked_to_unique = + std::move(encoded_contents.children[0]); // INT32 map (keep for kernel) + auto unique_keys = std::move(encoded_contents.children[1]); // compact keys, owned on _mr + + // Remap every row's index onto the compact key space in place. Null rows carry a zero index + // (fill_pruned_offsets); the shift keeps them in range and the null mask (carried by + // `indices_owner`) still nullifies them in `decode`. + // + // These H2D copies are synchronous + auto const d_row_offsets = cudf::detail::make_device_uvector( + chunk_row_offsets, _stream, get_current_device_resource_ref()); + auto const d_key_counts_prefix = cudf::detail::make_device_uvector( + key_counts_prefix, _stream, get_current_device_resource_ref()); + remap_dict_indices_by_chunk( + indices_owner->mutable_view().data(), + num_row_vals, + cudf::device_span{d_row_offsets.data(), d_row_offsets.size()}, + cudf::device_span{d_key_counts_prefix.data(), d_key_counts_prefix.size()}, + cudf::device_span{stacked_to_unique->view().data(), + static_cast(stacked_to_unique->size())}, + _stream); + + out_columns[out_idx] = cudf::make_dictionary_column( + std::move(unique_keys), std::move(indices_owner), _stream, _mr); }); } diff --git a/cpp/tests/io/parquet_reader_dict_test.cpp b/cpp/tests/io/parquet_reader_dict_test.cpp index 80e2cf6dcf5..a490d4907ee 100644 --- a/cpp/tests/io/parquet_reader_dict_test.cpp +++ b/cpp/tests/io/parquet_reader_dict_test.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -57,9 +58,9 @@ std::string make_value_string(int value) return std::string{utf8_prefixes[value % utf8_prefixes.size()]} + "_" + std::to_string(value); } -cudf::test::strings_column_wrapper make_low_cardinality_strings() +cudf::test::strings_column_wrapper make_low_cardinality_strings(unsigned int col_seed = seed) { - std::mt19937 engine(seed); + std::mt19937 engine(col_seed); std::uniform_int_distribution value_dist(0, cardinality - 1); std::bernoulli_distribution null_dist(null_probability); @@ -486,3 +487,72 @@ TEST_F(ParquetReaderDictTest, MultiColumnMixedEligibility) ASSERT_EQ(read_key.type().id(), cudf::type_id::INT32); CUDF_TEST_EXPECT_COLUMNS_EQUAL(key_col, read_key); } + +// A low-cardinality flat string column spanning multiple row groups must transcode to a +// DICTIONARY32 with unique keys. Check if deduplication works correctly by checking for unique +// keys. +TEST_F(ParquetReaderDictTest, MultiRowGroupKeysAreUnique) +{ + auto input_col = make_low_cardinality_strings(); + + auto const input_tbl = cudf::table_view{{input_col}}; + auto const filepath = temp_env->get_temp_filepath("MultiRowGroupKeysAreUnique.parquet"); + write_parquet(input_tbl, filepath); + + auto const read_table = read_parquet_as_dict(filepath).tbl; + ASSERT_EQ(read_table->num_columns(), 1); + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::DICTIONARY32); + + cudf::dictionary_column_view const dict_view(read_col); + auto const keys = dict_view.keys(); + + // Keys must be unique and no larger than the source cardinality; a stacked-but-not-deduplicated + // dictionary would carry up to (number of row groups) times more keys. + auto const num_distinct = + cudf::distinct_count(keys, cudf::null_policy::INCLUDE, cudf::nan_policy::NAN_IS_VALID); + EXPECT_EQ(num_distinct, keys.size()); + EXPECT_LE(keys.size(), cardinality); + + // Check if the decoded column is equal to the original input. + auto const decoded = cudf::dictionary::decode(dict_view); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(input_col, decoded->view()); +} + +// Two flat, low-cardinality string columns across multiple row groups. Their per-row-group +// dictionaries interleave in the reader's shared key buffer (`pass.str_dict_index`), so each +// column's keys are strided within it. Both columns must transcode to DICTIONARY32 +// and decode back to their (distinct) inputs. +TEST_F(ParquetReaderDictTest, MultiStringColumnsDictTranscode) +{ + auto col_a = make_low_cardinality_strings(); // default seed + auto col_b = make_low_cardinality_strings(seed ^ 0xBE'EF01u); // distinct data + + auto const input_tbl = cudf::table_view{{col_a, col_b}}; + auto const filepath = temp_env->get_temp_filepath("MultiStringColumnsDictTranscode.parquet"); + write_parquet(input_tbl, filepath); // row_group_size rows/group -> multiple row groups + + auto const read_table = read_parquet_as_dict(filepath).tbl; + ASSERT_EQ(read_table->num_rows(), num_rows); + ASSERT_EQ(read_table->num_columns(), 2); + + auto const read_a = read_table->view().column(0); + auto const read_b = read_table->view().column(1); + ASSERT_EQ(read_a.type().id(), cudf::type_id::DICTIONARY32); + ASSERT_EQ(read_b.type().id(), cudf::type_id::DICTIONARY32); + + // Keys must be deduplicated in both columns -- the strided multi-column branch must produce + // unique keys (no larger than the cardinality) just like the contiguous single-column path. + for (auto const& read_col : {read_a, read_b}) { + auto const keys = cudf::dictionary_column_view(read_col).keys(); + auto const num_distinct = + cudf::distinct_count(keys, cudf::null_policy::INCLUDE, cudf::nan_policy::NAN_IS_VALID); + EXPECT_EQ(num_distinct, keys.size()); + EXPECT_LE(keys.size(), cardinality); + } + + auto const decoded_a = cudf::dictionary::decode(cudf::dictionary_column_view(read_a)); + auto const decoded_b = cudf::dictionary::decode(cudf::dictionary_column_view(read_b)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col_a, decoded_a->view()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col_b, decoded_b->view()); +}