Skip to content

Experimental Parquet reader supports retention vectors #23402

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 5 commits into from
Jul 23, 2026
Merged

Experimental Parquet reader supports retention vectors #23402

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7c2e007
Support retention vectors in parquet experimental reader
mhaseeb123 Jul 22, 2026
f6231f6
Minor bug fix
mhaseeb123 Jul 22, 2026
99d1436
Add gtests
mhaseeb123 Jul 22, 2026
e305f18
Copyright header fix
mhaseeb123 Jul 22, 2026
00c95c2
Bug fixing
mhaseeb123 Jul 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions cpp/include/cudf/io/experimental/deletion_vectors.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,9 @@ struct deletion_vector_info {
std::vector<size_t> row_group_offsets;
/// Number of rows in each row group to be read from the Parquet source(s)
std::vector<size_type> row_group_num_rows;

/// Whether the roaring bitmaps represent retention vectors
bool are_retention_vectors = false;
Comment thread
coderabbitai[bot] marked this conversation as resolved.
};

/**
Expand Down Expand Up @@ -147,6 +150,7 @@ class chunked_parquet_reader {
std::queue<size_type> _deletion_vector_row_counts;
size_t _start_row;
bool _is_unspecified_row_group_data;
bool _are_retentions;
rmm::cuda_stream_view _stream;
rmm::device_async_resource_ref _mr;
rmm::device_async_resource_ref _table_mr;
Expand Down
18 changes: 12 additions & 6 deletions cpp/src/io/parquet/experimental/deletion_vectors.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,11 +104,15 @@ namespace detail {
deletion_vector_row_counts,
stream,
cudf::get_current_device_resource_ref());

auto const mask_type = deletion_vector_info.are_retention_vectors
? cudf::detail::mask_type::RETENTION
: cudf::detail::mask_type::DELETION;

// Filter the table using the deletion vector
return table_with_metadata{
return cudf::io::table_with_metadata{
// Supply user-provided mr to apply deletion mask to allocate output table's memory
cudf::detail::apply_mask(
table_with_index->view(), row_mask->view(), cudf::detail::mask_type::DELETION, stream, mr),
cudf::detail::apply_mask(table_with_index->view(), row_mask->view(), mask_type, stream, mr),
std::move(metadata)};
}

Expand Down Expand Up @@ -200,6 +204,7 @@ chunked_parquet_reader::chunked_parquet_reader(std::size_t chunk_read_limit,
rmm::device_async_resource_ref mr)
: _start_row{0},
_is_unspecified_row_group_data{deletion_vector_info.row_group_offsets.empty()},
_are_retentions{deletion_vector_info.are_retention_vectors},
_stream{stream},
_mr{mr},
// Use default mr for the internal chunked reader and row index column if we will
Expand Down Expand Up @@ -317,10 +322,11 @@ table_with_metadata chunked_parquet_reader::read_chunk()
_deletion_vector_row_counts,
_stream,
cudf::get_current_device_resource_ref());
return table_with_metadata{
auto const mask_type =
_are_retentions ? cudf::detail::mask_type::RETENTION : cudf::detail::mask_type::DELETION;
return cudf::io::table_with_metadata{
// Supply user-provided mr to apply deletion mask to allocate output table's memory
cudf::detail::apply_mask(
table_with_index->view(), row_mask->view(), cudf::detail::mask_type::DELETION, _stream, _mr),
cudf::detail::apply_mask(table_with_index->view(), row_mask->view(), mask_type, _stream, _mr),
std::move(metadata)};
}

Expand Down
190 changes: 118 additions & 72 deletions cpp/tests/io/parquet_deletion_vectors_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/

Expand Down Expand Up @@ -136,14 +136,17 @@ auto build_expected_row_indices(cudf::host_span<std::size_t const> row_group_off
* @param num_rows Number of rows in the table
* @param deletion_probability The probability of a row being deleted
* @param row_indices Host vector of row indices
* @param are_retention_vectors Whether to add retained, rather than deleted, row indices to the
* bitmap
*
* @return A pair of a deletion vector and a host row mask vector
* @return A pair of a roaring bitmap and a host row mask vector
*/
auto build_deletion_vector_and_expected_row_mask(cudf::size_type num_rows,
float deletion_probability,
cudf::host_span<std::size_t const> row_indices,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
auto build_roaring_bitmap_and_expected_row_mask(cudf::size_type num_rows,
float deletion_probability,
cudf::host_span<std::size_t const> row_indices,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr,
bool are_retention_vectors = false)
{
static constexpr auto seed = 0xbaLL;
std::mt19937 engine{seed};
Expand All @@ -164,8 +167,9 @@ auto build_deletion_vector_and_expected_row_mask(cudf::size_type num_rows,
std::for_each(cuda::counting_iterator<size_t>(0),
cuda::counting_iterator<size_t>(num_rows),
[&](auto row_idx) {
// Insert provided host row index if the row is deleted in the row mask
if (not expected_row_mask[row_idx]) {
// For retention vectors, retain rows selected in the row mask. For deletion
// vectors, delete the remaining rows.
if (expected_row_mask[row_idx] == are_retention_vectors) {
roaring::api::roaring64_bitmap_add_bulk(
deletion_vector, &roaring64_context, row_indices[row_idx]);
}
Expand Down Expand Up @@ -222,7 +226,7 @@ std::unique_ptr<cudf::table> build_expected_table(
* well as the `cudf::io::parquet::experimental::chunked_parquet_reader`
*/
template <cudf::size_type num_concat = 1>
void test_read_parquet_and_apply_deletion_vector(
void test_read_parquet_and_apply_mask(
cudf::host_span<char const> parquet_buffer,
cudf::io::parquet::experimental::deletion_vector_info const& deletion_vector_info,
cudf::table_view const& input_table_view,
Expand All @@ -245,6 +249,7 @@ void test_read_parquet_and_apply_deletion_vector(
.deletion_vector_row_counts = std::vector(num_concat, input_table_view.num_rows()),
.row_group_offsets = {},
.row_group_num_rows = {},
.are_retention_vectors = deletion_vector_info.are_retention_vectors,
};

// Vector to hold the Parquet buffer spans
Expand Down Expand Up @@ -296,7 +301,7 @@ void test_read_parquet_and_apply_deletion_vector(
local_expected_row_indices, cudf::type_id::UINT64, stream, mr);

auto [local_deletion_vector, local_expected_row_mask_column] =
build_deletion_vector_and_expected_row_mask(
build_roaring_bitmap_and_expected_row_mask(
num_input_rows, deletion_probability, local_expected_row_indices, stream, mr);
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

// Insert the expected table, the corresponding deletion vector and its data span
Expand Down Expand Up @@ -383,37 +388,58 @@ TEST_F(ParquetDeletionVectorsTest, NoRowIndexColumn)
auto expected_row_index_column = build_column_from_host_data<std::size_t>(
expected_row_indices, cudf::type_id::UINT64, stream, mr);

// Build deletion vector and the expected row mask column
auto [deletion_vector, expected_row_mask_column] = build_deletion_vector_and_expected_row_mask(
num_rows, deletion_probability, expected_row_indices, stream, mr);

// Use num_concat = 1 here since the row index column is simply a sequence and input table
// concatenation won't properly reset it.
auto constexpr num_concat = 1;
cudf::io::parquet::experimental::deletion_vector_info deletion_vector_info{
.serialized_roaring_bitmaps = {deletion_vector},
.deletion_vector_row_counts = {input_table->view().num_rows()}};
test_read_parquet_and_apply_deletion_vector<num_concat>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);

// Build deletion vector and the expected row mask column
{
auto [deletion_vector, expected_row_mask_column] = build_roaring_bitmap_and_expected_row_mask(
num_rows, deletion_probability, expected_row_indices, stream, mr);
auto deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{
.serialized_roaring_bitmaps = {deletion_vector},
.deletion_vector_row_counts = {input_table->view().num_rows()}};
test_read_parquet_and_apply_mask<num_concat>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
}

// Retention vectors retain the rows present in the bitmap.
{
auto constexpr are_retention_vectors = true;
auto [retention_vector, retention_row_mask_column] = build_roaring_bitmap_and_expected_row_mask(
num_rows, deletion_probability, expected_row_indices, stream, mr, are_retention_vectors);
auto retention_vector_info = cudf::io::parquet::experimental::deletion_vector_info{
.serialized_roaring_bitmaps = {retention_vector},
.deletion_vector_row_counts = {input_table->view().num_rows()},
.are_retention_vectors = are_retention_vectors};
test_read_parquet_and_apply_mask<num_concat>(parquet_buffer,
retention_vector_info,
input_table->view(),
retention_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
}

// Test no row index column and no deletion vector
{
// Build expected row mask column containing all true values
auto expected_row_mask = thrust::host_vector<bool>(num_rows, true);
auto expected_row_mask_column =
build_column_from_host_data<bool>(expected_row_mask, cudf::type_id::BOOL8, stream, mr);
deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{};
test_read_parquet_and_apply_deletion_vector<num_concat>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
auto deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{};
test_read_parquet_and_apply_mask<num_concat>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
}
}

Expand Down Expand Up @@ -477,7 +503,7 @@ TEST_F(ParquetDeletionVectorsTest, CustomRowIndexColumn)
expected_row_indices, cudf::type_id::UINT64, stream, mr);

// Build deletion vector and the expected row mask column
auto [deletion_vector, expected_row_mask_column] = build_deletion_vector_and_expected_row_mask(
auto [deletion_vector, expected_row_mask_column] = build_roaring_bitmap_and_expected_row_mask(
num_rows, deletion_probability, expected_row_indices, stream, mr);

// Don't concatenate the input table and test with single deletion vector
Expand All @@ -486,38 +512,58 @@ TEST_F(ParquetDeletionVectorsTest, CustomRowIndexColumn)
.deletion_vector_row_counts = {input_table->view().num_rows()},
.row_group_offsets = row_group_offsets,
.row_group_num_rows = row_group_num_rows};
test_read_parquet_and_apply_deletion_vector<1>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
test_read_parquet_and_apply_mask<1>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);

// Concatenate input table and test with multiple deletion vectors
test_read_parquet_and_apply_deletion_vector<4>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
test_read_parquet_and_apply_mask<4>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);

// Concatenate input table and test with many deletion vectors (>= stream fork threshold of 8)
test_read_parquet_and_apply_deletion_vector<8>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
test_read_parquet_and_apply_deletion_vector<16>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
test_read_parquet_and_apply_mask<8>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
test_read_parquet_and_apply_mask<16>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);

// Retention vectors with custom row indices.
{
auto constexpr are_retention_vectors = true;
auto [retention_vector, retention_row_mask_column] = build_roaring_bitmap_and_expected_row_mask(
num_rows, deletion_probability, expected_row_indices, stream, mr, are_retention_vectors);
deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{
.serialized_roaring_bitmaps = {retention_vector},
.deletion_vector_row_counts = {input_table->view().num_rows()},
.row_group_offsets = row_group_offsets,
.row_group_num_rows = row_group_num_rows,
.are_retention_vectors = are_retention_vectors};
test_read_parquet_and_apply_mask<1>(parquet_buffer,
deletion_vector_info,
input_table->view(),
retention_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
}

// Test custom row index column and no deletion vector
{
Expand All @@ -528,13 +574,13 @@ TEST_F(ParquetDeletionVectorsTest, CustomRowIndexColumn)

auto deletion_vector_info = cudf::io::parquet::experimental::deletion_vector_info{
.row_group_offsets = row_group_offsets, .row_group_num_rows = row_group_num_rows};
test_read_parquet_and_apply_deletion_vector<1>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
test_read_parquet_and_apply_mask<1>(parquet_buffer,
deletion_vector_info,
input_table->view(),
expected_row_mask_column->view(),
expected_row_index_column->view(),
stream,
mr);
}
}

Expand Down Expand Up @@ -562,7 +608,7 @@ TEST_F(DeletionVectorsCountTests, NoRowIndex)
auto row_indices = thrust::host_vector<size_t>(num_rows);
std::iota(row_indices.begin(), row_indices.end(), size_t{0});

auto [deletion_vector, expected_row_mask_column] = build_deletion_vector_and_expected_row_mask(
auto [deletion_vector, expected_row_mask_column] = build_roaring_bitmap_and_expected_row_mask(
num_rows, deletion_probability, row_indices, stream, mr);

auto const expected_row_mask = cudf::detail::make_host_vector(
Expand Down Expand Up @@ -628,7 +674,7 @@ TEST_F(DeletionVectorsCountTests, CustomRowIndex)
auto expected_row_indices =
build_expected_row_indices(row_group_offsets, row_group_num_rows, num_rows);

auto [deletion_vector, expected_row_mask_column] = build_deletion_vector_and_expected_row_mask(
auto [deletion_vector, expected_row_mask_column] = build_roaring_bitmap_and_expected_row_mask(
num_rows, deletion_probability, expected_row_indices, stream, mr);

auto const expected_row_mask = cudf::detail::make_host_vector(
Expand Down Expand Up @@ -681,7 +727,7 @@ TEST_F(DeletionVectorsCountTests, MultipleDeletionVectors)
auto local_indices =
cudf::host_span<size_t const>(expected_row_indices.data() + span_start, num_rows_per_dv);

auto [dv, mask_col] = build_deletion_vector_and_expected_row_mask(
auto [dv, mask_col] = build_roaring_bitmap_and_expected_row_mask(
num_rows_per_dv, deletion_probability, local_indices, stream, mr);

auto const host_mask = cudf::detail::make_host_vector(
Expand Down
Loading