From 9dde4349bd38434205c8d4b0ce127fa55a719591 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Thu, 18 May 2023 17:07:16 -0500 Subject: [PATCH] JNI api for cudf::chunked_pack (#13278) This PR is the standalone JNI side of https://github.com/rapidsai/cudf/pull/13260. Therefore it doesn't build in as is, but I am putting this up as draft for the Java reviewers to start taking a look. This implements a `ChunkedPack` java class mirroring the interface of `cudf::chunked_pack`. This is an iterator-like class that can be used to invoke `cudf::pack` (aka `cudf::contigous_split` without splits) over several iterations against a bounce buffer. In order to create a `ChunkedPack`, the user calls `makeChunkedPack` from a `Table` instance. During this call the user can also pass an `RmmDeviceMemoryResource` to be used internally by `cudf::chunked_pack` exclusively for scratch/temporary data. Authors: - Alessandro Bellina (https://github.com/abellina) Approvers: - Robert (Bobby) Evans (https://github.com/revans2) - Nghia Truong (https://github.com/ttnghia) URL: https://github.com/rapidsai/cudf/pull/13278 --- .../main/java/ai/rapids/cudf/ChunkedPack.java | 99 +++++++++++++++++++ .../java/ai/rapids/cudf/ContiguousTable.java | 27 ++--- .../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 | 75 ++++++++++++++ .../main/native/src/ContiguousTableJni.cpp | 20 ---- .../native/src/PackedColumnMetadataJni.cpp | 41 ++++++++ java/src/main/native/src/TableJni.cpp | 21 ++++ .../test/java/ai/rapids/cudf/TableTest.java | 76 ++++++++++++++ 10 files changed, 437 insertions(+), 38 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..90ec05cdbbf --- /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 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 + * 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 contiguous 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 contiguous 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 new PackedColumnMetadata(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 long 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..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. @@ -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..5ee4ddcf80b --- /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 a9cfe783aac..b553cf9913b 100644 --- a/java/src/main/java/ai/rapids/cudf/Table.java +++ b/java/src/main/java/ai/rapids/cudf/Table.java @@ -182,6 +182,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 tempMemoryResource); + private static native long[] partition(long inputTable, long partitionView, int numberOfPartitions, int[] outputOffsets); @@ -2136,6 +2138,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 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 tempMemoryResource) { + long tempMemoryResourceHandle = tempMemoryResource.getHandle(); + return new ChunkedPack( + makeChunkedPack(nativeHandle, bounceBufferSize, tempMemoryResourceHandle)); + } + + /** + * 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..746a67e1463 --- /dev/null +++ b/java/src/main/native/src/ChunkedPackJni.cpp @@ -0,0 +1,75 @@ +/* + * 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 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 reinterpret_cast(result.release()); + } + CATCH_STD(env, 0); +} + +} // 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/PackedColumnMetadataJni.cpp b/java/src/main/native/src/PackedColumnMetadataJni.cpp new file mode 100644 index 00000000000..7ec3e1294ce --- /dev/null +++ b/java/src/main/native/src/PackedColumnMetadataJni.cpp @@ -0,0 +1,41 @@ +/* + * 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 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 178845ee0d7..c21650bc202 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" @@ -3210,6 +3211,26 @@ 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); + // `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); +} + 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/test/java/ai/rapids/cudf/TableTest.java b/java/src/test/java/ai/rapids/cudf/TableTest.java index 6638e2a7ec6..1569a6eb8e3 100644 --- a/java/src/test/java/ai/rapids/cudf/TableTest.java +++ b/java/src/test/java/ai/rapids/cudf/TableTest.java @@ -3100,6 +3100,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; @@ -3129,6 +3181,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;