WIP Table.concatenatePacked

java/src/main/java/ai/rapids/cudf/Table.java

@@ -774,6 +774,8 @@ private static native long[] repeatColumnCount(long tableHandle,
 
   private static native long[] columnViewsFromPacked(ByteBuffer metadata, long dataAddress);
 
+  public static native long[] concatenatePacked(ByteBuffer[] metadata, long[] dataAddresses);
+
   private static native ContigSplitGroupByResult contiguousSplitGroups(long inputTable,
                                                                 int[] keyIndices,
                                                                 boolean ignoreNullKeys,
@@ -1944,6 +1946,35 @@ public static Table concatenate(Table... tables) {
     return new Table(concatenate(tableHandles));
   }
 
+  public static Table concatenatePacked(ByteBuffer[] metadata, DeviceMemoryBuffer[] data) {
+    // Ensure the metadata buffer is direct so it can be passed to JNI
+    ByteBuffer[] directBuffers = new ByteBuffer[metadata.length];
+    for (int i = 0; i < metadata.length; i++) {
+      if (!metadata[i].isDirect()) {
+        directBuffers[i]= ByteBuffer.allocateDirect(metadata[i].remaining());
+        directBuffers[i].put(metadata[i]);
+        directBuffers[i].flip();
+      } else {
+        directBuffers[i] = metadata[i];
+      }
+    }
+
+    if (metadata.length < 2) {
+      throw new IllegalArgumentException("concatenate requires 2 or more tables");
+    }
+
+    if (metadata.length != data.length) {
+      throw new IllegalArgumentException("metadata/data arrays length not matching");
+    }
+
+    long[] dataAddressess = new long[data.length];
+    for (int i = 0; i < data.length; i++) {
+      dataAddressess[i] = data[i].address;
+    }
+
+    return new Table(concatenatePacked(directBuffers, dataAddressess));
+  }
+
   /**
    * Interleave all columns into a single column. Columns must all have the same data type and length.
    *
@@ -3663,34 +3694,40 @@ public static Table fromPackedTable(ByteBuffer metadata, DeviceMemoryBuffer data
       directBuffer.flip();
     }
 
-    long[] columnViewAddresses = columnViewsFromPacked(directBuffer, data.getAddress());
-    ColumnVector[] columns = new ColumnVector[columnViewAddresses.length];
-    Table result = null;
-    try {
-      for (int i = 0; i < columns.length; i++) {
-        long columnViewAddress = columnViewAddresses[i];
-        // setting address to zero, so we don't clean it in case of an exception as it
-        // will be cleaned up by the ColumnView  constructor
-        columnViewAddresses[i] = 0;
-        columns[i] = ColumnVector.fromViewWithContiguousAllocation(columnViewAddress, data);
-      }
-      result = new Table(columns);
-    } catch (Throwable t) {
+    long[] columnViewAddresses;
+
+    try(NvtxRange r = new NvtxRange("columnViewsFromPacked", NvtxColor.YELLOW)) {
+      columnViewAddresses = columnViewsFromPacked(directBuffer, data.getAddress());
+    }
+
+    try(NvtxRange r = new NvtxRange("fromViewWithContiguousAllocation", NvtxColor.ORANGE)) {
+      ColumnVector[] columns = new ColumnVector[columnViewAddresses.length];
+      Table result = null;
       try {
-        ColumnView.cleanupColumnViews(columnViewAddresses, columns, t);
-      } catch (Throwable s){
-        t.addSuppressed(s);
-      } finally {
-        throw t;
+        for (int i = 0; i < columns.length; i++) {
+          long columnViewAddress = columnViewAddresses[i];
+          // setting address to zero, so we don't clean it in case of an exception as it
+          // will be cleaned up by the ColumnView  constructor
+          columnViewAddresses[i] = 0;
+          columns[i] = ColumnVector.fromViewWithContiguousAllocation(columnViewAddress, data);
+        }
+        result = new Table(columns);
+      } catch (Throwable t) {
+        try {
+          ColumnView.cleanupColumnViews(columnViewAddresses, columns, t);
+        } catch (Throwable s){
+          t.addSuppressed(s);
+        } finally {
+          throw t;
+        }
       }
-    }
 
-    // close columns to leave the resulting table responsible for freeing underlying columns
-    for (ColumnVector column : columns) {
-      column.close();
+      // close columns to leave the resulting table responsible for freeing underlying columns
+      for (ColumnVector column : columns) {
+        column.close();
+      }
+      return result;
     }
-
-    return result;
   }
 
 

java/src/main/native/src/TableJni.cpp

@@ -1008,6 +1008,29 @@ JNIEXPORT jlongArray JNICALL Java_ai_rapids_cudf_Table_columnViewsFromPacked(JNI
   CATCH_STD(env, nullptr);
 }
 
+JNIEXPORT jlongArray JNICALL Java_ai_rapids_cudf_Table_concatenatePacked(JNIEnv *env, jclass,
+                                                                         jobjectArray j_buffer_objects,
+                                                                         jlongArray j_data_addresses) {
+  // The GPU data address can be null when the table is empty, so it is not null-checked here.
+  try {
+    cudf::jni::auto_set_device(env);
+    int num_buffers = env->GetArrayLength(j_buffer_objects);
+    cudf::jni::native_jlongArray data_addresses(env, j_data_addresses);
+
+    std::vector<cudf::table_view> tables;
+    for (int i = 0; i < num_buffers; i++) {
+      jobject buffer_obj = env->GetObjectArrayElement(j_buffer_objects, i);
+      void const *metadata_address = env->GetDirectBufferAddress(buffer_obj);
+      jlong j_data_address = data_addresses[i];
+      JNI_NULL_CHECK(env, metadata_address, "metadata buffer address is null", nullptr);
+      tables.push_back(cudf::unpack(static_cast<uint8_t const *>(metadata_address),
+                       reinterpret_cast<uint8_t const *>(j_data_address)));
+    }
+    return convert_table_for_return(env, cudf::concatenate(tables));
+  }
+  CATCH_STD(env, nullptr);
+}
+
 JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Table_sortOrder(JNIEnv *env, jclass,
                                                             jlong j_input_table,
                                                             jlongArray j_sort_keys_columns,