Add batched computeNumDeletedRows Java API

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.
@@ -109,23 +111,73 @@ 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 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,
+   *     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");
+    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 (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(deletionVectorInfo.serializedBitmap),
-        new int[] {deletionVectorInfo.totalNumRows},
-        deletionVectorInfo.rowGroupOffsets,
-        deletionVectorInfo.rowGroupNumRows,
-        deletionVectorInfo.isRetention,
+        bitmapAddrsSizes,
+        deletionVectorRowCounts,
+        rowGroupOffsets.build().toArray(),
+        rowGroupNumRows.build().toArray(),
+        areRetentionVectors,
         maxChunkRows);
   }
 

java/src/test/java/ai/rapids/cudf/DeletionVectorTableTest.java

@@ -181,6 +181,34 @@ 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 {
+    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));
+    }
+  }
+
+  /**
+   * Verifies invalid row-count arguments are rejected with the expected messages.
+   */
   @Test
   void testComputeNumDeletedRowsInvalidArguments() throws IOException {
     byte[] bitmapData = TableTestUtils.arrayFrom(DELETED_ROWS_FILE1);
@@ -196,8 +224,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 +251,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());
     }
   }