From 3fe1140fcaa9016de938231e94a23259156bbdcd Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Tue, 25 Apr 2023 00:41:08 -0500 Subject: [PATCH 1/8] JNI api for cudf::chunked_pack Signed-off-by: Alessandro Bellina --- .../main/java/ai/rapids/cudf/ChunkedPack.java | 99 +++++++++++++++++++ .../java/ai/rapids/cudf/ContiguousTable.java | 25 ++--- .../ai/rapids/cudf/PackedColumnMetadata.java | 74 ++++++++++++++ java/src/main/java/ai/rapids/cudf/Table.java | 40 ++++++++ java/src/main/native/CMakeLists.txt | 2 + java/src/main/native/src/ChunkedPackJni.cpp | 73 ++++++++++++++ .../main/native/src/ContiguousTableJni.cpp | 20 ---- java/src/main/native/src/CudfJni.cpp | 11 ++- .../native/src/PackedColumnMetadataJni.cpp | 90 +++++++++++++++++ java/src/main/native/src/TableJni.cpp | 18 ++++ java/src/main/native/src/cudf_jni_apis.hpp | 5 +- .../test/java/ai/rapids/cudf/TableTest.java | 76 ++++++++++++++ 12 files changed, 494 insertions(+), 39 deletions(-) create mode 100644 java/src/main/java/ai/rapids/cudf/ChunkedPack.java create mode 100644 java/src/main/java/ai/rapids/cudf/PackedColumnMetadata.java create mode 100644 java/src/main/native/src/ChunkedPackJni.cpp create mode 100644 java/src/main/native/src/PackedColumnMetadataJni.cpp diff --git a/java/src/main/java/ai/rapids/cudf/ChunkedPack.java b/java/src/main/java/ai/rapids/cudf/ChunkedPack.java new file mode 100644 index 00000000000..291e6e41ce7 --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/ChunkedPack.java @@ -0,0 +1,99 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package ai.rapids.cudf; + +/** + * JNI interface to cudf::chunked_pack. + * + * ChunkedPack is an iterator-like interface with the familiar `hasNext` and `next` + * interface. `next` should be used in a loop until `hasNext` returns false. + * + * However, `ChunkedPack.next` is special because it takes a `DeviceMemoryBuffer` as a + * parameter, which means that the caller can call `next` giving any bounce buffer it + * may have previously allocated. No requirement exists that the bounce buffer be the + * same each time, the only requirement is that their sizes are all the same, and match + * the size that was passed to `Table.makeChunkedPack` (which instantiates this class). + * + * The user of `ChunkedPack` must close `.close()` when done using it to clear up both + * host and device resources. + */ +public class ChunkedPack implements AutoCloseable { + long nativePtr; + + /** + * This constructor is invoked by `Table.makeChunkedPack` after creating a native + * `cudf::chunked_pack`. + * @param nativePtr pointer to a `cudf::chunked_pack` + */ + public ChunkedPack(long nativePtr) { + this.nativePtr = nativePtr; + } + + /** + * Get the final contiguous size of the table we are packing. This is + * the size that the final buffer should be, just like if the user called + * `cudf::pack` instead. + * @return the total number of bytes for the table in contigous layout + */ + public long getTotalContiguousSize() { + return chunkedPackGetTotalContiguousSize(nativePtr); + } + + /** + * Method to be called to ensure that `ChunkedPack` has work left. + * This method should be invoked followed by a call to `next`, until + * `hasNext` returns false. + * @return true if there is work left to be done (`next` should be called), + * false otherwise. + */ + public boolean hasNext() { + return chunkedPackHasNext(nativePtr); + } + + /** + * Place the next contigous chunk of our table into `userPtr`. + * + * This method throws if `hasNext` is false. + * @param userPtr the bounce buffer to use for this iteration + * @return the number of bytes that we were able to place in `userPtr`. This is + * at most `userPtr.getLength()`. + */ + public long next(DeviceMemoryBuffer userPtr) { + return chunkedPackNext(nativePtr, userPtr.getAddress(), userPtr.getLength()); + } + + /** + * Generates opaque table metadata that can be unpacked via `cudf::unpack` + * at a later time. + * @return a `PackedColumnMetadata` instance referencing cuDF packed table metadata + */ + public PackedColumnMetadata buildMetadata() { + return chunkedPackBuildMetadata(nativePtr); + } + + @Override + public void close() { + chunkedPackDelete(nativePtr); + } + + private static native long chunkedPackGetTotalContiguousSize(long nativePtr); + private static native boolean chunkedPackHasNext(long nativePtr); + private static native long chunkedPackNext(long nativePtr, long userPtr, long userPtrSize); + private static native PackedColumnMetadata chunkedPackBuildMetadata(long nativePtr); + private static native void chunkedPackDelete(long nativePtr); +} diff --git a/java/src/main/java/ai/rapids/cudf/ContiguousTable.java b/java/src/main/java/ai/rapids/cudf/ContiguousTable.java index 87a3f5f0ddf..a1aadebd1bd 100644 --- a/java/src/main/java/ai/rapids/cudf/ContiguousTable.java +++ b/java/src/main/java/ai/rapids/cudf/ContiguousTable.java @@ -25,11 +25,11 @@ * much simpler. */ public final class ContiguousTable implements AutoCloseable { - private long metadataHandle = 0; private Table table = null; private DeviceMemoryBuffer buffer; - private ByteBuffer metadataBuffer = null; private final long rowCount; + private PackedColumnMetadata meta; + private ByteBuffer metadataBuffer; // This method is invoked by JNI static ContiguousTable fromPackedTable(long metadataHandle, @@ -43,8 +43,8 @@ static ContiguousTable fromPackedTable(long metadataHandle, /** Construct a contiguous table instance given a table and the device buffer backing it. */ ContiguousTable(Table table, DeviceMemoryBuffer buffer) { - this.metadataHandle = createPackedMetadata(table.getNativeView(), - buffer.getAddress(), buffer.getLength()); + this.meta = new PackedColumnMetadata(createPackedMetadata(table.getNativeView(), + buffer.getAddress(), buffer.getLength())); this.table = table; this.buffer = buffer; this.rowCount = table.getRowCount(); @@ -57,7 +57,7 @@ static ContiguousTable fromPackedTable(long metadataHandle, * @param rowCount number of rows in the table */ ContiguousTable(long metadataHandle, DeviceMemoryBuffer buffer, long rowCount) { - this.metadataHandle = metadataHandle; + this.meta = new PackedColumnMetadata(metadataHandle); this.buffer = buffer; this.rowCount = rowCount; } @@ -94,18 +94,14 @@ public DeviceMemoryBuffer getBuffer() { * or data corruption. */ public ByteBuffer getMetadataDirectBuffer() { - if (metadataBuffer == null) { - metadataBuffer = createMetadataDirectBuffer(metadataHandle); - } - return metadataBuffer.asReadOnlyBuffer(); + return meta.getMetadataDirectBuffer(); } /** Close the contiguous table instance and its underlying resources. */ @Override public void close() { - if (metadataHandle != 0) { - closeMetadata(metadataHandle); - metadataHandle = 0; + if (meta != null) { + meta.close(); } if (table != null) { @@ -122,9 +118,4 @@ public void close() { // create packed metadata for a table backed by a single data buffer private static native long createPackedMetadata(long tableView, long dataAddress, long dataSize); - // create a DirectByteBuffer for the packed table metadata - private static native ByteBuffer createMetadataDirectBuffer(long metadataHandle); - - // release the native metadata resources for a packed table - private static native void closeMetadata(long metadataHandle); } diff --git a/java/src/main/java/ai/rapids/cudf/PackedColumnMetadata.java b/java/src/main/java/ai/rapids/cudf/PackedColumnMetadata.java new file mode 100644 index 00000000000..1b173e7e2aa --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/PackedColumnMetadata.java @@ -0,0 +1,74 @@ +/* + * + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package ai.rapids.cudf; + +import java.nio.ByteBuffer; + +/** + * Metadata for a table that is backed by a single contiguous device buffer. + */ +public final class PackedColumnMetadata implements AutoCloseable { + private long metadataHandle = 0; + private ByteBuffer metadataBuffer = null; + + // This method is invoked by JNI + static PackedColumnMetadata fromPackedColumnMeta(long metadataHandle) { + return new PackedColumnMetadata(metadataHandle); + } + + /** + * Construct the PackedColumnMetadata instance given a metadata handle. + * @param metadataHandle address of the cudf packed_table host-based metadata instance + */ + PackedColumnMetadata(long metadataHandle) { + this.metadataHandle = metadataHandle; + } + + /** + * Get the byte buffer containing the host metadata describing the schema and layout of the + * contiguous table. + *

+ * NOTE: This is a direct byte buffer that is backed by the underlying native metadata instance + * and therefore is only valid to be used while this PackedColumnMetadata instance is valid. + * Attempts to cache and access the resulting buffer after this instance has been destroyed + * will result in undefined behavior including the possibility of segmentation faults + * or data corruption. + */ + public ByteBuffer getMetadataDirectBuffer() { + if (metadataBuffer == null) { + metadataBuffer = createMetadataDirectBuffer(metadataHandle); + } + return metadataBuffer.asReadOnlyBuffer(); + } + + /** Close the PackedColumnMetadata instance and its underlying resources. */ + @Override + public void close() { + if (metadataHandle != 0) { + closeMetadata(metadataHandle); + metadataHandle = 0; + } + } + + // create a DirectByteBuffer for the packed metadata + private static native ByteBuffer createMetadataDirectBuffer(long metadataHandle); + + // release the native metadata resources for a packed table + private static native void closeMetadata(long metadataHandle); +} diff --git a/java/src/main/java/ai/rapids/cudf/Table.java b/java/src/main/java/ai/rapids/cudf/Table.java index 9abc2dbcd7c..9ce383d0a9f 100644 --- a/java/src/main/java/ai/rapids/cudf/Table.java +++ b/java/src/main/java/ai/rapids/cudf/Table.java @@ -187,6 +187,8 @@ static Table removeNullMasksIfNeeded(Table table) { private static native ContiguousTable[] contiguousSplit(long inputTable, int[] indices); + private static native long makeChunkedPack(long inputTable, long bounceBufferSize, long memoryResource); + private static native long[] partition(long inputTable, long partitionView, int numberOfPartitions, int[] outputOffsets); @@ -2171,6 +2173,44 @@ public ContiguousTable[] contiguousSplit(int... indices) { return contiguousSplit(nativeHandle, indices); } + /** + * Create an instance of `ChunkedPack` which can be used to pack this table + * contiguously in memory utilizing a bounce buffer of size `bounceBufferSize`. + * + * This version of `makeChunkedPack` takes a `RmmDviceMemoryResource`, which can be used + * to pre-allocate all scratch and temporary space required for the state of `cudf::chunked_pack`. + * + * The caller is responsible for calling close on the returned `ChunkedPack` object. + * + * @param bounceBufferSize The size of bounce buffer that will be utilized to pack into + * @param deviceMemoryResource A memory resource that is used to satisfy allocations for + * temporary and thrust scratch space. + * @return An instance of `ChunkedPack` that the caller must use to finish the operation. + */ + public ChunkedPack makeChunkedPack( + long bounceBufferSize, RmmDeviceMemoryResource deviceMemoryResource) { + long memoryResourceNativeHandle = deviceMemoryResource.getHandle(); + return new ChunkedPack( + makeChunkedPack(nativeHandle, bounceBufferSize, memoryResourceNativeHandle)); + } + + /** + * Create an instance of `ChunkedPack` which can be used to pack this table + * contiguously in memory utilizing a bounce buffer of size `bounceBufferSize`. + * + * This version of `makeChunkedPack` makes use of the default per-device memory resource, + * for scratch and temporary space required for the state of `cudf::chunked_pack`. + * + * The caller is responsible for calling close on the returned `ChunkedPack` object. + * + * @param bounceBufferSize The size of bounce buffer that will be utilized to pack into + * @return An instance of `ChunkedPack` that the caller must use to finish the operation. + */ + public ChunkedPack makeChunkedPack(long bounceBufferSize) { + return new ChunkedPack( + makeChunkedPack(nativeHandle, bounceBufferSize, 0)); + } + /** * Explodes a list column's elements. * diff --git a/java/src/main/native/CMakeLists.txt b/java/src/main/native/CMakeLists.txt index 57372bf7c09..37b7194258d 100644 --- a/java/src/main/native/CMakeLists.txt +++ b/java/src/main/native/CMakeLists.txt @@ -126,6 +126,7 @@ add_library( cudfjni src/Aggregation128UtilsJni.cpp src/AggregationJni.cpp + src/ChunkedPackJni.cpp src/ChunkedReaderJni.cpp src/CudfJni.cpp src/CudaJni.cpp @@ -139,6 +140,7 @@ add_library( src/NvcompJni.cpp src/NvtxRangeJni.cpp src/NvtxUniqueRangeJni.cpp + src/PackedColumnMetadataJni.cpp src/RmmJni.cpp src/ScalarJni.cpp src/TableJni.cpp diff --git a/java/src/main/native/src/ChunkedPackJni.cpp b/java/src/main/native/src/ChunkedPackJni.cpp new file mode 100644 index 00000000000..b8d6211c590 --- /dev/null +++ b/java/src/main/native/src/ChunkedPackJni.cpp @@ -0,0 +1,73 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cudf_jni_apis.hpp" + +extern "C" { +JNIEXPORT void JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackDelete( + JNIEnv *env, jclass, jlong chunked_pack) { + try { + cudf::jni::auto_set_device(env); + auto cs = reinterpret_cast(chunked_pack); + delete cs; + } + CATCH_STD(env, ); +} + +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackGetTotalContiguousSize( + JNIEnv *env, jclass, jlong chunked_pack) { + try { + cudf::jni::auto_set_device(env); + auto cs = reinterpret_cast(chunked_pack); + return cs->get_total_contiguous_size(); + } + CATCH_STD(env, 0); +} + +JNIEXPORT jboolean JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackHasNext( + JNIEnv *env, jclass, jlong chunked_pack) { + try { + cudf::jni::auto_set_device(env); + auto cs = reinterpret_cast(chunked_pack); + return cs->has_next(); + } + CATCH_STD(env, 0); +} + +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackNext( + JNIEnv *env, jclass, jlong chunked_pack, jlong user_ptr, jlong user_ptr_size) { + try { + cudf::jni::auto_set_device(env); + auto cs = reinterpret_cast(chunked_pack); + auto user_buffer_span = cudf::device_span(reinterpret_cast(user_ptr), + static_cast(user_ptr_size)); + return cs->next(user_buffer_span); + } + CATCH_STD(env, 0); +} + +JNIEXPORT jobject JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackBuildMetadata( + JNIEnv *env, jclass, jlong chunked_pack) { + try { + cudf::jni::auto_set_device(env); + auto cs = reinterpret_cast(chunked_pack); + std::unique_ptr> result = cs->build_metadata(); + return cudf::jni::packed_column_metadata_from(env, std::move(result)); + } + CATCH_STD(env, NULL); +} + +} // extern "C" diff --git a/java/src/main/native/src/ContiguousTableJni.cpp b/java/src/main/native/src/ContiguousTableJni.cpp index 85a5a24262e..7eddea2a895 100644 --- a/java/src/main/native/src/ContiguousTableJni.cpp +++ b/java/src/main/native/src/ContiguousTableJni.cpp @@ -149,24 +149,4 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ContiguousTable_createPackedMetadata CATCH_STD(env, 0); } -JNIEXPORT jobject JNICALL Java_ai_rapids_cudf_ContiguousTable_createMetadataDirectBuffer( - JNIEnv *env, jclass, jlong j_metadata_ptr) { - JNI_NULL_CHECK(env, j_metadata_ptr, "metadata is null", nullptr); - try { - auto metadata = reinterpret_cast *>(j_metadata_ptr); - return env->NewDirectByteBuffer(const_cast(metadata->data()), metadata->size()); - } - CATCH_STD(env, nullptr); -} - -JNIEXPORT void JNICALL Java_ai_rapids_cudf_ContiguousTable_closeMetadata(JNIEnv *env, jclass, - jlong j_metadata_ptr) { - JNI_NULL_CHECK(env, j_metadata_ptr, "metadata is null", ); - try { - auto metadata = reinterpret_cast *>(j_metadata_ptr); - delete metadata; - } - CATCH_STD(env, ); -} - } // extern "C" diff --git a/java/src/main/native/src/CudfJni.cpp b/java/src/main/native/src/CudfJni.cpp index 1fdadb4b7e2..66f491f51fd 100644 --- a/java/src/main/native/src/CudfJni.cpp +++ b/java/src/main/native/src/CudfJni.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2022, NVIDIA CORPORATION. + * Copyright (c) 2019-2023, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -164,6 +164,14 @@ JNIEXPORT jint JNI_OnLoad(JavaVM *vm, void *) { return JNI_ERR; } + if (!cudf::jni::cache_packed_column_meta_jni(env)) { + if (!env->ExceptionCheck()) { + env->ThrowNew(env->FindClass("java/lang/RuntimeException"), + "Unable to locate packed column metadata methods needed by JNI"); + } + return JNI_ERR; + } + if (!cudf::jni::cache_contig_split_group_by_result_jni(env)) { if (!env->ExceptionCheck()) { env->ThrowNew(env->FindClass("java/lang/RuntimeException"), @@ -193,6 +201,7 @@ JNIEXPORT void JNI_OnUnload(JavaVM *vm, void *) { cudf::jni::release_contiguous_table_jni(env); cudf::jni::release_contig_split_group_by_result_jni(env); cudf::jni::release_host_memory_buffer_jni(env); + cudf::jni::release_packed_column_meta_jni(env); } JNIEXPORT jboolean JNICALL Java_ai_rapids_cudf_Cuda_isPtdsEnabled(JNIEnv *env, jclass, jlong size) { diff --git a/java/src/main/native/src/PackedColumnMetadataJni.cpp b/java/src/main/native/src/PackedColumnMetadataJni.cpp new file mode 100644 index 00000000000..0047652be67 --- /dev/null +++ b/java/src/main/native/src/PackedColumnMetadataJni.cpp @@ -0,0 +1,90 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cudf_jni_apis.hpp" + +namespace { + +#define PACKED_COLUMN_META_CLASS "ai/rapids/cudf/PackedColumnMetadata" +#define PACKED_COLUMN_META_FACTORY_SIG(param_sig) "(" param_sig ")L" PACKED_COLUMN_META_CLASS ";" + +jclass Packed_columns_meta_jclass; +jmethodID From_packed_column_meta_method; + +} // anonymous namespace + +namespace cudf { +namespace jni { + +bool cache_packed_column_meta_jni(JNIEnv *env) { + jclass cls = env->FindClass(PACKED_COLUMN_META_CLASS); + if (cls == nullptr) { + return false; + } + + From_packed_column_meta_method = + env->GetStaticMethodID(cls, "fromPackedColumnMeta", PACKED_COLUMN_META_FACTORY_SIG("J")); + if (From_packed_column_meta_method == nullptr) { + return false; + } + + // Convert local reference to global so it cannot be garbage collected. + Packed_columns_meta_jclass = static_cast(env->NewGlobalRef(cls)); + if (Packed_columns_meta_jclass == nullptr) { + return false; + } + return true; +} + +jobject packed_column_metadata_from(JNIEnv *env, std::unique_ptr> meta) { + jlong metadata_address = reinterpret_cast(meta.release()); + return env->CallStaticObjectMethod(Packed_columns_meta_jclass, From_packed_column_meta_method, + metadata_address); +} + +void release_packed_column_meta_jni(JNIEnv *env) { + if (Packed_columns_meta_jclass != nullptr) { + env->DeleteGlobalRef(Packed_columns_meta_jclass); + Packed_columns_meta_jclass = nullptr; + } +} + +} // namespace jni +} // namespace cudf + +extern "C" { + +JNIEXPORT jobject JNICALL Java_ai_rapids_cudf_PackedColumnMetadata_createMetadataDirectBuffer( + JNIEnv *env, jclass, jlong j_metadata_ptr) { + JNI_NULL_CHECK(env, j_metadata_ptr, "metadata is null", nullptr); + try { + auto metadata = reinterpret_cast *>(j_metadata_ptr); + return env->NewDirectByteBuffer(const_cast(metadata->data()), metadata->size()); + } + CATCH_STD(env, nullptr); +} + +JNIEXPORT void JNICALL +Java_ai_rapids_cudf_PackedColumnMetadata_closeMetadata(JNIEnv *env, jclass, jlong j_metadata_ptr) { + JNI_NULL_CHECK(env, j_metadata_ptr, "metadata is null", ); + try { + auto metadata = reinterpret_cast *>(j_metadata_ptr); + delete metadata; + } + CATCH_STD(env, ); +} + +} // extern "C" diff --git a/java/src/main/native/src/TableJni.cpp b/java/src/main/native/src/TableJni.cpp index 51c2e92492f..ba0f4af6175 100644 --- a/java/src/main/native/src/TableJni.cpp +++ b/java/src/main/native/src/TableJni.cpp @@ -45,6 +45,7 @@ #include #include #include +#include #include #include "csv_chunked_writer.hpp" @@ -3158,6 +3159,23 @@ JNIEXPORT jobjectArray JNICALL Java_ai_rapids_cudf_Table_contiguousSplit(JNIEnv CATCH_STD(env, NULL); } +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Table_makeChunkedPack( + JNIEnv *env, jclass, jlong input_table, jlong bounce_buffer_size, jlong memoryResourceHandle) { + JNI_NULL_CHECK(env, input_table, "native handle is null", 0); + + try { + cudf::jni::auto_set_device(env); + cudf::table_view *n_table = reinterpret_cast(input_table); + auto mr = memoryResourceHandle != 0 ? + reinterpret_cast(memoryResourceHandle) : + rmm::mr::get_current_device_resource(); + auto chunked_contig_split = + cudf::make_chunked_pack(*n_table, bounce_buffer_size, mr); + return reinterpret_cast(chunked_contig_split.release()); + } + CATCH_STD(env, 0); +} + JNIEXPORT jlongArray JNICALL Java_ai_rapids_cudf_Table_rollingWindowAggregate( JNIEnv *env, jclass, jlong j_input_table, jintArray j_keys, jlongArray j_default_output, jintArray j_aggregate_column_indices, jlongArray j_agg_instances, jintArray j_min_periods, diff --git a/java/src/main/native/src/cudf_jni_apis.hpp b/java/src/main/native/src/cudf_jni_apis.hpp index 18993aea294..2005fdfc446 100644 --- a/java/src/main/native/src/cudf_jni_apis.hpp +++ b/java/src/main/native/src/cudf_jni_apis.hpp @@ -50,10 +50,13 @@ convert_table_for_return(JNIEnv *env, std::unique_ptr &&table_resul // ContiguousTable APIs // -bool cache_contiguous_table_jni(JNIEnv *env); +bool cache_packed_column_meta_jni(JNIEnv *env); +void release_packed_column_meta_jni(JNIEnv *env); +bool cache_contiguous_table_jni(JNIEnv *env); void release_contiguous_table_jni(JNIEnv *env); +jobject packed_column_metadata_from(JNIEnv *env, std::unique_ptr> meta); jobject contiguous_table_from(JNIEnv *env, cudf::packed_columns &split, long row_count); native_jobjectArray contiguous_table_array(JNIEnv *env, jsize length); diff --git a/java/src/test/java/ai/rapids/cudf/TableTest.java b/java/src/test/java/ai/rapids/cudf/TableTest.java index 6b03079fa81..8ccda6df079 100644 --- a/java/src/test/java/ai/rapids/cudf/TableTest.java +++ b/java/src/test/java/ai/rapids/cudf/TableTest.java @@ -3099,6 +3099,58 @@ void testContiguousSplit() { } } + @Test + void testChunkedPackBasic() { + try (Table t1 = new Table.TestBuilder() + .column(10, 12, 14, 16, 18, 20, 22, 24, null, 28) + .column(50, 52, 54, 56, 58, 60, 62, 64, 66, null) + .decimal32Column(-3, 10, 12, 14, 16, 18, 20, 22, 24, null, 28) + .decimal64Column(-8, 50L, 52L, 54L, 56L, 58L, 60L, 62L, 64L, 66L, null) + .build(); + DeviceMemoryBuffer bounceBuffer = DeviceMemoryBuffer.allocate(10L*1024*1024); + ChunkedPack cp = t1.makeChunkedPack(10L*1024*1024); + PackedColumnMetadata meta = cp.buildMetadata()) { + + // unpack to bounce buffer + assertEquals(true, cp.hasNext()); + assertEquals(cp.getTotalContiguousSize(), cp.next(bounceBuffer)); + assertEquals(false, cp.hasNext()); + + try (Table unpacked = Table.fromPackedTable(meta.getMetadataDirectBuffer(), bounceBuffer)) { + assertTablesAreEqual(t1, unpacked); + } + } + } + + @Test + void testChunkedPackTwoPasses() { + // this test packes ~2MB worth of long into a 1MB bounce buffer + // this is 3 iterations because of the validity buffer + Long[] longs = new Long[256*1024]; + try (Table t1 = new Table.TestBuilder().column(longs).build(); + DeviceMemoryBuffer bounceBuffer = DeviceMemoryBuffer.allocate(1L*1024*1024); + ChunkedPack cp = t1.makeChunkedPack(1L*1024*1024); + PackedColumnMetadata meta = cp.buildMetadata(); + DeviceMemoryBuffer target = DeviceMemoryBuffer.allocate(cp.getTotalContiguousSize())) { + long offset = 0; + + // unpack to bounce buffer + assertEquals(true, cp.hasNext()); + while (cp.hasNext()) { + long copied = cp.next(bounceBuffer); + target.copyFromDeviceBufferAsync( + offset, target, 0, copied, Cuda.DEFAULT_STREAM); + offset += copied; + } + + assertEquals(offset, cp.getTotalContiguousSize()); + + try (Table unpacked = Table.fromPackedTable(meta.getMetadataDirectBuffer(), target)) { + assertTablesAreEqual(t1, unpacked); + } + } + } + @Test void testContiguousSplitWithStrings() { ContiguousTable[] splits = null; @@ -3128,6 +3180,30 @@ void testContiguousSplitWithStrings() { } } + @Test + void testContiguousSplitWithStringsChunked() { + try (Table t1 = new Table.TestBuilder() + .column(10, 12, 14, 16, 18, 20, 22, 24, null, 28) + .column(50, 52, 54, 56, 58, 60, 62, 64, 66, null) + .column("A", "B", "C", "D", "E", "F", "G", "H", "I", "J") + .decimal32Column(-3, 10, 12, 14, 16, 18, 20, 22, 24, null, 28) + .decimal64Column(-8, 50L, 52L, 54L, 56L, 58L, 60L, 62L, 64L, 66L, null) + .build(); + DeviceMemoryBuffer bounceBuffer = DeviceMemoryBuffer.allocate(2L*1024*1024); + ChunkedPack cp = t1.makeChunkedPack(2L*1024*1024); + PackedColumnMetadata meta = cp.buildMetadata()) { + + // unpack to bounce buffer + assertEquals(true, cp.hasNext()); + assertEquals(cp.getTotalContiguousSize(), cp.next(bounceBuffer)); + assertEquals(false, cp.hasNext()); + + try (Table unpacked = Table.fromPackedTable(meta.getMetadataDirectBuffer(), bounceBuffer)) { + assertTablesAreEqual(t1, unpacked); + } + } + } + @Test void testPartStability() { final int PARTS = 5; From 6e1a1f47a678719fa6cca105001185e830a77d5a Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 3 May 2023 09:34:41 -0500 Subject: [PATCH 2/8] Update copyright in ContiguousTable.java --- java/src/main/java/ai/rapids/cudf/ContiguousTable.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java/src/main/java/ai/rapids/cudf/ContiguousTable.java b/java/src/main/java/ai/rapids/cudf/ContiguousTable.java index a1aadebd1bd..8193e4f943b 100644 --- a/java/src/main/java/ai/rapids/cudf/ContiguousTable.java +++ b/java/src/main/java/ai/rapids/cudf/ContiguousTable.java @@ -1,6 +1,6 @@ /* * - * Copyright (c) 2019-2021, NVIDIA CORPORATION. + * Copyright (c) 2019-2023, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. From b14126444a9bcc83f286179a82051161cd5b4cad Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Sun, 7 May 2023 23:15:10 -0500 Subject: [PATCH 3/8] Address review comments --- .../main/java/ai/rapids/cudf/ChunkedPack.java | 8 +-- java/src/main/native/src/ChunkedPackJni.cpp | 6 +-- java/src/main/native/src/CudfJni.cpp | 9 ---- .../native/src/PackedColumnMetadataJni.cpp | 49 ------------------- java/src/main/native/src/cudf_jni_apis.hpp | 4 -- 5 files changed, 7 insertions(+), 69 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/ChunkedPack.java b/java/src/main/java/ai/rapids/cudf/ChunkedPack.java index 291e6e41ce7..65c8c69e3b6 100644 --- a/java/src/main/java/ai/rapids/cudf/ChunkedPack.java +++ b/java/src/main/java/ai/rapids/cudf/ChunkedPack.java @@ -20,8 +20,8 @@ /** * JNI interface to cudf::chunked_pack. * - * ChunkedPack is an iterator-like interface with the familiar `hasNext` and `next` - * interface. `next` should be used in a loop until `hasNext` returns false. + * ChunkedPack has an Iterator-like API with the familiar `hasNext` and `next` + * methods. `next` should be used in a loop until `hasNext` returns false. * * However, `ChunkedPack.next` is special because it takes a `DeviceMemoryBuffer` as a * parameter, which means that the caller can call `next` giving any bounce buffer it @@ -83,7 +83,7 @@ public long next(DeviceMemoryBuffer userPtr) { * @return a `PackedColumnMetadata` instance referencing cuDF packed table metadata */ public PackedColumnMetadata buildMetadata() { - return chunkedPackBuildMetadata(nativePtr); + return new PackedColumnMetadata(chunkedPackBuildMetadata(nativePtr)); } @Override @@ -94,6 +94,6 @@ public void close() { private static native long chunkedPackGetTotalContiguousSize(long nativePtr); private static native boolean chunkedPackHasNext(long nativePtr); private static native long chunkedPackNext(long nativePtr, long userPtr, long userPtrSize); - private static native PackedColumnMetadata chunkedPackBuildMetadata(long nativePtr); + private static native long chunkedPackBuildMetadata(long nativePtr); private static native void chunkedPackDelete(long nativePtr); } diff --git a/java/src/main/native/src/ChunkedPackJni.cpp b/java/src/main/native/src/ChunkedPackJni.cpp index b8d6211c590..a496b24467b 100644 --- a/java/src/main/native/src/ChunkedPackJni.cpp +++ b/java/src/main/native/src/ChunkedPackJni.cpp @@ -59,15 +59,15 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackNext( CATCH_STD(env, 0); } -JNIEXPORT jobject JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackBuildMetadata( +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackBuildMetadata( JNIEnv *env, jclass, jlong chunked_pack) { try { cudf::jni::auto_set_device(env); auto cs = reinterpret_cast(chunked_pack); std::unique_ptr> result = cs->build_metadata(); - return cudf::jni::packed_column_metadata_from(env, std::move(result)); + return reinterpret_cast(result.release()); } - CATCH_STD(env, NULL); + CATCH_STD(env, 0); } } // extern "C" diff --git a/java/src/main/native/src/CudfJni.cpp b/java/src/main/native/src/CudfJni.cpp index 66f491f51fd..119f9ed7a06 100644 --- a/java/src/main/native/src/CudfJni.cpp +++ b/java/src/main/native/src/CudfJni.cpp @@ -164,14 +164,6 @@ JNIEXPORT jint JNI_OnLoad(JavaVM *vm, void *) { return JNI_ERR; } - if (!cudf::jni::cache_packed_column_meta_jni(env)) { - if (!env->ExceptionCheck()) { - env->ThrowNew(env->FindClass("java/lang/RuntimeException"), - "Unable to locate packed column metadata methods needed by JNI"); - } - return JNI_ERR; - } - if (!cudf::jni::cache_contig_split_group_by_result_jni(env)) { if (!env->ExceptionCheck()) { env->ThrowNew(env->FindClass("java/lang/RuntimeException"), @@ -201,7 +193,6 @@ JNIEXPORT void JNI_OnUnload(JavaVM *vm, void *) { cudf::jni::release_contiguous_table_jni(env); cudf::jni::release_contig_split_group_by_result_jni(env); cudf::jni::release_host_memory_buffer_jni(env); - cudf::jni::release_packed_column_meta_jni(env); } JNIEXPORT jboolean JNICALL Java_ai_rapids_cudf_Cuda_isPtdsEnabled(JNIEnv *env, jclass, jlong size) { diff --git a/java/src/main/native/src/PackedColumnMetadataJni.cpp b/java/src/main/native/src/PackedColumnMetadataJni.cpp index 0047652be67..7ec3e1294ce 100644 --- a/java/src/main/native/src/PackedColumnMetadataJni.cpp +++ b/java/src/main/native/src/PackedColumnMetadataJni.cpp @@ -16,55 +16,6 @@ #include "cudf_jni_apis.hpp" -namespace { - -#define PACKED_COLUMN_META_CLASS "ai/rapids/cudf/PackedColumnMetadata" -#define PACKED_COLUMN_META_FACTORY_SIG(param_sig) "(" param_sig ")L" PACKED_COLUMN_META_CLASS ";" - -jclass Packed_columns_meta_jclass; -jmethodID From_packed_column_meta_method; - -} // anonymous namespace - -namespace cudf { -namespace jni { - -bool cache_packed_column_meta_jni(JNIEnv *env) { - jclass cls = env->FindClass(PACKED_COLUMN_META_CLASS); - if (cls == nullptr) { - return false; - } - - From_packed_column_meta_method = - env->GetStaticMethodID(cls, "fromPackedColumnMeta", PACKED_COLUMN_META_FACTORY_SIG("J")); - if (From_packed_column_meta_method == nullptr) { - return false; - } - - // Convert local reference to global so it cannot be garbage collected. - Packed_columns_meta_jclass = static_cast(env->NewGlobalRef(cls)); - if (Packed_columns_meta_jclass == nullptr) { - return false; - } - return true; -} - -jobject packed_column_metadata_from(JNIEnv *env, std::unique_ptr> meta) { - jlong metadata_address = reinterpret_cast(meta.release()); - return env->CallStaticObjectMethod(Packed_columns_meta_jclass, From_packed_column_meta_method, - metadata_address); -} - -void release_packed_column_meta_jni(JNIEnv *env) { - if (Packed_columns_meta_jclass != nullptr) { - env->DeleteGlobalRef(Packed_columns_meta_jclass); - Packed_columns_meta_jclass = nullptr; - } -} - -} // namespace jni -} // namespace cudf - extern "C" { JNIEXPORT jobject JNICALL Java_ai_rapids_cudf_PackedColumnMetadata_createMetadataDirectBuffer( diff --git a/java/src/main/native/src/cudf_jni_apis.hpp b/java/src/main/native/src/cudf_jni_apis.hpp index 2005fdfc446..575fd88a805 100644 --- a/java/src/main/native/src/cudf_jni_apis.hpp +++ b/java/src/main/native/src/cudf_jni_apis.hpp @@ -50,13 +50,9 @@ convert_table_for_return(JNIEnv *env, std::unique_ptr &&table_resul // ContiguousTable APIs // -bool cache_packed_column_meta_jni(JNIEnv *env); -void release_packed_column_meta_jni(JNIEnv *env); - bool cache_contiguous_table_jni(JNIEnv *env); void release_contiguous_table_jni(JNIEnv *env); -jobject packed_column_metadata_from(JNIEnv *env, std::unique_ptr> meta); jobject contiguous_table_from(JNIEnv *env, cudf::packed_columns &split, long row_count); native_jobjectArray contiguous_table_array(JNIEnv *env, jsize length); From bc0ce4afd4912696155dde101082bcad5933931d Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Thu, 18 May 2023 13:22:32 -0500 Subject: [PATCH 4/8] Update to use chunked_pack::create --- java/src/main/native/src/TableJni.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/java/src/main/native/src/TableJni.cpp b/java/src/main/native/src/TableJni.cpp index ba0f4af6175..f608e69e357 100644 --- a/java/src/main/native/src/TableJni.cpp +++ b/java/src/main/native/src/TableJni.cpp @@ -3169,9 +3169,9 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Table_makeChunkedPack( auto mr = memoryResourceHandle != 0 ? reinterpret_cast(memoryResourceHandle) : rmm::mr::get_current_device_resource(); - auto chunked_contig_split = - cudf::make_chunked_pack(*n_table, bounce_buffer_size, mr); - return reinterpret_cast(chunked_contig_split.release()); + auto chunked_pack = + cudf::chunked_pack::create(*n_table, bounce_buffer_size, mr); + return reinterpret_cast(chunked_pack.release()); } CATCH_STD(env, 0); } From f9a69fdd05bddeb08e722e5549f04c1cde60a4d7 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Thu, 18 May 2023 13:30:30 -0500 Subject: [PATCH 5/8] Fix styles and bring tempMemoryResource to JNI --- .../main/java/ai/rapids/cudf/ChunkedPack.java | 32 +++++++++---------- .../ai/rapids/cudf/PackedColumnMetadata.java | 2 +- java/src/main/java/ai/rapids/cudf/Table.java | 30 ++++++++--------- java/src/main/native/src/ChunkedPackJni.cpp | 18 ++++++----- java/src/main/native/src/TableJni.cpp | 17 ++++++---- .../test/java/ai/rapids/cudf/TableTest.java | 2 +- 6 files changed, 53 insertions(+), 48 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/ChunkedPack.java b/java/src/main/java/ai/rapids/cudf/ChunkedPack.java index 65c8c69e3b6..c1671deebde 100644 --- a/java/src/main/java/ai/rapids/cudf/ChunkedPack.java +++ b/java/src/main/java/ai/rapids/cudf/ChunkedPack.java @@ -19,16 +19,16 @@ /** * JNI interface to cudf::chunked_pack. - * + * * ChunkedPack has an Iterator-like API with the familiar `hasNext` and `next` - * methods. `next` should be used in a loop until `hasNext` returns false. - * - * However, `ChunkedPack.next` is special because it takes a `DeviceMemoryBuffer` as a - * parameter, which means that the caller can call `next` giving any bounce buffer it - * may have previously allocated. No requirement exists that the bounce buffer be the - * same each time, the only requirement is that their sizes are all the same, and match + * methods. `next` should be used in a loop until `hasNext` returns false. + * + * However, `ChunkedPack.next` is special because it takes a `DeviceMemoryBuffer` as a + * parameter, which means that the caller can call `next` giving any bounce buffer it + * may have previously allocated. No requirement exists that the bounce buffer be the + * same each time, the only requirement is that their sizes are all the same, and match * the size that was passed to `Table.makeChunkedPack` (which instantiates this class). - * + * * The user of `ChunkedPack` must close `.close()` when done using it to clear up both * host and device resources. */ @@ -43,10 +43,10 @@ public class ChunkedPack implements AutoCloseable { public ChunkedPack(long nativePtr) { this.nativePtr = nativePtr; } - + /** * Get the final contiguous size of the table we are packing. This is - * the size that the final buffer should be, just like if the user called + * the size that the final buffer should be, just like if the user called * `cudf::pack` instead. * @return the total number of bytes for the table in contigous layout */ @@ -55,8 +55,8 @@ public long getTotalContiguousSize() { } /** - * Method to be called to ensure that `ChunkedPack` has work left. - * This method should be invoked followed by a call to `next`, until + * Method to be called to ensure that `ChunkedPack` has work left. + * This method should be invoked followed by a call to `next`, until * `hasNext` returns false. * @return true if there is work left to be done (`next` should be called), * false otherwise. @@ -66,11 +66,11 @@ public boolean hasNext() { } /** - * Place the next contigous chunk of our table into `userPtr`. - * + * Place the next contigous chunk of our table into `userPtr`. + * * This method throws if `hasNext` is false. * @param userPtr the bounce buffer to use for this iteration - * @return the number of bytes that we were able to place in `userPtr`. This is + * @return the number of bytes that we were able to place in `userPtr`. This is * at most `userPtr.getLength()`. */ public long next(DeviceMemoryBuffer userPtr) { @@ -78,7 +78,7 @@ public long next(DeviceMemoryBuffer userPtr) { } /** - * Generates opaque table metadata that can be unpacked via `cudf::unpack` + * Generates opaque table metadata that can be unpacked via `cudf::unpack` * at a later time. * @return a `PackedColumnMetadata` instance referencing cuDF packed table metadata */ diff --git a/java/src/main/java/ai/rapids/cudf/PackedColumnMetadata.java b/java/src/main/java/ai/rapids/cudf/PackedColumnMetadata.java index 1b173e7e2aa..5ee4ddcf80b 100644 --- a/java/src/main/java/ai/rapids/cudf/PackedColumnMetadata.java +++ b/java/src/main/java/ai/rapids/cudf/PackedColumnMetadata.java @@ -21,7 +21,7 @@ import java.nio.ByteBuffer; /** - * Metadata for a table that is backed by a single contiguous device buffer. + * Metadata for a table that is backed by a single contiguous device buffer. */ public final class PackedColumnMetadata implements AutoCloseable { private long metadataHandle = 0; diff --git a/java/src/main/java/ai/rapids/cudf/Table.java b/java/src/main/java/ai/rapids/cudf/Table.java index bf92aef655a..3b46e2f07c9 100644 --- a/java/src/main/java/ai/rapids/cudf/Table.java +++ b/java/src/main/java/ai/rapids/cudf/Table.java @@ -182,7 +182,7 @@ static Table removeNullMasksIfNeeded(Table table) { private static native ContiguousTable[] contiguousSplit(long inputTable, int[] indices); - private static native long makeChunkedPack(long inputTable, long bounceBufferSize, long memoryResource); + private static native long makeChunkedPack(long inputTable, long bounceBufferSize, long tempMemoryResource); private static native long[] partition(long inputTable, long partitionView, int numberOfPartitions, int[] outputOffsets); @@ -2169,35 +2169,35 @@ public ContiguousTable[] contiguousSplit(int... indices) { } /** - * Create an instance of `ChunkedPack` which can be used to pack this table + * Create an instance of `ChunkedPack` which can be used to pack this table * contiguously in memory utilizing a bounce buffer of size `bounceBufferSize`. - * + * * This version of `makeChunkedPack` takes a `RmmDviceMemoryResource`, which can be used * to pre-allocate all scratch and temporary space required for the state of `cudf::chunked_pack`. - * + * * The caller is responsible for calling close on the returned `ChunkedPack` object. - * + * * @param bounceBufferSize The size of bounce buffer that will be utilized to pack into - * @param deviceMemoryResource A memory resource that is used to satisfy allocations for - * temporary and thrust scratch space. + * @param tempMemoryResource A memory resource that is used to satisfy allocations for + * temporary and thrust scratch space. * @return An instance of `ChunkedPack` that the caller must use to finish the operation. */ public ChunkedPack makeChunkedPack( - long bounceBufferSize, RmmDeviceMemoryResource deviceMemoryResource) { - long memoryResourceNativeHandle = deviceMemoryResource.getHandle(); + long bounceBufferSize, RmmDeviceMemoryResource tempMemoryResource) { + long tempMemoryResourceHandle = tempMemoryResource.getHandle(); return new ChunkedPack( - makeChunkedPack(nativeHandle, bounceBufferSize, memoryResourceNativeHandle)); + makeChunkedPack(nativeHandle, bounceBufferSize, tempMemoryResourceHandle)); } /** - * Create an instance of `ChunkedPack` which can be used to pack this table + * Create an instance of `ChunkedPack` which can be used to pack this table * contiguously in memory utilizing a bounce buffer of size `bounceBufferSize`. - * - * This version of `makeChunkedPack` makes use of the default per-device memory resource, + * + * This version of `makeChunkedPack` makes use of the default per-device memory resource, * for scratch and temporary space required for the state of `cudf::chunked_pack`. - * + * * The caller is responsible for calling close on the returned `ChunkedPack` object. - * + * * @param bounceBufferSize The size of bounce buffer that will be utilized to pack into * @return An instance of `ChunkedPack` that the caller must use to finish the operation. */ diff --git a/java/src/main/native/src/ChunkedPackJni.cpp b/java/src/main/native/src/ChunkedPackJni.cpp index a496b24467b..746a67e1463 100644 --- a/java/src/main/native/src/ChunkedPackJni.cpp +++ b/java/src/main/native/src/ChunkedPackJni.cpp @@ -17,8 +17,8 @@ #include "cudf_jni_apis.hpp" extern "C" { -JNIEXPORT void JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackDelete( - JNIEnv *env, jclass, jlong chunked_pack) { +JNIEXPORT void JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackDelete(JNIEnv *env, jclass, + jlong chunked_pack) { try { cudf::jni::auto_set_device(env); auto cs = reinterpret_cast(chunked_pack); @@ -37,8 +37,8 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackGetTotalConti CATCH_STD(env, 0); } -JNIEXPORT jboolean JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackHasNext( - JNIEnv *env, jclass, jlong chunked_pack) { +JNIEXPORT jboolean JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackHasNext(JNIEnv *env, jclass, + jlong chunked_pack) { try { cudf::jni::auto_set_device(env); auto cs = reinterpret_cast(chunked_pack); @@ -47,8 +47,10 @@ JNIEXPORT jboolean JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackHasNext( CATCH_STD(env, 0); } -JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackNext( - JNIEnv *env, jclass, jlong chunked_pack, jlong user_ptr, jlong user_ptr_size) { +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackNext(JNIEnv *env, jclass, + jlong chunked_pack, + jlong user_ptr, + jlong user_ptr_size) { try { cudf::jni::auto_set_device(env); auto cs = reinterpret_cast(chunked_pack); @@ -59,8 +61,8 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackNext( CATCH_STD(env, 0); } -JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ChunkedPack_chunkedPackBuildMetadata( - JNIEnv *env, jclass, jlong chunked_pack) { +JNIEXPORT jlong JNICALL +Java_ai_rapids_cudf_ChunkedPack_chunkedPackBuildMetadata(JNIEnv *env, jclass, jlong chunked_pack) { try { cudf::jni::auto_set_device(env); auto cs = reinterpret_cast(chunked_pack); diff --git a/java/src/main/native/src/TableJni.cpp b/java/src/main/native/src/TableJni.cpp index f608e69e357..39a6ef3dea8 100644 --- a/java/src/main/native/src/TableJni.cpp +++ b/java/src/main/native/src/TableJni.cpp @@ -3159,18 +3159,21 @@ JNIEXPORT jobjectArray JNICALL Java_ai_rapids_cudf_Table_contiguousSplit(JNIEnv CATCH_STD(env, NULL); } -JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Table_makeChunkedPack( - JNIEnv *env, jclass, jlong input_table, jlong bounce_buffer_size, jlong memoryResourceHandle) { +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Table_makeChunkedPack(JNIEnv *env, jclass, + jlong input_table, + jlong bounce_buffer_size, + jlong memoryResourceHandle) { JNI_NULL_CHECK(env, input_table, "native handle is null", 0); try { cudf::jni::auto_set_device(env); cudf::table_view *n_table = reinterpret_cast(input_table); - auto mr = memoryResourceHandle != 0 ? - reinterpret_cast(memoryResourceHandle) : - rmm::mr::get_current_device_resource(); - auto chunked_pack = - cudf::chunked_pack::create(*n_table, bounce_buffer_size, mr); + // `temp_mr` is the memory resource that `cudf::chunked_pack` will use to create temporary + // and scratch memory only. + auto temp_mr = memoryResourceHandle != 0 ? + reinterpret_cast(memoryResourceHandle) : + rmm::mr::get_current_device_resource(); + auto chunked_pack = cudf::chunked_pack::create(*n_table, bounce_buffer_size, temp_mr); return reinterpret_cast(chunked_pack.release()); } CATCH_STD(env, 0); diff --git a/java/src/test/java/ai/rapids/cudf/TableTest.java b/java/src/test/java/ai/rapids/cudf/TableTest.java index 8ccda6df079..94a886563a5 100644 --- a/java/src/test/java/ai/rapids/cudf/TableTest.java +++ b/java/src/test/java/ai/rapids/cudf/TableTest.java @@ -3121,7 +3121,7 @@ void testChunkedPackBasic() { } } } - + @Test void testChunkedPackTwoPasses() { // this test packes ~2MB worth of long into a 1MB bounce buffer From 6600dbd3e74276e073c8481ce3e463ccf891a8fe Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Thu, 18 May 2023 14:20:51 -0500 Subject: [PATCH 6/8] Fix spelling --- java/src/main/java/ai/rapids/cudf/ChunkedPack.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/ChunkedPack.java b/java/src/main/java/ai/rapids/cudf/ChunkedPack.java index c1671deebde..90ec05cdbbf 100644 --- a/java/src/main/java/ai/rapids/cudf/ChunkedPack.java +++ b/java/src/main/java/ai/rapids/cudf/ChunkedPack.java @@ -48,7 +48,7 @@ public ChunkedPack(long nativePtr) { * Get the final contiguous size of the table we are packing. This is * the size that the final buffer should be, just like if the user called * `cudf::pack` instead. - * @return the total number of bytes for the table in contigous layout + * @return the total number of bytes for the table in contiguous layout */ public long getTotalContiguousSize() { return chunkedPackGetTotalContiguousSize(nativePtr); @@ -66,7 +66,7 @@ public boolean hasNext() { } /** - * Place the next contigous chunk of our table into `userPtr`. + * Place the next contiguous chunk of our table into `userPtr`. * * This method throws if `hasNext` is false. * @param userPtr the bounce buffer to use for this iteration From f984fbddaa9ad2a7f3a4cf57d9fe710240f6f510 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Thu, 18 May 2023 15:20:16 -0500 Subject: [PATCH 7/8] Remove unnecessary changes --- java/src/main/native/src/CudfJni.cpp | 2 +- java/src/main/native/src/cudf_jni_apis.hpp | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/java/src/main/native/src/CudfJni.cpp b/java/src/main/native/src/CudfJni.cpp index 119f9ed7a06..ee5b0f7a31b 100644 --- a/java/src/main/native/src/CudfJni.cpp +++ b/java/src/main/native/src/CudfJni.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2023, NVIDIA CORPORATION. + * Copyright (c) 2019, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/java/src/main/native/src/cudf_jni_apis.hpp b/java/src/main/native/src/cudf_jni_apis.hpp index 575fd88a805..18993aea294 100644 --- a/java/src/main/native/src/cudf_jni_apis.hpp +++ b/java/src/main/native/src/cudf_jni_apis.hpp @@ -51,6 +51,7 @@ convert_table_for_return(JNIEnv *env, std::unique_ptr &&table_resul // bool cache_contiguous_table_jni(JNIEnv *env); + void release_contiguous_table_jni(JNIEnv *env); jobject contiguous_table_from(JNIEnv *env, cudf::packed_columns &split, long row_count); From 83ac0238dcaa608c2dc64313309a31fc8c7c82e9 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Thu, 18 May 2023 15:21:13 -0500 Subject: [PATCH 8/8] Fix copyright --- java/src/main/native/src/CudfJni.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java/src/main/native/src/CudfJni.cpp b/java/src/main/native/src/CudfJni.cpp index ee5b0f7a31b..1fdadb4b7e2 100644 --- a/java/src/main/native/src/CudfJni.cpp +++ b/java/src/main/native/src/CudfJni.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019, NVIDIA CORPORATION. + * Copyright (c) 2019-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License.