Skip to content

JNI bridge for the retention filter support for deletion vectors #23423

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
gforsyth merged 3 commits into from
Jul 27, 2026
Merged

JNI bridge for the retention filter support for deletion vectors #23423

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
34de9b4
JNI implementation for deletion vector retention filter support
jihoonson Jul 24, 2026
f2439ed
new tests to test retention filter
jihoonson Jul 24, 2026
bd3b5c4
license and style fix
jihoonson Jul 24, 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
30 changes: 24 additions & 6 deletions java/src/main/java/ai/rapids/cudf/DeletionVector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*
*
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
*/
Expand Down Expand Up @@ -44,6 +44,8 @@ public static class DeletionVectorInfo {
*/
public final HostMemoryBuffer serializedBitmap;

public final boolean isRetention;

/**
* Row index offsets for each row group to read. Can be null to read all row groups.
*/
Expand All @@ -59,8 +61,9 @@ public static class DeletionVectorInfo {
*/
public final int totalNumRows;

public DeletionVectorInfo(HostMemoryBuffer serializedBitmap, long[] rowGroupOffsets, int[] rowGroupNumRows) {
public DeletionVectorInfo(HostMemoryBuffer serializedBitmap, boolean isRetention, long[] rowGroupOffsets, int[] rowGroupNumRows) {
Comment thread
igorpeshansky marked this conversation as resolved.
this.serializedBitmap = serializedBitmap;
this.isRetention = isRetention;
this.rowGroupOffsets = rowGroupOffsets;
this.rowGroupNumRows = rowGroupNumRows;
this.totalNumRows = computeTotalNumRows();
Expand All @@ -81,6 +84,16 @@ private int computeTotalNumRows() {
}
}

private static boolean getDeletionVectorTypes(DeletionVectorInfo[] deletionVectorInfos) {
boolean isRetention = deletionVectorInfos[0].isRetention;
for (DeletionVectorInfo info : deletionVectorInfos) {
if (info.isRetention != isRetention) {
throw new IllegalArgumentException("All DeletionVectorInfo objects must have the same isRetention value.");
}
}
return isRetention;
}

/**
* Reads a Parquet file with deletion vector support.
*
Expand Down Expand Up @@ -174,6 +187,7 @@ private static Table readParquet(ParquetOptions opts,
List<Integer> deletionVectorRowCountsList = new ArrayList<>(deletionVectorInfos.length);
List<Long> rowGroupOffsetsList = new ArrayList<>(deletionVectorInfos.length);
List<Integer> rowGroupNumRowsList = new ArrayList<>(deletionVectorInfos.length);
boolean areRetentionVectors = getDeletionVectorTypes(deletionVectorInfos);
if (deletionVectorInfos != null) {
for (DeletionVectorInfo info : deletionVectorInfos) {
serializedBitmapList.add(info.serializedBitmap);
Expand Down Expand Up @@ -202,7 +216,8 @@ private static Table readParquet(ParquetOptions opts,
bitmapAddrsSizes,
deletionVectorRowCounts,
rowGroupOffsets,
rowGroupNumRows);
rowGroupNumRows,
areRetentionVectors);
return new Table(columnHandles);
}

Expand Down Expand Up @@ -325,6 +340,7 @@ private ParquetChunkedReader(long chunkSizeByteLimit, long passReadLimit,
List<Integer> deletionVectorRowCountsList = new ArrayList<>(deletionVectorInfos.length);
List<Long> rowGroupOffsetsList = new ArrayList<>(deletionVectorInfos.length);
List<Integer> rowGroupNumRowsList = new ArrayList<>(deletionVectorInfos.length);
boolean areRetentionVectors = getDeletionVectorTypes(deletionVectorInfos);
if (deletionVectorInfos != null) {
for (DeletionVectorInfo info : deletionVectorInfos) {
serializedBitmapList.add(info.serializedBitmap);
Expand All @@ -347,7 +363,7 @@ private ParquetChunkedReader(long chunkSizeByteLimit, long passReadLimit,
long[] handles = createParquetChunkedReader(chunkSizeByteLimit, passReadLimit,
opts.getIncludeColumnNames(), opts.getReadBinaryAsString(), inputFilePaths,
dataBufferAddrsSizes, rowGroups, opts.timeUnit().typeId.getNativeId(),
bitmapAddrsSizes, deletionVectorRowCounts, rowGroupOffsets, rowGroupNumRows);
bitmapAddrsSizes, deletionVectorRowCounts, rowGroupOffsets, rowGroupNumRows, areRetentionVectors);
readerHandle = handles[0];
if (readerHandle == 0) {
throw new IllegalStateException("Cannot create native chunked Parquet reader object.");
Expand Down Expand Up @@ -431,7 +447,8 @@ private static native long[] readParquet(String[] filterColumnNames,
long[] serializedRoaring64,
int[] deletionVectorRowCounts,
long[] rowGroupOffsets,
int[] rowGroupNumRows)
int[] rowGroupNumRows,
boolean areRetentionVectors)
throws CudfException;

private static native long[] createParquetChunkedReader(long chunkReadLimit,
Expand All @@ -445,7 +462,8 @@ private static native long[] createParquetChunkedReader(long chunkReadLimit,
long[] serializedRoaringBitmaps,
int[] deletionVectorRowCounts,
long[] rowGroupOffsets,
int[] rowGroupNumRows)
int[] rowGroupNumRows,
boolean areRetentionVectors)
throws CudfException;

private static native boolean parquetChunkedReaderHasNext(long readerHandle) throws CudfException;
Expand Down
28 changes: 20 additions & 8 deletions java/src/main/native/src/DeletionVectorJni.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/

Expand Down Expand Up @@ -82,7 +82,8 @@ std::unique_ptr<cudf::io::parquet::experimental::deletion_vector_info> make_dele
jlongArray const& serialized_roaring64,
jintArray const& deletion_vector_row_counts,
jlongArray const& row_group_offsets,
jintArray const& row_group_num_rows)
jintArray const& row_group_num_rows,
jboolean are_retention_vectors)
{
cudf::jni::native_jlongArray n_serialized_roaring64(env, serialized_roaring64);
std::vector<cudf::host_span<cuda::std::byte const>> serialized_bitmaps =
Expand All @@ -97,6 +98,7 @@ std::unique_ptr<cudf::io::parquet::experimental::deletion_vector_info> make_dele
dv_info->deletion_vector_row_counts = n_deletion_vector_row_counts.to_vector();
dv_info->row_group_num_rows = n_row_group_num_rows.to_vector();
dv_info->row_group_offsets.reserve(n_row_group_offsets.size());
dv_info->are_retention_vectors = are_retention_vectors;
std::transform(n_row_group_offsets.begin(),
n_row_group_offsets.end(),
std::back_inserter(dv_info->row_group_offsets),
Expand Down Expand Up @@ -141,7 +143,8 @@ Java_ai_rapids_cudf_DeletionVector_readParquet(JNIEnv* env,
jlongArray serialized_roaring64,
jintArray deletion_vector_row_counts,
jlongArray row_group_offsets,
jintArray row_group_num_rows)
jintArray row_group_num_rows,
jboolean are_retention_vectors)
{
bool read_buffer = true;
if (addrs_and_sizes == nullptr) {
Expand Down Expand Up @@ -175,8 +178,12 @@ Java_ai_rapids_cudf_DeletionVector_readParquet(JNIEnv* env,
cudf::io::parquet_reader_options opts = make_parquet_reader_options(
env, filter_col_names, col_binary_read, row_groups, std::move(source), unit);

auto dv_info = make_deletion_vector_info(
env, serialized_roaring64, deletion_vector_row_counts, row_group_offsets, row_group_num_rows);
auto dv_info = make_deletion_vector_info(env,
serialized_roaring64,
deletion_vector_row_counts,
row_group_offsets,
row_group_num_rows,
are_retention_vectors);

auto tbl = cudf::io::parquet::experimental::read_parquet(opts, *dv_info).tbl;
return cudf::jni::convert_table_for_return(env, tbl);
Expand Down Expand Up @@ -231,7 +238,8 @@ Java_ai_rapids_cudf_DeletionVector_createParquetChunkedReader(JNIEnv* env,
jlongArray serialized_roaring64,
jintArray deletion_vector_row_counts,
jlongArray row_group_offsets,
jintArray row_group_num_rows)
jintArray row_group_num_rows,
jboolean are_retention_vectors)
{
bool read_buffer = true;
if (addrs_sizes == nullptr) {
Expand Down Expand Up @@ -265,8 +273,12 @@ Java_ai_rapids_cudf_DeletionVector_createParquetChunkedReader(JNIEnv* env,
cudf::io::parquet_reader_options opts = make_parquet_reader_options(
env, filter_col_names, col_binary_read, row_groups, std::move(source), unit);

auto dv_info = make_deletion_vector_info(
env, serialized_roaring64, deletion_vector_row_counts, row_group_offsets, row_group_num_rows);
auto dv_info = make_deletion_vector_info(env,
serialized_roaring64,
deletion_vector_row_counts,
row_group_offsets,
row_group_num_rows,
are_retention_vectors);

// Create the chunked reader with pass read limit and multiple deletion vectors
auto reader = new cudf::io::parquet::experimental::chunked_parquet_reader(
Expand Down
58 changes: 45 additions & 13 deletions java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*
*
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
*/
Expand All @@ -11,13 +11,16 @@
import ai.rapids.cudf.DeletionVector.DeletionVectorInfo;
import ai.rapids.cudf.DeletionVector.ParquetChunkedReader;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.CsvSource;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;

import static ai.rapids.cudf.AssertUtils.assertTableTypes;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;


class DeletionVectorTableTest extends CudfTestBase {
Expand All @@ -29,8 +32,9 @@ class DeletionVectorTableTest extends CudfTestBase {
private static final int DELETED_ROWS_COUNT2 = 3959;
private static final int DELETED_ROWS_COUNT2_RGS_1_AND_3 = 1974;

@Test
void testReadParquetReadAllRowGroups() throws IOException {
@ParameterizedTest(name = "isRetention={0}")
@CsvSource({"false", "true"})
void testReadParquetReadAllRowGroups(boolean isRetention) throws IOException {
ParquetOptions opts = ParquetOptions.builder()
.includeColumn("loan_id")
.includeColumn("zip")
Expand All @@ -40,10 +44,12 @@ void testReadParquetReadAllRowGroups() throws IOException {
byte[] bitmapData = TableTestUtils.arrayFrom(DELETED_ROWS_FILE1);
try (HostMemoryBufferArray array = TableTestUtils.buffersFrom(data);
HostMemoryBufferArray bitmapArray = TableTestUtils.buffersFrom(new byte[][] { bitmapData })) {
DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], null, null);
DeletionVectorInfo dvInfo = new DeletionVectorInfo(
bitmapArray.buffers[0], isRetention, null, null);
try (Table table = DeletionVector.readParquet(opts, array.buffers, new DeletionVectorInfo[] { dvInfo })) {
long rows = table.getRowCount();
assertEquals(1000 - DELETED_ROWS_COUNT1, rows);
long expectedRows =
isRetention ? DELETED_ROWS_COUNT1 : 1000 - DELETED_ROWS_COUNT1;
assertEquals(expectedRows, table.getRowCount());
assertTableTypes(new DType[]{DType.UINT64, DType.INT64, DType.INT32, DType.INT32}, table);
}
}
Expand All @@ -59,7 +65,7 @@ void testReadParquetReadSomeRowGroups() throws IOException {
HostMemoryBufferArray bitmapArray = TableTestUtils.buffersFrom(new byte[][] { bitmapData })) {
long[] rowGroupOffsets = Arrays.stream(rowGroups[0]).mapToLong(i -> i * 10000L).toArray();
int[] rowGroupNumRows = Arrays.stream(rowGroups[0]).map(i -> 10000).toArray();
DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], rowGroupOffsets, rowGroupNumRows);
DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], false, rowGroupOffsets, rowGroupNumRows);
try (Table table = DeletionVector.readParquet(opts, array.buffers, rowGroups, new DeletionVectorInfo[] { dvInfo })) {
long rows = table.getRowCount();
assertEquals(20000 - DELETED_ROWS_COUNT2_RGS_1_AND_3, rows);
Expand All @@ -76,7 +82,7 @@ void testChunkedReadParquetAllRowGroups() throws Exception {
try (HostMemoryBufferArray array = TableTestUtils.buffersFrom(data);
HostMemoryBufferArray bitmapArray = TableTestUtils.buffersFrom(new byte[][] { bitmapData })) {
ParquetOptions opts = ParquetOptions.DEFAULT;
DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], null, null);
DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], false, null, null);
try (ParquetChunkedReader reader = DeletionVector.newParquetChunkedReader(240000, 0, opts, array.buffers,
new DeletionVectorInfo[] { dvInfo })) {
int numChunks = 0;
Expand All @@ -94,8 +100,9 @@ void testChunkedReadParquetAllRowGroups() throws Exception {
}
}

@Test
void testChunkedReadParquetSomeRowGroups() throws Exception {
@ParameterizedTest(name = "isRetention={0}")
@CsvSource({"false", "true"})
void testChunkedReadParquetSomeRowGroups(boolean isRetention) throws Exception {
byte[][] data = TableTestUtils.sliceBytes(TableTestUtils.arrayFrom(TEST_FILE2), 2);
byte[] bitmapData = TableTestUtils.arrayFrom(DELETED_ROWS_FILE2);
int[][] rowGroups = new int[][] { {1, 3} };
Expand All @@ -104,7 +111,8 @@ void testChunkedReadParquetSomeRowGroups() throws Exception {
ParquetOptions opts = ParquetOptions.DEFAULT;
long[] rowGroupOffsets = Arrays.stream(rowGroups[0]).mapToLong(i -> i * 10000L).toArray();
int[] rowGroupNumRows = Arrays.stream(rowGroups[0]).map(i -> 10000).toArray();
DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], rowGroupOffsets, rowGroupNumRows);
DeletionVectorInfo dvInfo = new DeletionVectorInfo(
bitmapArray.buffers[0], isRetention, rowGroupOffsets, rowGroupNumRows);
try (ParquetChunkedReader reader = DeletionVector.newParquetChunkedReader(120000, 0, opts, array.buffers,
rowGroups, new DeletionVectorInfo[] { dvInfo })) {
int numChunks = 0;
Expand All @@ -117,7 +125,10 @@ void testChunkedReadParquetSomeRowGroups() throws Exception {
}
}
assertEquals(2, numChunks);
assertEquals(20000 - DELETED_ROWS_COUNT2_RGS_1_AND_3, totalRows);
long expectedRows = isRetention
? DELETED_ROWS_COUNT2_RGS_1_AND_3
: 20000 - DELETED_ROWS_COUNT2_RGS_1_AND_3;
assertEquals(expectedRows, totalRows);
}
}
}
Expand All @@ -132,7 +143,7 @@ void testChunkedReadParquetMultiFiles() throws Exception {
ParquetOptions opts = ParquetOptions.DEFAULT;
long[] rowGroupOffsets = new long[] { 30000L, 10000L };
int[] rowGroupNumRows = new int[] { 10000, 10000 };
DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], rowGroupOffsets, rowGroupNumRows);
DeletionVectorInfo dvInfo = new DeletionVectorInfo(bitmapArray.buffers[0], false, rowGroupOffsets, rowGroupNumRows);
try (ParquetChunkedReader reader = DeletionVector.newParquetChunkedReader(120000, 0, opts, new String[] {
TEST_FILE2.getAbsolutePath(),
TEST_FILE2.getAbsolutePath()
Expand All @@ -152,4 +163,25 @@ void testChunkedReadParquetMultiFiles() throws Exception {
}
}
}

@Test
void testMixedDeletionAndRetentionVectorsRejected() throws IOException {
byte[][] data = TableTestUtils.sliceBytes(TableTestUtils.arrayFrom(TEST_FILE1), 10);
byte[] bitmapData = TableTestUtils.arrayFrom(DELETED_ROWS_FILE1);
try (HostMemoryBufferArray array = TableTestUtils.buffersFrom(data);
HostMemoryBufferArray bitmapArray =
TableTestUtils.buffersFrom(new byte[][] { bitmapData })) {
DeletionVectorInfo[] mixedVectorInfos = new DeletionVectorInfo[] {
new DeletionVectorInfo(bitmapArray.buffers[0], false, null, null),
new DeletionVectorInfo(bitmapArray.buffers[0], true, null, null)
};

assertThrows(IllegalArgumentException.class,
() -> DeletionVector.readParquet(
ParquetOptions.DEFAULT, array.buffers, mixedVectorInfos));
assertThrows(IllegalArgumentException.class,
() -> DeletionVector.newParquetChunkedReader(
0, 0, ParquetOptions.DEFAULT, array.buffers, mixedVectorInfos));
}
}
}
Loading