From 0abbcd3c859573abb45015ad4b74829cfb57e9d9 Mon Sep 17 00:00:00 2001 From: Rahul Prabhu Date: Tue, 8 Sep 2026 15:13:07 -0700 Subject: [PATCH 1/3] Add batched Java API Signed-off-by: Rahul Prabhu --- .../java/ai/rapids/cudf/DeletionVector.java | 54 ++++++++++++++----- .../rapids/cudf/DeletionVectorTableTest.java | 37 ++++++++++++- 2 files changed, 78 insertions(+), 13 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/DeletionVector.java b/java/src/main/java/ai/rapids/cudf/DeletionVector.java index ea51086f58c3..1dc9dfe9b415 100644 --- a/java/src/main/java/ai/rapids/cudf/DeletionVector.java +++ b/java/src/main/java/ai/rapids/cudf/DeletionVector.java @@ -109,23 +109,53 @@ public static long computeNumDeletedRows( if (deletionVectorInfo == null) { throw new NullPointerException("Expected non-null deletionVectorInfo"); } + return computeNumDeletedRows(new DeletionVectorInfo[] {deletionVectorInfo}, maxChunkRows); + } + + /** + * Computes (on the GPU) the total number of rows deleted by serialized deletion vectors. + * + * @param deletionVectorInfos deletion vectors and row-group metadata + * @param maxChunkRows maximum number of row indexes to process at once + * @return total number of deleted rows in the specified row groups + * @throws NullPointerException if {@code deletionVectorInfos} or one of its elements is null + * @throws IllegalArgumentException if no deletion vectors are supplied, row-group metadata is + * missing, empty, or contains a negative value, deletion and retention vectors are mixed, + * or {@code maxChunkRows} is not positive + */ + public static long computeNumDeletedRows( + DeletionVectorInfo[] deletionVectorInfos, int maxChunkRows) { + if (deletionVectorInfos == null) { + throw new NullPointerException("Expected non-null deletionVectorInfos"); + } + if (deletionVectorInfos.length == 0) { + throw new IllegalArgumentException("deletionVectorInfos must be non-empty"); + } if (maxChunkRows <= 0) { throw new IllegalArgumentException("maxChunkRows must be positive"); } - if (deletionVectorInfo.rowGroupOffsets == null || - deletionVectorInfo.rowGroupOffsets.length == 0) { - throw new IllegalArgumentException("row-group metadata must be non-empty"); - } - if (Arrays.stream(deletionVectorInfo.rowGroupOffsets).anyMatch(value -> value < 0) || - Arrays.stream(deletionVectorInfo.rowGroupNumRows).anyMatch(value -> value < 0)) { - throw new IllegalArgumentException("row-group metadata values must be non-negative"); + for (DeletionVectorInfo info : deletionVectorInfos) { + if (info == null) { + throw new NullPointerException("Expected non-null deletionVectorInfo"); + } + if (info.rowGroupOffsets == null || info.rowGroupOffsets.length == 0) { + throw new IllegalArgumentException("row-group metadata must be non-empty"); + } + if (Arrays.stream(info.rowGroupOffsets).anyMatch(value -> value < 0) || + Arrays.stream(info.rowGroupNumRows).anyMatch(value -> value < 0)) { + throw new IllegalArgumentException("row-group metadata values must be non-negative"); + } } + return computeNumDeletedRows( - getAddrsAndSizes(deletionVectorInfo.serializedBitmap), - new int[] {deletionVectorInfo.totalNumRows}, - deletionVectorInfo.rowGroupOffsets, - deletionVectorInfo.rowGroupNumRows, - deletionVectorInfo.isRetention, + getAddrsAndSizes(Arrays.stream(deletionVectorInfos) + .map(info -> info.serializedBitmap).toArray(HostMemoryBuffer[]::new)), + Arrays.stream(deletionVectorInfos).mapToInt(info -> info.totalNumRows).toArray(), + Arrays.stream(deletionVectorInfos) + .flatMapToLong(info -> Arrays.stream(info.rowGroupOffsets)).toArray(), + Arrays.stream(deletionVectorInfos) + .flatMapToInt(info -> Arrays.stream(info.rowGroupNumRows)).toArray(), + getDeletionVectorTypes(deletionVectorInfos), maxChunkRows); } diff --git a/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java b/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java index e9914f094055..11bdea63b229 100644 --- a/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java +++ b/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java @@ -181,6 +181,26 @@ void testComputeNumDeletedRows(boolean isRetention) throws IOException { } } + @ParameterizedTest(name = "isRetention={0}") + @CsvSource({"false", "true"}) + void testComputeNumDeletedRowsBatch(boolean isRetention) throws IOException { + byte[] bitmapData = TableTestUtils.arrayFrom(DELETED_ROWS_FILE2); + long[] rowGroupOffsets = new long[] {10000L, 30000L}; + int[] rowGroupNumRows = new int[] {10000, 10000}; + try (HostMemoryBufferArray bitmapArray = + TableTestUtils.buffersFrom(new byte[][] {bitmapData, bitmapData})) { + DeletionVectorInfo[] dvInfos = Arrays.stream(bitmapArray.buffers) + .map(bitmap -> new DeletionVectorInfo( + bitmap, isRetention, rowGroupOffsets, rowGroupNumRows)) + .toArray(DeletionVectorInfo[]::new); + long expectedRowsDeletedPerVector = isRetention + ? 20000 - DELETED_ROWS_COUNT2_RGS_1_AND_3 + : DELETED_ROWS_COUNT2_RGS_1_AND_3; + assertEquals(2 * expectedRowsDeletedPerVector, + DeletionVector.computeNumDeletedRows(dvInfos, 5000)); + } + } + @Test void testComputeNumDeletedRowsInvalidArguments() throws IOException { byte[] bitmapData = TableTestUtils.arrayFrom(DELETED_ROWS_FILE1); @@ -196,8 +216,19 @@ void testComputeNumDeletedRowsInvalidArguments() throws IOException { bitmapArray.buffers[0], false, new long[] {-1}, new int[] {1000}); DeletionVectorInfo negativeRowCount = new DeletionVectorInfo( bitmapArray.buffers[0], false, new long[] {0}, new int[] {-1}); + DeletionVectorInfo retentionInfo = new DeletionVectorInfo( + bitmapArray.buffers[0], true, new long[] {0}, new int[] {1000}); assertEquals("Expected non-null deletionVectorInfo", assertThrows(NullPointerException.class, - () -> DeletionVector.computeNumDeletedRows(null, 1000)).getMessage()); + () -> DeletionVector.computeNumDeletedRows( + (DeletionVectorInfo) null, 1000)).getMessage()); + assertEquals("Expected non-null deletionVectorInfos", + assertThrows(NullPointerException.class, + () -> DeletionVector.computeNumDeletedRows( + (DeletionVectorInfo[]) null, 1000)).getMessage()); + assertEquals("deletionVectorInfos must be non-empty", + assertThrows(IllegalArgumentException.class, + () -> DeletionVector.computeNumDeletedRows( + new DeletionVectorInfo[0], 1000)).getMessage()); assertEquals("maxChunkRows must be positive", assertThrows(IllegalArgumentException.class, () -> DeletionVector.computeNumDeletedRows(dvInfo, 0)).getMessage()); assertEquals("row-group metadata must be non-empty", @@ -212,6 +243,10 @@ void testComputeNumDeletedRowsInvalidArguments() throws IOException { assertEquals("row-group metadata values must be non-negative", assertThrows(IllegalArgumentException.class, () -> DeletionVector.computeNumDeletedRows(negativeRowCount, 1000)).getMessage()); + assertEquals("All DeletionVectorInfo objects must have the same isRetention value.", + assertThrows(IllegalArgumentException.class, + () -> DeletionVector.computeNumDeletedRows( + new DeletionVectorInfo[] {dvInfo, retentionInfo}, 1000)).getMessage()); } } From 2a6e152f74b9643d288f9fd1bdfb1d2bc67179d0 Mon Sep 17 00:00:00 2001 From: Rahul Prabhu Date: Wed, 9 Sep 2026 09:30:21 -0700 Subject: [PATCH 2/3] Added docstrings --- .../test/java/ai/rapids/cudf/DeletionVectorTableTest.java | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java b/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java index 11bdea63b229..6caaf593f546 100644 --- a/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java +++ b/java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java @@ -181,6 +181,11 @@ void testComputeNumDeletedRows(boolean isRetention) throws IOException { } } + /** + * Verifies batched row counting for chunked deletion and retention vectors. + * + * @param isRetention whether the input bitmaps identify retained rows + */ @ParameterizedTest(name = "isRetention={0}") @CsvSource({"false", "true"}) void testComputeNumDeletedRowsBatch(boolean isRetention) throws IOException { @@ -201,6 +206,9 @@ void testComputeNumDeletedRowsBatch(boolean isRetention) throws IOException { } } + /** + * Verifies invalid row-count arguments are rejected with the expected messages. + */ @Test void testComputeNumDeletedRowsInvalidArguments() throws IOException { byte[] bitmapData = TableTestUtils.arrayFrom(DELETED_ROWS_FILE1); From 567fe39e8c29af2b025ac666ab3ebbcf66a4e6b3 Mon Sep 17 00:00:00 2001 From: Rahul Prabhu Date: Wed, 9 Sep 2026 13:26:25 -0700 Subject: [PATCH 3/3] iterate dv infos once --- .../java/ai/rapids/cudf/DeletionVector.java | 48 ++++++++++++++----- 1 file changed, 35 insertions(+), 13 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/DeletionVector.java b/java/src/main/java/ai/rapids/cudf/DeletionVector.java index 1dc9dfe9b415..748b1dc24b03 100644 --- a/java/src/main/java/ai/rapids/cudf/DeletionVector.java +++ b/java/src/main/java/ai/rapids/cudf/DeletionVector.java @@ -10,6 +10,8 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.stream.IntStream; +import java.util.stream.LongStream; /** * Provides JNI wrappers for reading Parquet files with deletion vector support. @@ -117,7 +119,7 @@ public static long computeNumDeletedRows( * * @param deletionVectorInfos deletion vectors and row-group metadata * @param maxChunkRows maximum number of row indexes to process at once - * @return total number of deleted rows in the specified row groups + * @return total number of deleted rows in the specified row groups across all deletion vectors. * @throws NullPointerException if {@code deletionVectorInfos} or one of its elements is null * @throws IllegalArgumentException if no deletion vectors are supplied, row-group metadata is * missing, empty, or contains a negative value, deletion and retention vectors are mixed, @@ -134,28 +136,48 @@ public static long computeNumDeletedRows( if (maxChunkRows <= 0) { throw new IllegalArgumentException("maxChunkRows must be positive"); } - for (DeletionVectorInfo info : deletionVectorInfos) { + long[] bitmapAddrsSizes = new long[deletionVectorInfos.length * 2]; + int[] deletionVectorRowCounts = new int[deletionVectorInfos.length]; + LongStream.Builder rowGroupOffsets = LongStream.builder(); + IntStream.Builder rowGroupNumRows = IntStream.builder(); + boolean areRetentionVectors = false; + for (int i = 0; i < deletionVectorInfos.length; i++) { + DeletionVectorInfo info = deletionVectorInfos[i]; if (info == null) { throw new NullPointerException("Expected non-null deletionVectorInfo"); } if (info.rowGroupOffsets == null || info.rowGroupOffsets.length == 0) { throw new IllegalArgumentException("row-group metadata must be non-empty"); } - if (Arrays.stream(info.rowGroupOffsets).anyMatch(value -> value < 0) || - Arrays.stream(info.rowGroupNumRows).anyMatch(value -> value < 0)) { - throw new IllegalArgumentException("row-group metadata values must be non-negative"); + if (i == 0) { + areRetentionVectors = info.isRetention; + } else if (info.isRetention != areRetentionVectors) { + throw new IllegalArgumentException( + "All DeletionVectorInfo objects must have the same isRetention value."); + } + bitmapAddrsSizes[i * 2] = info.serializedBitmap.getAddress(); + bitmapAddrsSizes[(i * 2) + 1] = info.serializedBitmap.getLength(); + deletionVectorRowCounts[i] = info.totalNumRows; + for (int rowGroupIndex = 0; + rowGroupIndex < info.rowGroupOffsets.length; + rowGroupIndex++) { + long offset = info.rowGroupOffsets[rowGroupIndex]; + int numRows = info.rowGroupNumRows[rowGroupIndex]; + if (offset < 0 || numRows < 0) { + throw new IllegalArgumentException( + "row-group metadata values must be non-negative"); + } + rowGroupOffsets.add(offset); + rowGroupNumRows.add(numRows); } } return computeNumDeletedRows( - getAddrsAndSizes(Arrays.stream(deletionVectorInfos) - .map(info -> info.serializedBitmap).toArray(HostMemoryBuffer[]::new)), - Arrays.stream(deletionVectorInfos).mapToInt(info -> info.totalNumRows).toArray(), - Arrays.stream(deletionVectorInfos) - .flatMapToLong(info -> Arrays.stream(info.rowGroupOffsets)).toArray(), - Arrays.stream(deletionVectorInfos) - .flatMapToInt(info -> Arrays.stream(info.rowGroupNumRows)).toArray(), - getDeletionVectorTypes(deletionVectorInfos), + bitmapAddrsSizes, + deletionVectorRowCounts, + rowGroupOffsets.build().toArray(), + rowGroupNumRows.build().toArray(), + areRetentionVectors, maxChunkRows); }