Repository: NVIDIA/spark-rapids-jni
Pull request #3100: Java APIs for GPU kudo
Author: login: revans2
Created: 2025-04-03T20:57:13Z | Merged: 2025-05-14T18:56:03Z

Please note that this updates the CPU format to match the GPU format for kudo in places where the GPU appears to be more self consistent or the size of the data was smaller.

The main changes are padding and buffer order. Padding only happens at the end of each section, instead of in between each buffer. The order of the buffers for a struct has changed so that the child buffers come after the buffers for the struct. This makes it consistent with lists, which also had them come after. I also updated the code so that if a validity buffer was skipped because there were no rows, then the flag to say if the validity buffer is there or not indicates that it is not there. Before the CPU code would set it to say that the buffer was there in some cases, but know that it was not there because the row could was 0. This just makes it more consistent and should be matching what the GPU is doing.


Diff:
diff --git a/src/main/cpp/CMakeLists.txt b/src/main/cpp/CMakeLists.txt
index 50c7d6f67c..f4367c030a 100644
--- a/src/main/cpp/CMakeLists.txt
+++ b/src/main/cpp/CMakeLists.txt
@@ -193,6 +193,7 @@ add_library(
   src/HostTableJni.cpp
   src/HyperLogLogPlusPlusHostUDFJni.cpp
   src/JSONUtilsJni.cpp
+  src/KudoGpuSerializerJni.cpp
   src/ListSliceJni.cpp
   src/MapJni.cpp
   src/NativeParquetJni.cpp
diff --git a/src/main/cpp/src/KudoGpuSerializerJni.cpp b/src/main/cpp/src/KudoGpuSerializerJni.cpp
new file mode 100644
index 0000000000..5134015f25
--- /dev/null
+++ b/src/main/cpp/src/KudoGpuSerializerJni.cpp
@@ -0,0 +1,116 @@
+/*
+ * Copyright (c) 2025, 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"
+#include "shuffle_split.hpp"
+
+extern "C" {
+
+JNIEXPORT jlongArray JNICALL
+Java_com_nvidia_spark_rapids_jni_kudo_KudoGpuSerializer_splitAndSerializeToDevice(
+  JNIEnv* env, jclass, jlong j_table_view, jintArray j_splits)
+{
+  JNI_NULL_CHECK(env, j_table_view, "table is null", NULL);
+  JNI_NULL_CHECK(env, j_splits, "splits is null", NULL);
+  try {
+    cudf::jni::auto_set_device(env);
+
+    auto table = reinterpret_cast<cudf::table_view const*>(j_table_view);
+    const cudf::jni::native_jintArray n_splits(env, j_splits);
+    std::vector<cudf::size_type> splits = n_splits.to_vector<int>();
+
+    auto [split_result, split_meta] = spark_rapids_jni::shuffle_split(
+      *table, splits, cudf::get_default_stream(), cudf::get_current_device_resource());
+
+    // The following code is ugly. We need to return two device buffers to java, but
+    // there is no good way to do this. For this we return three values for each buffer.
+    // These values are {data_address, data_size, rmm::device_buffer*}
+    // These values are then returned to java and we call into `DeviceMemoryBuffer.fromRmm`
+    // That creates a new DeviceMemoryBuffer that takes ownership of the rmm::device_buffer*
+    // and will free it when the DeviceMemoryBuffer is closed.
+    // To make this work it looks like we pull data out of the rmm::device_buffer and
+    // then either leak it or release the memory held by it, but that is not technically
+    // the case.
+    cudf::jni::native_jlongArray result(env, 6);
+    result[0] = reinterpret_cast<jlong>(split_result.partitions->data());
+    result[1] = static_cast<jlong>(split_result.partitions->size());
+    result[2] = reinterpret_cast<jlong>(split_result.partitions.release());
+
+    // split_result.offsets is an rmm::device_uvector<size_t> so we have to
+    // pull out the rmm::device_buffer * from inside it to return the data in a way that
+    // java can handle it.
+    auto offsets = std::make_unique<rmm::device_buffer>(std::move(split_result.offsets.release()));
+    result[3]    = reinterpret_cast<jlong>(offsets->data());
+    result[4]    = static_cast<jlong>(offsets->size());
+    result[5]    = reinterpret_cast<jlong>(offsets.release());
+
+    return result.get_jArray();
+  }
+  CATCH_STD(env, NULL);
+}
+
+JNIEXPORT jlongArray JNICALL
+Java_com_nvidia_spark_rapids_jni_kudo_KudoGpuSerializer_assembleFromDeviceRawNative(
+  JNIEnv* env,
+  jclass,
+  jlong part_addr,
+  jlong part_len,
+  jlong offset_addr,
+  jlong offset_len,
+  jintArray flat_num_children,
+  jintArray flat_type_ids,
+  jintArray flat_scale)
+{
+  JNI_NULL_CHECK(env, part_addr, "part_addr is null", NULL);
+  JNI_NULL_CHECK(env, offset_addr, "offset_addr is null", NULL);
+  JNI_NULL_CHECK(env, flat_num_children, "num_children is null", NULL);
+  JNI_NULL_CHECK(env, flat_type_ids, "type_ids is null", NULL);
+  JNI_NULL_CHECK(env, flat_scale, "scale is null", NULL);
+
+  try {
+    cudf::jni::auto_set_device(env);
+
+    cudf::device_span<uint8_t const> partitions(reinterpret_cast<uint8_t*>(part_addr), part_len);
+    cudf::device_span<size_t const> offsets(reinterpret_cast<size_t*>(offset_addr),
+                                            offset_len / sizeof(size_t));
+
+    cudf::jni::native_jintArray nnc(env, flat_num_children);
+    cudf::jni::native_jintArray nti(env, flat_type_ids);
+    cudf::jni::native_jintArray ns(env, flat_scale);
+
+    spark_rapids_jni::shuffle_split_metadata meta;
+    meta.col_info.reserve(nnc.size());
+
+    for (int i = 0; i < nnc.size(); ++i) {
+      auto tid          = static_cast<cudf::type_id>(nti[i]);
+      auto scale        = ns[i];
+      auto num_children = static_cast<cudf::size_type>(nnc[i]);
+      cudf::size_type param =
+        spark_rapids_jni::is_fixed_point(cudf::data_type{tid, scale}) ? scale : num_children;
+      meta.col_info.emplace_back(tid, param);
+    }
+
+    return cudf::jni::convert_table_for_return(
+      env,
+      shuffle_assemble(meta,
+                       partitions,
+                       offsets,
+                       cudf::get_default_stream(),
+                       cudf::get_current_device_resource()));
+  }
+  CATCH_STD(env, NULL);
+}
+}
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoGpuSerializer.java b/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoGpuSerializer.java
new file mode 100644
index 0000000000..405c042e92
--- /dev/null
+++ b/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoGpuSerializer.java
@@ -0,0 +1,81 @@
+/*
+ * Copyright (c) 2025, 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 com.nvidia.spark.rapids.jni.kudo;
+
+import ai.rapids.cudf.*;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * Right now this is just to provide access to the underlying C++ APIs. In the future it should hopefully
+ * look more like the CPU KudoSerializer.
+ */
+public class KudoGpuSerializer {
+  static {
+    NativeDepsLoader.loadNativeDeps();
+  }
+
+  public static DeviceMemoryBuffer[] splitAndSerializeToDevice(Table table, int... splits) {
+    DeviceMemoryBuffer[] ret = new DeviceMemoryBuffer[2];
+    boolean success = false;
+    try {
+      long[] values = splitAndSerializeToDevice(table.getNativeView(), splits);
+      ret[0] = DeviceMemoryBuffer.fromRmm(values[0], values[1], values[2]);
+      ret[1] = DeviceMemoryBuffer.fromRmm(values[3], values[4], values[5]);
+      success = true;
+      return ret;
+    } finally {
+      if (!success) {
+        if (ret[0] != null) {
+          ret[0].close();
+        }
+        if (ret[1] != null) {
+          ret[1].close();
+        }
+      }
+    }
+  }
+
+  public static Table assembleFromDeviceRaw(Schema schema,
+                                            DeviceMemoryBuffer partitions,
+                                            DeviceMemoryBuffer offsets) {
+    return new Table(assembleFromDeviceRawNative(
+        partitions.getAddress(), partitions.getLength(),
+        offsets.getAddress(), offsets.getLength(),
+        schema.getFlattenedNumChildren(),
+        schema.getFlattenedTypeIds(),
+        schema.getFlattenedTypeScales()));
+  }
+
+  /**
+   * Split the input table and serialize it to a device buffer.
+   * @param tableNativeView the native view of the table to split
+   * @param splits the row indices to split around
+   * @return an array of 6 longs. These represent two device buffers each with 3 values.
+   * (address, length, rmmAddress) The first tuple holds all the serialized data. The second tuple holds
+   * size_t offsets into the first to show where the splits are.
+   */
+  private static native long[] splitAndSerializeToDevice(long tableNativeView, int[] splits);
+
+
+  private static native long[] assembleFromDeviceRawNative(long partAddr, long partLen,
+                                                           long offsetAddr, long offsetLen,
+                                                           int[] flattenedNumChildren,
+                                                           int[] flattenedTypeIds,
+                                                           int[] flattenedTypeScales);
+
+}
\ No newline at end of file
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoSerializer.java b/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoSerializer.java
index 58e5cc145b..156bbcde48 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoSerializer.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoSerializer.java
@@ -67,6 +67,7 @@
  * <h1>Format</h1>
  * <p>
  * Similar to {@link JCudfSerialization}, it still consists of two parts: header and body.
+ * <p>
  *
  * <h2>Header</h2>
  * <p>
@@ -126,6 +127,14 @@
  *           int bit = col<sub>i</sub> % 8; <br/>
  *           return (hasValidityBuffer[pos] & (1 << bit)) != 0;
  *         </code>
+ *         The order of the bits is the same as the order of the buffers in the body. They are depth-first
+ *         when walking the schema, but for structs and arrays the validity buffer for that object itself
+ *         comes before its children.
+ *         <br/>
+ *         In all cases if hasValidityBuffer indicates that validity is present at least 1 byte must be
+ *         output in the body for that. If because of nesting a buffer would have 0 rows, then hasValidityBuffer
+ *         should either indicate that there is no validity or insert in a byte that can be ignored. The first
+ *         option is preferable.
  *         </td>
  *     </tr>
  * </table>
@@ -134,13 +143,17 @@
  * <p>
  * The body consists of three part:
  * <ol>
- *     <li>Validity buffers for every column with validity in depth-first ordering of schema columns. Each buffer of
- *     each column is 4 bytes padded.
+ *     <li>Validity buffers for every column with validity in depth-first ordering of schema columns. Just like with
+ *     hasValidityBuffer the validity for structs and arrays comes before their children. The entire validity part
+ *     is padded to 4 byte alignment. Because the header is not padded, this takes the header length into account
+ *     when padding.
  *     </li>
  *     <li>Offset buffers for every column with offsets in depth-first ordering of schema columns. Each buffer of each
- *     column is 4 bytes padded.</li>
- *     <li>Data buffers for every column with data in depth-first ordering of schema columns. Each buffer of each
- *     column is 4 bytes padded.</li>
+ *     column is inherently 4-byte aligned because offsets are 4-byte values and the validity if 4-byte aligned.
+ *     Because of nesting it is possible for an offset to have a length of 0, if there are 0 rows.</li>
+ *     <li>Data buffers for every column with data in depth-first ordering of schema columns. The entire part
+ *     will also be padded to 4 byte alignment, but the buffers within the part have no alignment guarantees.
+ *     Because of nesting it is possible of a data buffer to have a length of 0, if there are 0 rows.</li>
  * </ol>
  *
  * <h1>Serialization</h1>
@@ -432,7 +445,8 @@ private WriteMetrics writeSliced(HostColumnVector[] columns, DataWriter out, int
     for (BufferType bufferType : ALL_BUFFER_TYPES) {
       SlicedBufferSerializer serializer = new SlicedBufferSerializer(rowOffset,
           numRows, bufferType,
-          out, metrics, measureCopyBufferTime);
+          out, metrics, measureCopyBufferTime,
+          header.getSerializedSize());
       Visitors.visitColumns(columns, serializer);
       bytesWritten += serializer.getTotalDataLen();
       metrics.addWrittenBytes(serializer.getTotalDataLen());
@@ -461,9 +475,27 @@ private static DataWriter writerFrom(OutputStream out) {
     }
   }
 
+  /**
+   * Based on cudf::util::round_up_safe
+   */
+  private static long roundUpSafe(long toRound, long modulus) {
+    long remainder = toRound % modulus;
+    if (remainder == 0) {
+      return toRound;
+    }
+    long roundedUp = toRound - remainder + modulus;
+    if (roundedUp < toRound) {
+      throw new IllegalStateException("Overflow when rounding up");
+    }
+    return roundedUp;
+  }
 
   static long padForHostAlignment(long orig) {
-    return ((orig + 3) / 4) * 4;
+    return roundUpSafe(orig, 4);
+  }
+
+  static long padForValidityAlignment(long orig, long headerSize) {
+    return roundUpSafe(orig + headerSize, 4) - headerSize;
   }
 
   static long padForHostAlignment(DataWriter out, long bytes) throws IOException {
@@ -474,6 +506,14 @@ static long padForHostAlignment(DataWriter out, long bytes) throws IOException {
     return paddedBytes;
   }
 
+  static long padForValidityAlignment(DataWriter out, long bytes, long headerSize) throws IOException {
+    final long paddedBytes = padForValidityAlignment(bytes, headerSize);
+    if (paddedBytes > bytes) {
+      out.write(PADDING, 0, (int) (paddedBytes - bytes));
+    }
+    return paddedBytes;
+  }
+
   static long padFor64byteAlignment(long orig) {
     return ((orig + 63) / 64) * 64;
   }
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeader.java b/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeader.java
index e60a237b23..5808a17437 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeader.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeader.java
@@ -17,6 +17,8 @@
 package com.nvidia.spark.rapids.jni.kudo;
 
 import ai.rapids.cudf.BufferType;
+import ai.rapids.cudf.HostMemoryBuffer;
+
 import java.io.DataInputStream;
 import java.io.EOFException;
 import java.io.IOException;
@@ -155,6 +157,10 @@ public int getSerializedSize() {
     return 7 * Integer.BYTES + hasValidityBuffer.length;
   }
 
+  public static int getSerializedSize(int hasValiditySize) {
+    return 7 * Integer.BYTES + hasValiditySize;
+  }
+
   public int getNumColumns() {
     return numColumns;
   }
@@ -193,6 +199,27 @@ public void writeTo(DataWriter dout) throws IOException {
     dout.write(hasValidityBuffer, 0, hasValidityBuffer.length);
   }
 
+  private static void setInt(HostMemoryBuffer hmb, long hmbOffset, int v) {
+    byte[] bytes = new byte[4];
+    bytes[0] = (byte) ((v >>> 24) & 0xFF);
+    bytes[1] = (byte) ((v >>> 16) & 0xFF);
+    bytes[2] = (byte) ((v >>> 8) & 0xFF);
+    bytes[3] = (byte) (v & 0xFF);
+    hmb.setBytes(hmbOffset, bytes, 0, bytes.length);
+  }
+
+  public long writeTo(HostMemoryBuffer hmb, long hmbOffset) {
+    setInt(hmb, hmbOffset, SER_FORMAT_MAGIC_NUMBER);
+    setInt(hmb, hmbOffset + 4, offset);
+    setInt(hmb, hmbOffset + 8, numRows);
+    setInt(hmb, hmbOffset + 12, validityBufferLen);
+    setInt(hmb, hmbOffset + 16, offsetBufferLen);
+    setInt(hmb, hmbOffset + 20, totalDataLen);
+    setInt(hmb, hmbOffset + 24, numColumns);
+    hmb.setBytes(hmbOffset + 28, hasValidityBuffer, 0 , hasValidityBuffer.length);
+    return hmbOffset + getSerializedSize();
+  }
+
   @Override
   public String toString() {
     return "SerializedTableHeader{" +
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java b/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
index e3800fe066..3457d35f73 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
@@ -17,14 +17,15 @@
 package com.nvidia.spark.rapids.jni.kudo;
 
 import ai.rapids.cudf.DType;
+import ai.rapids.cudf.HostColumnVector;
 import ai.rapids.cudf.HostColumnVectorCore;
 import com.nvidia.spark.rapids.jni.schema.HostColumnsVisitor;
 
 import java.util.ArrayDeque;
 import java.util.Deque;
-import java.util.List;
 
 import static com.nvidia.spark.rapids.jni.kudo.KudoSerializer.padForHostAlignment;
+import static com.nvidia.spark.rapids.jni.kudo.KudoSerializer.padForValidityAlignment;
 import static java.lang.Math.toIntExact;
 
 /**
@@ -41,14 +42,14 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   private final byte[] bitset;
   private long validityBufferLen;
   private long offsetBufferLen;
-  private long totalDataLen;
+  private long dataOnlyLen;
   private int nextColIdx;
 
   private Deque<SliceInfo> sliceInfos = new ArrayDeque<>();
 
   KudoTableHeaderCalc(int rowOffset, int numRows, int numFlattenedCols) {
     this.root = new SliceInfo(rowOffset, numRows);
-    this.totalDataLen = 0;
+    this.dataOnlyLen = 0;
     sliceInfos.addLast(this.root);
     this.bitset = new byte[(numFlattenedCols + 7) / 8];
     this.numFlattenedCols = numFlattenedCols;
@@ -56,50 +57,62 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   }
 
   public KudoTableHeader getHeader() {
+    int headerSize = KudoTableHeader.getSerializedSize(bitset.length);
+    // The validity is a bit odd because we want to pad it for 4 byte alignment
+    // But that is relative to the header, not the payload buffer
+    long paddedValiditySize = padForValidityAlignment(validityBufferLen, headerSize);
+    long paddedOffsetsSize = padForHostAlignment(offsetBufferLen);
+    long paddedDataSize = padForHostAlignment(dataOnlyLen);
     return new KudoTableHeader(toIntExact(root.offset),
         toIntExact(root.rowCount),
-        toIntExact(validityBufferLen),
-        toIntExact(offsetBufferLen),
-        toIntExact(totalDataLen),
+        toIntExact(paddedValiditySize),
+        toIntExact(paddedOffsetsSize),
+        toIntExact(paddedValiditySize +
+            paddedOffsetsSize +
+            paddedDataSize),
         numFlattenedCols,
         bitset);
   }
 
   @Override
-  public void visitStruct(HostColumnVectorCore col) {
+  public void preVisitStruct(HostColumnVectorCore col) {
     SliceInfo parent = sliceInfos.getLast();
 
     long validityBufferLength = 0;
-    if (col.hasValidityVector()) {
-      validityBufferLength = padForHostAlignment(parent.getValidityBufferInfo().getBufferLength());
+    boolean includeValidity = col.hasValidityVector() && parent.rowCount > 0;
+    if (includeValidity) {
+      validityBufferLength = parent.getValidityBufferInfo().getBufferLength();
     }
 
     this.validityBufferLen += validityBufferLength;
+    this.setHasValidity(includeValidity);
+  }
 
-    totalDataLen += validityBufferLength;
-    this.setHasValidity(col.hasValidityVector());
+  @Override
+  public void visitStruct(HostColumnVectorCore col) {
+    // NOOP
   }
 
   @Override
   public void preVisitList(HostColumnVectorCore col) {
     SliceInfo parent = sliceInfos.getLast();
 
+    boolean includeValidity = col.hasValidityVector() && parent.rowCount > 0;
 
     long validityBufferLength = 0;
-    if (col.hasValidityVector() && parent.rowCount > 0) {
-      validityBufferLength = padForHostAlignment(parent.getValidityBufferInfo().getBufferLength());
+    if (includeValidity) {
+      validityBufferLength = parent.getValidityBufferInfo().getBufferLength();
     }
 
     long offsetBufferLength = 0;
     if (col.getOffsets() != null && parent.rowCount > 0) {
-      offsetBufferLength = padForHostAlignment((parent.rowCount + 1) * Integer.BYTES);
+      offsetBufferLength = (parent.rowCount + 1L) * Integer.BYTES;
     }
 
     this.validityBufferLen += validityBufferLength;
     this.offsetBufferLen += offsetBufferLength;
-    this.totalDataLen += validityBufferLength + offsetBufferLength;
 
-    this.setHasValidity(col.hasValidityVector());
+    this.setHasValidity(includeValidity);
 
     SliceInfo current;
 
@@ -130,11 +143,14 @@ public void visit(HostColumnVectorCore col) {
 
     this.validityBufferLen += validityBufferLen;
     this.offsetBufferLen += offsetBufferLen;
-    this.totalDataLen += validityBufferLen + offsetBufferLen + dataBufferLen;
+    this.dataOnlyLen += dataBufferLen;
 
-    this.setHasValidity(col.hasValidityVector());
+    this.setHasValidity(col.hasValidityVector() && parent.rowCount > 0);
   }
 
+  @Override
+  public void done() {}
+
   private void setHasValidity(boolean hasValidityBuffer) {
     if (hasValidityBuffer) {
       int bytePos = nextColIdx / 8;
@@ -146,7 +162,7 @@ private void setHasValidity(boolean hasValidityBuffer) {
 
   private static long dataLenOfValidityBuffer(HostColumnVectorCore col, SliceInfo info) {
     if (col.hasValidityVector() && info.getRowCount() > 0) {
-      return padForHostAlignment(info.getValidityBufferInfo().getBufferLength());
+      return info.getValidityBufferInfo().getBufferLength();
     } else {
       return 0;
     }
@@ -154,7 +170,7 @@ private static long dataLenOfValidityBuffer(HostColumnVectorCore col, SliceInfo
 
   private static long dataLenOfOffsetBuffer(HostColumnVectorCore col, SliceInfo info) {
     if (DType.STRING.equals(col.getType()) && info.getRowCount() > 0) {
-      return padForHostAlignment((info.rowCount + 1) * Integer.BYTES);
+      return (info.rowCount + 1L) * Integer.BYTES;
     } else {
       return 0;
     }
@@ -165,13 +181,13 @@ private static long dataLenOfDataBuffer(HostColumnVectorCore col, SliceInfo info
       if (col.getOffsets() != null) {
         long startByteOffset = col.getOffsets().getInt(info.offset * Integer.BYTES);
         long endByteOffset = col.getOffsets().getInt((info.offset + info.rowCount) * Integer.BYTES);
-        return padForHostAlignment(endByteOffset - startByteOffset);
+        return endByteOffset - startByteOffset;
       } else {
         return 0;
       }
     } else {
       if (col.getType().getSizeInBytes() > 0) {
-        return padForHostAlignment(col.getType().getSizeInBytes() * info.rowCount);
+        return col.getType().getSizeInBytes() * info.rowCount;
       } else {
         return 0;
       }
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableMerger.java b/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableMerger.java
index 1f14d61b68..56614c9395 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableMerger.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableMerger.java
@@ -105,7 +105,7 @@ public void visitTopSchema(Schema schema) {
   }
 
   @Override
-  public void visitStruct(Schema structType) {
+  public void preVisitStruct(Schema structType) {
     ColumnOffsetInfo offsetInfo = getCurColumnOffsets();
     int nullCount = deserializeValidityBuffer(offsetInfo);
     int totalRowCount = rowCounts[curColIdx];
@@ -116,13 +116,17 @@ public void visitStruct(Schema structType) {
     for (int i=0; i<kudoTables.length; i++) {
       KudoTableHeader header = kudoTables[i].getHeader();
       SliceInfo sliceInfo = sliceInfoOf(i);
-      if (header.hasValidityBuffer(curColIdx)) {
-        validityOffsets[i] += padForHostAlignment(sliceInfo.getValidityBufferInfo().getBufferLength());
+      if (header.hasValidityBuffer(curColIdx) && sliceInfo.getRowCount() > 0) {
+        validityOffsets[i] += sliceInfo.getValidityBufferInfo().getBufferLength();
       }
     }
     curColIdx++;
   }
 
+  public void visitStruct(Schema structType) {
+    // Noop
+  }
+
   @Override
   public void preVisitList(Schema listType) {
     ColumnOffsetInfo offsetInfo = getCurColumnOffsets();
@@ -136,11 +140,11 @@ public void preVisitList(Schema listType) {
     for (int i=0; i<kudoTables.length; i++) {
       KudoTableHeader header = kudoTables[i].getHeader();
       SliceInfo sliceInfo = sliceInfoOf(i);
-      if (header.hasValidityBuffer(curColIdx)) {
-        validityOffsets[i] += padForHostAlignment(sliceInfo.getValidityBufferInfo().getBufferLength());
+      if (header.hasValidityBuffer(curColIdx) && sliceInfo.getRowCount() > 0) {
+        validityOffsets[i] += sliceInfo.getValidityBufferInfo().getBufferLength();
       }
       if (sliceInfo.getRowCount() > 0) {
-        offsetOffsets[i] += padForHostAlignment((sliceInfo.getRowCount() + 1) * Integer.BYTES);
+        offsetOffsets[i] += (sliceInfo.getRowCount() + 1L) * Integer.BYTES;
       }
       sliceInfos[i].addLast(sliceInfoBuf[i]);
     }
@@ -173,23 +177,23 @@ public void visit(Schema primitiveType) {
       for (int i=0; i<kudoTables.length; i++) {
         KudoTableHeader header = kudoTables[i].getHeader();
         SliceInfo sliceInfo = sliceInfoOf(i);
-        if (header.hasValidityBuffer(curColIdx)) {
-          validityOffsets[i] += padForHostAlignment(sliceInfo.getValidityBufferInfo().getBufferLength());
+        if (header.hasValidityBuffer(curColIdx) && sliceInfo.getRowCount() > 0) {
+          validityOffsets[i] += sliceInfo.getValidityBufferInfo().getBufferLength();
         }
         if (sliceInfo.getRowCount() > 0) {
-          offsetOffsets[i] += padForHostAlignment((sliceInfo.getRowCount() + 1) * Integer.BYTES);
-          dataOffsets[i] += padForHostAlignment(sliceInfoBuf[i].getRowCount());
+          offsetOffsets[i] += (sliceInfo.getRowCount() + 1L) * Integer.BYTES;
+          dataOffsets[i] += sliceInfoBuf[i].getRowCount();
         }
       }
     } else {
       for (int i=0; i<kudoTables.length; i++) {
         KudoTableHeader header = kudoTables[i].getHeader();
         SliceInfo sliceInfo = sliceInfoOf(i);
-        if (header.hasValidityBuffer(curColIdx)) {
-          validityOffsets[i] += padForHostAlignment(sliceInfo.getValidityBufferInfo().getBufferLength());
+        if (header.hasValidityBuffer(curColIdx) && sliceInfo.getRowCount() > 0) {
+          validityOffsets[i] += sliceInfo.getValidityBufferInfo().getBufferLength();
         }
         if (sliceInfo.getRowCount() > 0) {
-          dataOffsets[i] += padForHostAlignment(primitiveType.getType().getSizeInBytes() * sliceInfo.getRowCount());
+          dataOffsets[i] += primitiveType.getType().getSizeInBytes() * sliceInfo.getRowCount();
         }
       }
     }
@@ -481,7 +485,7 @@ private int copySourceCaseOne(HostMemoryBuffer src, int srcOffset,
       int outputMask = (1 << curDestBitIdx) - 1;
       int destOutput = dest.getInt(curDestIntIdx) & outputMask;
 
-      int rawInput = src.getInt(curSrcIntIdx);
+      int rawInput = getIntSafe(src, curSrcIntIdx);
       int input = (rawInput >>> curSrcBitIdx) << curDestBitIdx;
       destOutput = input | destOutput;
       dest.setInt(curDestIntIdx, destOutput);
@@ -515,7 +519,7 @@ private int copySourceCaseOne(HostMemoryBuffer src, int srcOffset,
           break;
         }
 
-        src.getInts(inputBuf, 0, curSrcIntIdx, curArrLen);
+        getIntsSafe(src, inputBuf, 0, curSrcIntIdx, curArrLen);
 
         for (int i=0; i<curArrLen; i++) {
           outputBuf[i] = (inputBuf[i] << rshift) | lastValue;
@@ -555,7 +559,7 @@ private int copySourceCaseTwo(HostMemoryBuffer src, int srcOffset,
       int destMask = (1 << curDestBitIdx) - 1;
       int destOutput = dest.getInt(curDestIntIdx) & destMask;
 
-      int input = src.getInt(curSrcIntIdx);
+      int input = getIntSafe(src, curSrcIntIdx);
       if (srcIntBufLen == 1) {
         int leftRem = 32 - curSrcBitIdx - leftRowCount;
         assert leftRem >= 0;
@@ -581,7 +585,7 @@ private int copySourceCaseTwo(HostMemoryBuffer src, int srcOffset,
           leftRowCount = 0;
           break;
         }
-        src.getInts(inputBuf, 0, curSrcIntIdx, curArrLen);
+        getIntsSafe(src, inputBuf, 0, curSrcIntIdx, curArrLen);
         for (int i=0; i<curArrLen; i++) {
           outputBuf[i] = (inputBuf[i] << (32 - rshift)) | lastValue;
           nullCount += 32 - Integer.bitCount(outputBuf[i]);
@@ -601,6 +605,40 @@ private int copySourceCaseTwo(HostMemoryBuffer src, int srcOffset,
       return nullCount;
     }
 
+    private static int getIntAsBytes(HostMemoryBuffer src, long offset, long length) {
+      // We need to build up the int ourselves and pad it as needed
+      int ret = 0;
+      for (int at = 0; at < 4 && (offset + at) < length ; at++) {
+        int b = src.getByte(offset + at) & 0xFF;
+        ret |= b << (at * 8);
+      }
+      return ret;
+    }
+
+    private static int getIntSafe(HostMemoryBuffer src, long offset) {
+      long length = src.getLength();
+      if (offset + 4 < length) {
+        return src.getInt(offset);
+      } else {
+        return getIntAsBytes(src, offset, length);
+      }
+    }
+
+    private static void getIntsSafe(HostMemoryBuffer src, int[] dst, long dstIndex, long srcOffset, int count) {
+      long length = src.getLength();
+      if (srcOffset + (4L * count) < length) {
+        src.getInts(dst, dstIndex, srcOffset, count);
+      } else {
+        // Read as much as we can the fast way
+        int fastCount = (int)((length - srcOffset) / 4);
+        src.getInts(dst, dstIndex, srcOffset, fastCount);
+        // Then read the rest slowly...
+        for (int index = fastCount; index < count; index++) {
+          dst[index + (int)dstIndex] = getIntAsBytes(src, srcOffset + (index * 4L), length);
+        }
+      }
+    }
+
     private int copySourceCaseThree(HostMemoryBuffer src, int srcOffset,
                                     SliceInfo sliceInfo) {
       int leftRowCount = sliceInfo.getRowCount();
@@ -615,7 +653,7 @@ private int copySourceCaseThree(HostMemoryBuffer src, int srcOffset,
 
       // Process first element
       int mask = (1 << curDestBitIdx) - 1;
-      int firstInput = src.getInt(curSrcIntIdx);
+      int firstInput = getIntSafe(src, curSrcIntIdx);
       int destOutput = dest.getInt(curDestIntIdx);
       destOutput = (firstInput & ~mask) | (destOutput & mask);
       dest.setInt(curDestIntIdx, destOutput);
@@ -637,7 +675,7 @@ private int copySourceCaseThree(HostMemoryBuffer src, int srcOffset,
       while (leftRowCount > 0) {
         int curArrLen = min(min(inputBuf.length, outputBuf.length), srcIntBufLen - (curSrcIntIdx - srcOffset) / 4);
         assert curArrLen > 0;
-        src.getInts(inputBuf, 0, curSrcIntIdx, curArrLen);
+        getIntsSafe(src, inputBuf, 0, curSrcIntIdx, curArrLen);
         for (int i=0; i<curArrLen; i++) {
           nullCount += 32 - Integer.bitCount(inputBuf[i]);
           leftRowCount -= 32;
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/kudo/MergedInfoCalc.java b/src/main/java/com/nvidia/spark/rapids/jni/kudo/MergedInfoCalc.java
index 8a27577155..c37b055a42 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/kudo/MergedInfoCalc.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/kudo/MergedInfoCalc.java
@@ -118,7 +118,7 @@ public void visitTopSchema(Schema schema) {
     }
 
     @Override
-    public void visitStruct(Schema structType) {
+    public void preVisitStruct(Schema structType) {
         long validityOffset = INVALID_OFFSET;
         long validityBufferLen = 0;
 
@@ -133,6 +133,11 @@ public void visitStruct(Schema structType) {
         curColIdx++;
     }
 
+    @Override
+    public void visitStruct(Schema structType) {
+        // Noop
+    }
+
     @Override
     public void preVisitList(Schema listType) {
         long validityOffset = INVALID_OFFSET;
@@ -222,13 +227,18 @@ public void visitTopSchema(Schema schema) {
         }
 
         @Override
-        public void visitStruct(Schema structType) {
+        public void preVisitStruct(Schema structType) {
             SliceInfo sliceInfo = sliceInfos.getLast();
             rowCount[curColIdx] += sliceInfo.getRowCount();
 
             curColIdx++;
         }
 
+        @Override
+        public void visitStruct(Schema structType) {
+            // Noop
+        }
+
         @Override
         public void preVisitList(Schema listType) {
             SliceInfo sliceInfo = sliceInfos.getLast();
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/kudo/SlicedBufferSerializer.java b/src/main/java/com/nvidia/spark/rapids/jni/kudo/SlicedBufferSerializer.java
index d937af514d..9ae1ada4c9 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/kudo/SlicedBufferSerializer.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/kudo/SlicedBufferSerializer.java
@@ -18,6 +18,7 @@
 
 import static com.nvidia.spark.rapids.jni.kudo.KudoSerializer.KUDO_SANITY_CHECK;
 import static com.nvidia.spark.rapids.jni.kudo.KudoSerializer.padForHostAlignment;
+import static com.nvidia.spark.rapids.jni.kudo.KudoSerializer.padForValidityAlignment;
 
 import ai.rapids.cudf.BufferType;
 import ai.rapids.cudf.DType;
@@ -27,7 +28,6 @@
 import java.io.IOException;
 import java.util.ArrayDeque;
 import java.util.Deque;
-import java.util.List;
 
 /**
  * This class visits a list of columns and serialize one of the buffers (validity, offset, or data) into with kudo
@@ -51,9 +51,10 @@ class SlicedBufferSerializer implements HostColumnsVisitor {
   private final WriteMetrics metrics;
   private final boolean addCopyBufferTime;
   private long totalDataLen;
+  private long headerSize;
 
   SlicedBufferSerializer(int rowOffset, int numRows, BufferType bufferType, DataWriter writer,
-                         WriteMetrics metrics, boolean addCopyBufferTime) {
+                         WriteMetrics metrics, boolean addCopyBufferTime, long headerSize) {
     this.root = new SliceInfo(rowOffset, numRows);
     this.bufferType = bufferType;
     this.writer = writer;
@@ -61,6 +62,7 @@ class SlicedBufferSerializer implements HostColumnsVisitor {
     this.metrics = metrics;
     this.totalDataLen = 0;
     this.addCopyBufferTime = addCopyBufferTime;
+    this.headerSize = headerSize;
   }
 
   public long getTotalDataLen() {
@@ -68,7 +70,7 @@ public long getTotalDataLen() {
   }
 
   @Override
-  public void visitStruct(HostColumnVectorCore col) {
+  public void preVisitStruct(HostColumnVectorCore col) {
     SliceInfo parent = sliceInfos.peekLast();
 
     try {
@@ -88,6 +90,11 @@ public void visitStruct(HostColumnVectorCore col) {
     }
   }
 
+  @Override
+  public void visitStruct(HostColumnVectorCore col) {
+    // NOOP
+  }
+
   @Override
   public void preVisitList(HostColumnVectorCore col) {
     SliceInfo parent = sliceInfos.getLast();
@@ -156,12 +163,32 @@ public void visit(HostColumnVectorCore col) {
     }
   }
 
+  @Override
+  public void done() {
+    try {
+      switch (bufferType) {
+        case VALIDITY:
+          totalDataLen = padForValidityAlignment(writer, totalDataLen, headerSize);
+          return;
+        case OFFSET:
+          // fallthrough
+        case DATA:
+          totalDataLen = padForHostAlignment(writer, totalDataLen);
+          return;
+        default:
+          throw new IllegalArgumentException("Unexpected buffer type: " + bufferType);
+      }
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   private long copySlicedValidity(HostColumnVectorCore column, SliceInfo sliceInfo)
       throws IOException {
     if (column.getValidity() != null && sliceInfo.getRowCount() > 0) {
       HostMemoryBuffer buff = column.getValidity();
       long len = sliceInfo.getValidityBufferInfo().getBufferLength();
-      return copyBufferAndPadForHost(buff, sliceInfo.getValidityBufferInfo().getBufferOffset(), len);
+      return copyBuffer(buff, sliceInfo.getValidityBufferInfo().getBufferOffset(), len);
     } else {
       return 0;
     }
@@ -186,7 +213,7 @@ private long copySlicedOffset(HostColumnVectorCore column, SliceInfo sliceInfo)
       }
     }
 
-    return copyBufferAndPadForHost(column.getOffsets(), srcOffset, bytesToCopy);
+    return copyBuffer(column.getOffsets(), srcOffset, bytesToCopy);
   }
 
   private long copySlicedData(HostColumnVectorCore column, SliceInfo sliceInfo) throws IOException {
@@ -205,12 +232,12 @@ private long copySlicedData(HostColumnVectorCore column, SliceInfo sliceInfo) th
 
           return 0;
         } else {
-          return copyBufferAndPadForHost(column.getData(), startByteOffset, bytesToCopy);
+          return copyBuffer(column.getData(), startByteOffset, bytesToCopy);
         }
       } else if (type.getSizeInBytes() > 0) {
         long bytesToCopy = sliceInfo.rowCount * type.getSizeInBytes();
         long srcOffset = sliceInfo.offset * type.getSizeInBytes();
-        return copyBufferAndPadForHost(column.getData(), srcOffset, bytesToCopy);
+        return copyBuffer(column.getData(), srcOffset, bytesToCopy);
       } else {
         return 0;
       }
@@ -219,17 +246,16 @@ private long copySlicedData(HostColumnVectorCore column, SliceInfo sliceInfo) th
     }
   }
 
-  private long copyBufferAndPadForHost(HostMemoryBuffer buffer, long offset, long length)
+  private long copyBuffer(HostMemoryBuffer buffer, long offset, long length)
       throws IOException {
     if (addCopyBufferTime) {
       long now = System.nanoTime();
       writer.copyDataFrom(buffer, offset, length);
-      long ret = padForHostAlignment(writer, length);
       metrics.addCopyBufferTime(System.nanoTime() - now);
-      return ret;
+      return length;
     } else {
       writer.copyDataFrom(buffer, offset, length);
-      return padForHostAlignment(writer, length);
+      return length;
     }
   }
 }
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/kudo/TableBuilder.java b/src/main/java/com/nvidia/spark/rapids/jni/kudo/TableBuilder.java
index d33e1bb8ec..4b85afb9ed 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/kudo/TableBuilder.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/kudo/TableBuilder.java
@@ -68,14 +68,20 @@ public Table visitTopSchema(Schema schema, List<ColumnView> children) {
   }
 
   @Override
-  public ColumnView visitStruct(Schema structType, List<ColumnView> children) {
+  public ColumnViewInfo preVisitStruct(Schema structType) {
     ColumnViewInfo colViewInfo = getCurrentColumnViewInfo();
 
+    curViewInfoIdx += 1;
+    return colViewInfo;
+  }
+
+  @Override
+  public ColumnView visitStruct(Schema structType, ColumnViewInfo colViewInfo, List<ColumnView> children) {
     ColumnView[] childrenView = children.toArray(new ColumnView[0]);
+
     ColumnView columnView = colViewInfo.buildColumnView(buffer, childrenView);
     columnViewList[curViewIdx] = columnView;
     curViewIdx += 1;
-    curViewInfoIdx += 1;
     return columnView;
   }
 
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/schema/HostColumnsVisitor.java b/src/main/java/com/nvidia/spark/rapids/jni/schema/HostColumnsVisitor.java
index 9d296a25bb..d9ed0d0552 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/schema/HostColumnsVisitor.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/schema/HostColumnsVisitor.java
@@ -37,9 +37,10 @@
  *
  * The order of visiting will be:
  * <ol>
+ *     <li> Previsit struct column A</li>
  *     <li> Visit primitive column a1 </li>
  *     <li> Visit primitive column a2</li>
- *     <li> Visit struct column A, with results from a1 and a2</li>
+ *     <li> Visit struct column A</li>
  *     <li> Previsit list column B</li>
  *     <li> Visit primitive column b1</li>
  *     <li> Visit list column B</li>
@@ -50,31 +51,39 @@
  *
  */
 public interface HostColumnsVisitor {
+    /**
+     * Visit a struct column before any of its children.
+     * @param col the struct column to visit
+     */
+    void preVisitStruct(HostColumnVectorCore col);
+
+    // TODO no one uses this, do we even want it???
     /**
      * Visit a struct column.
      * @param col the struct column to visit
-     * @return the result of visiting the struct column
      */
     void visitStruct(HostColumnVectorCore col);
 
     /**
      * Visit a list column before actually visiting its child.
      * @param col the list column to visit
-     * @return the result of visiting the list column
      */
     void preVisitList(HostColumnVectorCore col);
 
     /**
      * Visit a list column after visiting its child.
      * @param col the list column to visit
-     * @return the result of visiting the list column
      */
     void visitList(HostColumnVectorCore col);
 
     /**
      * Visit a column that is a primitive type.
      * @param col the column to visit
-     * @return the result of visiting the column
      */
     void visit(HostColumnVectorCore col);
+
+    /**
+     * The processing is all done
+     */
+    void done();
 }
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/schema/SchemaVisitor.java b/src/main/java/com/nvidia/spark/rapids/jni/schema/SchemaVisitor.java
index c6b33e0fb4..d637424173 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/schema/SchemaVisitor.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/schema/SchemaVisitor.java
@@ -1,6 +1,6 @@
 /*
  *
- *  Copyright (c) 2024, NVIDIA CORPORATION.
+ *  Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  *  Licensed under the Apache License, Version 2.0 (the "License");
  *  you may not use this file except in compliance with the License.
@@ -87,13 +87,21 @@ public interface SchemaVisitor<T, P, R> {
      */
     R visitTopSchema(Schema schema, List<T> children);
 
+    /**
+     * Visit a struct schema before actually visiting its children.
+     * @param structType the struct schema to visit
+     * @return the result of visiting the struct schema
+     */
+    P preVisitStruct(Schema structType);
+
     /**
      * Visit a struct schema.
      * @param structType the struct schema to visit
+     * @param preVisitResult the result of visiting the struct schema before visiting its children
      * @param children the results of visiting the children
      * @return the result of visiting the struct schema
      */
-    T visitStruct(Schema structType, List<T> children);
+    T visitStruct(Schema structType, P preVisitResult, List<T> children);
 
     /**
      * Visit a list schema before actually visiting its child.
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/schema/SimpleSchemaVisitor.java b/src/main/java/com/nvidia/spark/rapids/jni/schema/SimpleSchemaVisitor.java
index 4a4d8e981b..9de73ca254 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/schema/SimpleSchemaVisitor.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/schema/SimpleSchemaVisitor.java
@@ -34,6 +34,12 @@ public interface SimpleSchemaVisitor {
      */
     void visitTopSchema(Schema schema);
 
+    /**
+     * Visit a struct schema before actually visiting its children.
+     * @param structType the struct schema to visit
+     */
+    void preVisitStruct(Schema structType);
+
     /**
      * Visit a struct schema.
      * @param structType the struct schema to visit
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/schema/Visitors.java b/src/main/java/com/nvidia/spark/rapids/jni/schema/Visitors.java
index 1329f2f381..fafa797bc0 100644
--- a/src/main/java/com/nvidia/spark/rapids/jni/schema/Visitors.java
+++ b/src/main/java/com/nvidia/spark/rapids/jni/schema/Visitors.java
@@ -66,6 +66,7 @@ public static void visitSchema(Schema schema, SimpleSchemaVisitor visitor) {
     private static void visitSchemaInner(Schema schema, SimpleSchemaVisitor visitor) {
         switch (schema.getType().getTypeId()) {
             case STRUCT:
+                visitor.preVisitStruct(schema);
                 for (int i=0; i<schema.getNumChildren(); i++) {
                     visitSchemaInner(schema.getChild(i), visitor);
                 }
@@ -84,15 +85,18 @@ private static void visitSchemaInner(Schema schema, SimpleSchemaVisitor visitor)
 
     private static <T, P, R> T visitSchemaInner(Schema schema, SchemaVisitor<T, P, R> visitor) {
         switch (schema.getType().getTypeId()) {
-            case STRUCT:
+            case STRUCT: {
+                P preVisitResult = visitor.preVisitStruct(schema);
                 List<T> children = IntStream.range(0, schema.getNumChildren())
-                        .mapToObj(childIdx -> visitSchemaInner(schema.getChild(childIdx), visitor))
-                        .collect(Collectors.toList());
-                return visitor.visitStruct(schema, children);
-            case LIST:
+                    .mapToObj(childIdx -> visitSchemaInner(schema.getChild(childIdx), visitor))
+                    .collect(Collectors.toList());
+                return visitor.visitStruct(schema, preVisitResult, children);
+            }
+            case LIST: {
                 P preVisitResult = visitor.preVisitList(schema);
                 T childResult = visitSchemaInner(schema.getChild(0), visitor);
                 return visitor.visitList(schema, preVisitResult, childResult);
+            }
             default:
                 return visitor.visit(schema);
         }
@@ -113,12 +117,13 @@ public static void visitColumns(HostColumnVector[] cols,
         for (HostColumnVector col : cols) {
             visitColumn(col, visitor);
         }
-
+        visitor.done();
     }
 
     private static void visitColumn(HostColumnVectorCore col, HostColumnsVisitor visitor) {
         switch (col.getType().getTypeId()) {
             case STRUCT:
+                visitor.preVisitStruct(col);
                 for (int i=0; i<col.getNumChildren(); i++) {
                     visitColumn(col.getChildColumnView(i), visitor);
                 }
diff --git a/src/test/java/com/nvidia/spark/rapids/jni/kudo/KudoGpuSerializerTest.java b/src/test/java/com/nvidia/spark/rapids/jni/kudo/KudoGpuSerializerTest.java
new file mode 100644
index 0000000000..a14f61174e
--- /dev/null
+++ b/src/test/java/com/nvidia/spark/rapids/jni/kudo/KudoGpuSerializerTest.java
@@ -0,0 +1,522 @@
+/*
+ * Copyright (c) 2025, 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 com.nvidia.spark.rapids.jni.kudo;
+
+import ai.rapids.cudf.*;
+import org.junit.jupiter.api.Test;
+
+import java.io.*;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+import static ai.rapids.cudf.AssertUtils.assertTablesAreEqual;
+import static com.nvidia.spark.rapids.jni.kudo.KudoSerializerTest.*;
+import static org.junit.jupiter.api.Assertions.*;
+
+public class KudoGpuSerializerTest {
+  /**
+   * Used in some performance tests.
+   */
+  private static class JustCountOutputStream extends OutputStream {
+    long length = 0;
+
+    @Override
+    public void write(byte[] i) {
+      this.length += i.length;
+    }
+
+    @Override
+    public void write(byte[] i, int start, int length) {
+      // Just assume that this is all correct for now...
+      this.length += length;
+    }
+
+    @Override
+    public void write(int i) {
+      this.length += 1;
+    }
+  }
+
+  private static int readIntHeader(int start, byte[] data) {
+    int b1 = start < data.length ? data[start] & 0xFF : 0;
+    int b2 = start + 1 < data.length ? data[start + 1] & 0xFF : 0;
+    int b3 = start + 2 < data.length ? data[start + 2] & 0xFF : 0;
+    int b4 = start + 3 < data.length ? data[start + 3] & 0xFF : 0;
+    return (b1 << 24) + (b2 << 16) + (b3 << 8) + b4;
+  }
+
+  private static int readIntOffset(int start, byte[] data) {
+    int b1 = start < data.length ? data[start] & 0xFF : 0;
+    int b2 = start + 1 < data.length ? data[start + 1] & 0xFF : 0;
+    int b3 = start + 2 < data.length ? data[start + 2] & 0xFF : 0;
+    int b4 = start + 3 < data.length ? data[start + 3] & 0xFF : 0;
+    return b1 + (b2 << 8) + (b3 << 16) + (b4 << 24);
+  }
+
+  private static String makeExtraEnd(String name, int start, byte[] hDataGPU, byte[] hDataCPU) {
+    int gpuNum = readIntHeader(start, hDataGPU);
+    int cpuNum = readIntHeader(start, hDataCPU);
+    return " <-- " + name + " END (" + gpuNum + "/" + cpuNum + ")";
+  }
+
+  private static String makeOffsetEnd(int index, int start, byte[] hDataGPU, byte[] hDataCPU) {
+    int gpuNum = readIntOffset(start, hDataGPU);
+    int cpuNum = readIntOffset(start, hDataCPU);
+    return " <-- OFFSET " + index + " (" + gpuNum + "/" + cpuNum + ")";
+  }
+
+  public static int logSinglePartitionComparison(String name, int startOffset, byte[] hDataGPU, byte[] hDataCPU) {
+    System.err.println(name + " COMP GPU VS CPU AT START OFFSET "+ startOffset);
+    int maxLength = Math.max(hDataGPU.length, hDataCPU.length) - startOffset;
+    int hasValidStart = -1;
+    int hasValidEnd = -1;
+    int validityBuffersStart = -1;
+    int validityBuffersEnd = -1;
+    int offsetBuffersStart = -1;
+    int offsetBuffersEnd = -1;
+    int dataBuffersStart = -1;
+    int dataBuffersEnd = -1;
+    for (int i = 0; i < maxLength; i++) {
+      Byte gpuByte = null;
+      Byte cpuByte = null;
+      String gpu = "N/A     ";
+      if (i + startOffset < hDataGPU.length) {
+        gpuByte = hDataGPU[i + startOffset];
+        gpu = String.format("0x%02X %03d", hDataGPU[i + startOffset] & 0xFF, hDataGPU[i + startOffset]);
+      }
+
+      String cpu = "N/A     ";
+      if (i + startOffset < hDataCPU.length) {
+        cpuByte = hDataCPU[i + startOffset];
+        cpu = String.format("0x%02X %03d", hDataCPU[i + startOffset] & 0xFF, hDataCPU[i + startOffset]);
+      }
+
+      String diff = "   ";
+      if (cpuByte != gpuByte) {
+        diff = "***";
+      }
+      String extra = "";
+      if (i == 0) {
+        extra = " <-- MAGIC START";
+      } else if (i == 3) {
+        extra = " <-- MAGIC END";
+      } else if (i == 4) {
+        extra = " <-- OFFSET START";
+      } else if (i == 7) {
+        extra = makeExtraEnd("OFFSET", 4 + startOffset, hDataGPU, hDataCPU);
+      } else if (i == 8) {
+        extra = " <-- NUM_ROWS START";
+      } else if (i == 11) {
+        extra = makeExtraEnd("NUM_ROWS", 8 + startOffset, hDataGPU, hDataCPU);
+      } else if (i == 12) {
+        extra = " <-- VALIDITY_BUF_LEN START";
+      } else if (i == 15) {
+        extra = makeExtraEnd("VALIDITY_BUF_LEN", 12 + startOffset, hDataGPU, hDataCPU);
+      } else if (i == 16) {
+        extra = " <-- OFFSET_BUF_LEN START";
+      } else if (i == 19) {
+        extra = makeExtraEnd("OFFSET_BUF_LEN", 16 + startOffset, hDataGPU, hDataCPU);
+      } else if (i == 20) {
+        extra = " <-- TOTAL_DATA_LEN START";
+      } else if (i == 23) {
+        extra = makeExtraEnd("TOTAL_DATA_LEN", 20 + startOffset, hDataGPU, hDataCPU);
+      } else if (i == 24) {
+        extra = " <-- NUM_COL START";
+      } else if (i == 27) {
+        extra = makeExtraEnd("NUM_COL", 24 + startOffset, hDataGPU, hDataCPU);
+        int numCol = readIntHeader(24 + startOffset, hDataCPU);
+        int hasValidLen = (numCol + 7) / 8;
+        if (hasValidLen > 0) {
+          hasValidStart = 28;
+          hasValidEnd = hasValidStart + hasValidLen - 1;
+        } else {
+          hasValidEnd = 28;
+        }
+        int validityBuffsLen = readIntHeader(12 + startOffset, hDataCPU);
+        if (validityBuffsLen > 0) {
+          validityBuffersStart = hasValidEnd + 1;
+          validityBuffersEnd = validityBuffersStart + validityBuffsLen - 1;
+        } else {
+          validityBuffersEnd = hasValidEnd;
+        }
+        int offsetBuffersLen = readIntHeader(16 + startOffset, hDataCPU);
+        if (offsetBuffersLen > 0) {
+          offsetBuffersStart = validityBuffersEnd + 1;
+          offsetBuffersEnd = offsetBuffersStart + offsetBuffersLen - 1;
+        } else {
+          offsetBuffersEnd = validityBuffersEnd;
+        }
+        int totalDataLen = readIntHeader(20 + startOffset, hDataCPU);
+        int dataLen = totalDataLen - offsetBuffersLen - validityBuffsLen;
+        if (dataLen > 0) {
+          dataBuffersStart = offsetBuffersEnd + 1;
+          dataBuffersEnd = dataBuffersStart + dataLen - 1;
+        } else {
+          dataBuffersEnd = offsetBuffersEnd;
+        }
+      } else if (hasValidStart == hasValidEnd && i == hasValidStart) {
+        extra = " <-- HAS_VALID_BUF";
+      } else if (hasValidStart == i) {
+        extra = " <-- HAS_VALID_BUF START";
+      } else if (hasValidStart > 0 && hasValidEnd == i) {
+        extra = " <-- HAS_VALID_BUF END";
+      } else if (validityBuffersStart == validityBuffersEnd && i == validityBuffersStart) {
+        extra = " <-- VALIDITY_BUFFERS";
+      } else if (validityBuffersStart == i) {
+        extra = " <-- VALIDITY_BUFFERS START";
+      } else if (validityBuffersStart > 0 && validityBuffersEnd == i) {
+        extra = " <-- VALIDITY_BUFFERS END";
+      } else if (offsetBuffersStart == i) {
+        extra = " <-- OFFSET_BUFFERS START";
+      } else if (offsetBuffersStart > 0 && i > offsetBuffersStart && i <= offsetBuffersEnd && (i - offsetBuffersStart + 1) % 4 == 0) {
+        extra = makeOffsetEnd(((i - offsetBuffersStart) / 4), i - 3 + startOffset, hDataGPU, hDataCPU);
+      } else if (dataBuffersStart == dataBuffersEnd && i == dataBuffersStart) {
+        extra = " <-- DATA_BUFFERS";
+      } else if (dataBuffersStart == i) {
+        extra = " <-- DATA_BUFFERS START";
+      } else if (dataBuffersStart > 0 && dataBuffersEnd == i) {
+        extra = " <-- DATA_BUFFERS END";
+      }
+      String i_str = String.format("%05d: ", i + startOffset);
+      String aligned = "        ";
+      if ((i + startOffset) % 4 == 0) {
+        aligned = "ALIGNED ";
+      }
+      System.err.println(aligned + i_str + diff + " " + gpu + " VS " + cpu  + extra);
+      if (i == dataBuffersEnd) {
+        return i + startOffset + 1;
+      }
+    }
+    return maxLength + startOffset + 1;
+  }
+
+  public static void logPartitionComparisons(String name, byte[] hDataGPU, byte[] hDataCPU) {
+    int maxLength = Math.max(hDataGPU.length, hDataCPU.length);
+    int offset = 0;
+    while (offset < maxLength) {
+      offset = logSinglePartitionComparison(name, offset, hDataGPU, hDataCPU);
+    }
+  }
+
+  public static int[] calcEvenSlices(Table table, int numSlices) {
+    assert table.getRowCount() >= numSlices : "slices must be <= number of rows " + table.getRowCount() + " < " + numSlices;
+    assert numSlices > 0 : "slices must be > 0 " + numSlices;
+    if (numSlices == 1) {
+      return new int[0];
+    } else {
+      int rowsPerSlice = (int)table.getRowCount() / numSlices;
+      int[] ret = new int[numSlices - 1];
+      int offset = 0;
+      for (int i = 0; i < ret.length; i++) {
+        offset += rowsPerSlice;
+        ret[i] = offset;
+      }
+      return ret;
+    }
+  }
+
+  public static byte[] writeGPU(Table table, int[] slices) {
+    DeviceMemoryBuffer[] buffers = KudoGpuSerializer.splitAndSerializeToDevice(table, slices);
+    try {
+      assertEquals(2, buffers.length);
+      DeviceMemoryBuffer data = buffers[0];
+      // Ignoring the offsets as we are not splitting it up...
+      try (HostMemoryBuffer hostBuffer = HostMemoryBuffer.allocate(data.getLength())) {
+        hostBuffer.copyFromDeviceBuffer(data);
+        assert hostBuffer.getLength() <= Integer.MAX_VALUE : "Serialized buffer is too large for test " + hostBuffer.getLength();
+        byte[] ret = new byte[(int)hostBuffer.getLength()];
+        hostBuffer.getBytes(ret, 0, 0, ret.length);
+        return ret;
+      }
+    } finally {
+      for (DeviceMemoryBuffer b : buffers) {
+        if (b != null) {
+          b.close();
+        }
+      }
+    }
+  }
+
+  public static void writeCPUCore(HostColumnVector[] columns, KudoSerializer serializer,
+                           int [] slices, int numRows, OutputStream out) {
+    int start = 0;
+    int end;
+    for (int i = 0; i < slices.length + 1; i++) {
+      if (i >= slices.length) {
+        end = numRows;
+      } else {
+        end = slices[i];
+      }
+      int len = end - start;
+      serializer.writeToStreamWithMetrics(columns, out, start, len);
+      start = end;
+    }
+  }
+
+  public static byte[] writeCPU(Table table, int[] slices) {
+    int numRows = (int)table.getRowCount();
+    int numColumns = table.getNumberOfColumns();
+    HostColumnVector[] columns = new HostColumnVector[numColumns];
+    try {
+      for (int c = 0; c < numColumns; c++) {
+        columns[c] = table.getColumn(c).copyToHost();
+      }
+      Schema s = KudoSerializerTest.schemaOf(table);
+      KudoSerializer serializer = new KudoSerializer(s);
+      ByteArrayOutputStream baos = new ByteArrayOutputStream();
+      writeCPUCore(columns, serializer, slices, numRows, baos);
+      return baos.toByteArray();
+    } finally {
+      for (HostColumnVector hcv: columns) {
+        if (hcv != null) {
+          hcv.close();
+        }
+      }
+    }
+  }
+
+  public void debugCompareWrite(String name, Table table, int[] slices) {
+    byte[] gpuData = writeGPU(table, slices);
+    byte[] cpuData = writeCPU(table, slices);
+    if (!Arrays.equals(gpuData, cpuData)) {
+      logPartitionComparisons(name, gpuData, cpuData);
+    }
+  }
+
+  public Table readGPU(String name, Schema s, byte[] bytes) throws IOException {
+    // We are going to parse the headers so we can set up the offsets properly.
+    DataInputStream din = new DataInputStream(new ByteArrayInputStream(bytes));
+    ArrayList<Long> offsets = new ArrayList<>();
+    long currentOffset = 0;
+    Optional<KudoTableHeader> header;
+    do {
+      header = KudoTableHeader.readFrom(din);
+      if (header.isPresent()) {
+        offsets.add(currentOffset);
+        KudoTableHeader h = header.get();
+        din.skipBytes(h.getTotalDataLen());
+        currentOffset += h.getTotalDataLen() + h.getSerializedSize();
+      }
+    } while (header.isPresent());
+    if (!offsets.isEmpty()) {
+      offsets.add(currentOffset);
+    }
+    // Okay now we have offsets, lets put all of this onto the GPU.
+    try (HostMemoryBuffer hostOffsets = HostMemoryBuffer.allocate(8L * offsets.size());
+         HostMemoryBuffer hostPayload = HostMemoryBuffer.allocate(bytes.length)) {
+      for (int i = 0; i < offsets.size(); i++) {
+        hostOffsets.setLong(i * 8L, offsets.get(i));
+      }
+      hostPayload.setBytes(0, bytes, 0, bytes.length);
+      try (DeviceMemoryBuffer devOffsets = DeviceMemoryBuffer.allocate(hostOffsets.getLength());
+           DeviceMemoryBuffer devPayload = DeviceMemoryBuffer.allocate(hostPayload.getLength())) {
+        devOffsets.copyFromHostBuffer(hostOffsets);
+        devPayload.copyFromHostBuffer(hostPayload);
+        return KudoGpuSerializer.assembleFromDeviceRaw(s, devPayload, devOffsets);
+      }
+    }
+  }
+
+  public Table readCPU(String name, Schema s, byte[] bytes) throws Exception {
+    KudoSerializer serializer = new KudoSerializer(s);
+    DataInputStream din = new DataInputStream(new ByteArrayInputStream(bytes));
+    ArrayList<KudoTable> kudoTables = new ArrayList<>();
+    Optional<KudoTable> kt;
+    do {
+      kt = KudoTable.from(din);
+      kt.ifPresent(kudoTables::add);
+    } while(kt.isPresent());
+    return serializer.mergeToTable(kudoTables.toArray(new KudoTable[kudoTables.size()]));
+  }
+
+  public void testRoundTrip(String name, Table table, int[] slices) throws Exception {
+    debugCompareWrite(name, table, slices);
+    testCPUOnlyRoundTrip(name + " CPU -> CPU", table, slices);
+    testGPUOnlyRoundTrip(name + " GPU -> GPU", table, slices);
+    testCPUWriteGPURead(name + " CPU -> GPU", table, slices);
+    testGPUWriteCPURead(name + " GPU -> CPU", table, slices);
+  }
+
+  public void testCPUOnlyRoundTrip(String name, Table table, int[] slices) throws Exception {
+    Schema s = KudoSerializerTest.schemaOf(table);
+    byte[] data = writeCPU(table, slices);
+    try (Table t = readCPU(name, s, data)) {
+      assertTablesAreEqual(table, t);
+    }
+  }
+
+  public void testGPUOnlyRoundTrip(String name, Table table, int[] slices) throws Exception {
+    Schema s = KudoSerializerTest.schemaOf(table);
+    byte[] data = writeGPU(table, slices);
+    try (Table t = readGPU(name, s, data)) {
+      assertTablesAreEqual(table, t);
+    }
+  }
+
+  public void testCPUWriteGPURead(String name, Table table, int[] slices) throws Exception {
+    Schema s = KudoSerializerTest.schemaOf(table);
+    byte[] data = writeCPU(table, slices);
+    try (Table t = readGPU(name, s, data)) {
+      assertTablesAreEqual(table, t);
+    }
+  }
+
+  public void testGPUWriteCPURead(String name, Table table, int[] slices) throws Exception {
+    Schema s = KudoSerializerTest.schemaOf(table);
+    byte[] data = writeGPU(table, slices);
+    try (Table t = readCPU(name, s, data)) {
+      assertTablesAreEqual(table, t);
+    }
+  }
+
+  @Test
+  public void testSimpleRoundTrip() throws Exception {
+    try (Table table = new Table.TestBuilder()
+        .column(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+            11, 12, 13, 14, 15, 16, 17, 18, 19, 20)
+        .build()) {
+      for (int numSlices = 1; numSlices < table.getRowCount(); numSlices++) {
+        System.err.println("TEST WITH "+ numSlices);
+        int[] slices = calcEvenSlices(table, numSlices);
+        testRoundTrip("simple", table, slices);
+      }
+    }
+  }
+
+  @Test
+  public void testEmptyStructRoundTrip() throws Exception {
+    try (Table table = KudoSerializerTest.buildEmptyStructTable()) {
+      for (int numSlices = 1; numSlices < table.getRowCount(); numSlices++) {
+        System.err.println("TEST WITH "+ numSlices);
+        int[] slices = calcEvenSlices(table, numSlices);
+        testRoundTrip("empty struct", table, slices);
+      }
+    }
+  }
+
+  public static Table buildMediumTable() {
+    HostColumnVector.StructType st = new HostColumnVector.StructType(
+        true,
+        new HostColumnVector.BasicType(true, DType.INT8),
+        new HostColumnVector.BasicType(true, DType.INT64)
+    );
+    return new Table.TestBuilder()
+        .column("1", null, "34", "45", "56", "67")
+        .column(new Integer[]{null},
+            new Integer[]{4},
+            new Integer[]{7},
+            null,
+            new Integer[]{},
+            new Integer[]{})
+        .column(st, new HostColumnVector.StructData(null, 11L),
+            new HostColumnVector.StructData((byte) 2, null),
+            new HostColumnVector.StructData((byte) 3, 33L),
+            new HostColumnVector.StructData((byte) 4, 44L),
+            new HostColumnVector.StructData((byte) 5, 55L),
+            null)
+        .column(null, 2, 3, 4, 5, 6)
+        .build();
+  }
+
+  @Test
+  public void testMediumRoundTrip() throws Exception {
+    try (Table table = buildMediumTable()) {
+      for (int numSlices = 1; numSlices < table.getRowCount(); numSlices++) {
+        System.err.println("TEST WITH "+ numSlices);
+        int[] slices = calcEvenSlices(table, numSlices);
+        testRoundTrip("medium", table, slices);
+      }
+    }
+  }
+
+  public static Table buildHalfEmptyStructTable() {
+    HostColumnVector.StructType st = new HostColumnVector.StructType(
+        true,
+        new HostColumnVector.BasicType(true, DType.INT32)
+    );
+    HostColumnVector.ListType lt = new HostColumnVector.ListType(true, st);
+    return new Table.TestBuilder()
+        .column(lt, structs(struct(1), struct(2), null),
+            structs(struct(4), struct(5), struct(6)),
+            structs(struct(7), struct(8), struct(9)),
+            null,
+            structs(),
+            structs())
+        .build();
+  }
+
+  @Test
+  public void testHalfEmptyStructRoundTrip() throws Exception {
+    try (Table table = buildHalfEmptyStructTable()) {
+      for (int numSlices = 1; numSlices < table.getRowCount(); numSlices++) {
+        System.err.println("TEST WITH "+ numSlices);
+        int[] slices = calcEvenSlices(table, numSlices);
+        testRoundTrip("half empty struct", table, slices);
+      }
+    }
+  }
+
+  @Test
+  public void testSimpleMultiColumnRoundTrip() throws Exception {
+    try (Table table = new Table.TestBuilder()
+        .column(null, (byte)0xF0, (byte)0x0F, (byte)0xAA, null)
+        .column((short)0xFFFF, (short)0xF0F0, null, (short)0xAAAA, (short)0x5555)
+        .column("A","B","C","D",null)
+        .column("0xFF", null, "0x0F", "0xAA", "0x55")
+        .build()) {
+      for (int numSlices = 1; numSlices < table.getRowCount(); numSlices++) {
+        System.err.println("TEST WITH "+ numSlices);
+        int[] slices = calcEvenSlices(table, numSlices);
+        testRoundTrip("simple multicolumn", table, slices);
+      }
+    }
+  }
+
+  @Test
+  public void testComplexRoundTrip() throws Exception {
+    try (Table table = KudoSerializerTest.buildTestTable()) {
+      int numSlices = 2;
+      //for (int numSlices = 1; numSlices < table.getRowCount(); numSlices++) {
+        System.err.println("TEST WITH "+ numSlices);
+        int[] slices = calcEvenSlices(table, numSlices);
+        testRoundTrip("complex", table, slices);
+      //}
+    }
+  }
+
+  static Table buildStringListTable() {
+    return new Table.TestBuilder()
+        .column(strings("*"), strings("*"), strings("****"),
+            strings("", "*", null))
+        .column(strings(null, null, null, null), strings(),
+            strings(null, null, null), strings())
+        .build();
+  }
+
+  @Test
+  public void testStringListRoundTrip() throws Exception {
+    try (Table table = buildStringListTable()) {
+      for (int numSlices = 1; numSlices < table.getRowCount(); numSlices++) {
+        System.err.println("TEST WITH "+ numSlices);
+        int[] slices = calcEvenSlices(table, numSlices);
+        testRoundTrip("string list table", table, slices);
+      }
+    }
+  }
+}
diff --git a/src/test/java/com/nvidia/spark/rapids/jni/kudo/KudoSerializerTest.java b/src/test/java/com/nvidia/spark/rapids/jni/kudo/KudoSerializerTest.java
index c614209099..e7b8a62f96 100644
--- a/src/test/java/com/nvidia/spark/rapids/jni/kudo/KudoSerializerTest.java
+++ b/src/test/java/com/nvidia/spark/rapids/jni/kudo/KudoSerializerTest.java
@@ -23,6 +23,8 @@
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.DataInputStream;
+import java.math.BigInteger;
+import java.math.RoundingMode;
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.OutputStream;
@@ -44,7 +46,7 @@ public class KudoSerializerTest {
   public void testSerializeAndDeserializeTable() {
     try(Table expected = buildTestTable()) {
       int rowCount = toIntExact(expected.getRowCount());
-      for (int sliceSize = 1; sliceSize <= rowCount; sliceSize++) {
+      for (int sliceSize = rowCount; sliceSize >= 1; sliceSize--) {
         List<TableSlice> tableSlices = new ArrayList<>();
         for (int startRow = 0; startRow < rowCount; startRow += sliceSize) {
           tableSlices.add(new TableSlice(startRow, Math.min(sliceSize, rowCount - startRow), expected));
@@ -74,6 +76,23 @@ public void testRowCountOnly() throws Exception {
     assertEquals(0, header.getTotalDataLen());
   }
 
+  @Test
+  public void testSerializeAndDeserializeEmptyStructTable() {
+    try(Table expected = buildEmptyStructTable()) {
+      int rowCount = toIntExact(expected.getRowCount());
+      for (int sliceSize = rowCount; sliceSize >= 1; sliceSize--) {
+        List<TableSlice> tableSlices = new ArrayList<>();
+        for (int startRow = 0; startRow < rowCount; startRow += sliceSize) {
+          tableSlices.add(new TableSlice(startRow, Math.min(sliceSize, rowCount - startRow), expected));
+        }
+
+        checkMergeTable(expected, tableSlices);
+      }
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   @Test
   public void testWriteSimple() throws Exception {
     KudoSerializer serializer = new KudoSerializer(buildSimpleTestSchema());
@@ -81,7 +100,8 @@ public void testWriteSimple() throws Exception {
     try (Table t = buildSimpleTable()) {
       OpenByteArrayOutputStream out = new OpenByteArrayOutputStream();
       long bytesWritten = serializer.writeToStreamWithMetrics(t, out, 0, 4).getWrittenBytes();
-      assertEquals(189, bytesWritten);
+
+      assertEquals(172, bytesWritten);
 
       ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
 
@@ -89,9 +109,9 @@ public void testWriteSimple() throws Exception {
       assertEquals(7, header.getNumColumns());
       assertEquals(0, header.getOffset());
       assertEquals(4, header.getNumRows());
-      assertEquals(24, header.getValidityBufferLen());
+      assertEquals(7, header.getValidityBufferLen());
       assertEquals(40, header.getOffsetBufferLen());
-      assertEquals(160, header.getTotalDataLen());
+      assertEquals(143, header.getTotalDataLen());
 
       // First integer column has no validity buffer
       assertFalse(header.hasValidityBuffer(0));
@@ -305,7 +325,7 @@ private static Schema buildSimpleTestSchema() {
     return builder.build();
   }
 
-  private static Table buildSimpleTable() {
+  static Table buildSimpleTable() {
     HostColumnVector.StructType st = new HostColumnVector.StructType(
         true,
         new HostColumnVector.BasicType(true, DType.INT8),
@@ -321,7 +341,27 @@ private static Table buildSimpleTable() {
         .build();
   }
 
-  private static Table buildTestTable() {
+  static Table buildEmptyStructTable() {
+    HostColumnVector.StructType st = new HostColumnVector.StructType(true);
+    return new Table.TestBuilder()
+        .column(st,
+            struct(), null, null, struct(), null, null, struct(), struct(),
+            null, struct(), struct(), null, struct(), struct(), null, null,
+            struct(), null, null, struct(), null, null, struct(), struct(),
+            null, struct(), struct(), null, struct(), struct(), null, null,
+            struct(), struct(), null, struct(), null, null, struct(), struct(),
+            null, struct(), struct(), null, struct(), null, null, null,
+            struct(), null, null, struct(), null, struct(), struct(), null,
+            null, struct(), struct(), null, struct(), struct(), null, null,
+            struct(), null, null, struct(), null, null, struct(), struct(),
+            null, struct(), struct(), null, struct(), struct(), null, null,
+            struct(), null, null, null, null, null, struct(), struct(),
+            null, struct(), struct(), null, struct(), struct(), null, null,
+            struct())
+        .build();
+  }
+
+  static Table buildTestTable() {
     HostColumnVector.ListType listMapType = new HostColumnVector.ListType(true,
         new HostColumnVector.ListType(true,
             new HostColumnVector.StructType(true,
@@ -363,6 +403,11 @@ private static Table buildTestTable() {
             null, 5, null, 7, null, 9, null, 11, null, 13, null, 15)
         .decimal64Column(-8, 1L, null, 1001L, 50L, -2000L, null, 1L, 2L, 3L,
             4L, null, 6L, 7L, 8L, 9L, null, 11L, 12L, 13L, 14L, null)
+        .decimal128Column(-2, RoundingMode.UNNECESSARY, new BigInteger("1"), null, new BigInteger("1001"),
+            new BigInteger("50"), new BigInteger("-2000"), null, new BigInteger("1"), new BigInteger("2"),
+            new BigInteger("3"), new BigInteger("4"), null, new BigInteger("6"), new BigInteger("7"),
+            new BigInteger("8"), new BigInteger("9"), null, new BigInteger("11"), new BigInteger("12"),
+            new BigInteger("13"), new BigInteger("14"), null)
         .column("A", "B", "C", "D", null, "TESTING", "1", "2", "3", "4",
             "5", "6", "7", null, "9", "10", "11", "12", "13", null, "15")
         .column("A", "A", "C", "C", "E", "TESTING", "1", "2", "3", "4", "5",
@@ -496,19 +541,19 @@ private static Integer[] integers(Integer... values) {
     return values;
   }
 
-  private static HostColumnVector.StructData struct(Object... values) {
+  public static HostColumnVector.StructData struct(Object... values) {
     return new HostColumnVector.StructData(values);
   }
 
-  private static List<HostColumnVector.StructData> structs(HostColumnVector.StructData... values) {
+  public static List<HostColumnVector.StructData> structs(HostColumnVector.StructData... values) {
     return asList(values);
   }
 
-  private static String[] strings(String... values) {
+  public static String[] strings(String... values) {
     return values;
   }
 
-  private static Schema schemaOf(Table t) {
+  public static Schema schemaOf(Table t) {
     Schema.Builder builder = Schema.builder();
 
     for (int i = 0; i < t.getNumberOfColumns(); i++) {


Reviews:
user: login: liurenjie1024
body: The serialization format changed a bit, and I'm surprised `KudoSerializerTest` could pass. [This part](https://github.com/NVIDIA/spark-rapids-jni/blob/b48b0c61cd14231738e5f0aa5fe750b1d1e4c400/src/test/java/com/nvidia/spark/rapids/jni/kudo/KudoConcatValidityTest.java#L35) should also change accordingly.
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T07:58:20Z
user: login: hyperbolic2346
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T15:53:00Z
user: login: revans2
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T18:17:23Z
user: login: revans2
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T20:10:17Z
user: login: revans2
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T20:29:35Z
user: login: liurenjie1024
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-07T01:38:11Z
user: login: revans2
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T19:47:25Z
user: login: revans2
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T19:53:04Z
user: login: liurenjie1024
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T08:34:07Z
user: login: liurenjie1024
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T08:35:18Z
user: login: revans2
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T16:38:48Z
user: login: liurenjie1024
body: Just some minor nits, others LGTM! A suggestion before merging is to run integration tests locally since currently we will always update spark-rapids-jni changes.
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T10:03:52Z
user: login: liurenjie1024
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T10:05:57Z
user: login: liurenjie1024
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-13T02:29:35Z

Review comments:
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoSerializer.java
user: login: liurenjie1024
body: This pr modified format, so it would be better to change the documentation of the class
created_at: 2025-05-06T07:07:39Z
author_association: COLLABORATOR
diff_hunk: @@ -193,6 +199,27 @@ public void writeTo(DataWriter dout) throws IOException {
     dout.write(hasValidityBuffer, 0, hasValidityBuffer.length);
   }
 
+  private static void setInt(HostMemoryBuffer hmb, long hmbOffset, int v) {
+    byte[] bytes = new byte[4];
+    bytes[0] = (byte) ((v >>> 24) & 0xFF);
+    bytes[1] = (byte) ((v >>> 16) & 0xFF);
+    bytes[2] = (byte) ((v >>> 8) & 0xFF);
+    bytes[3] = (byte) (v & 0xFF);
+    hmb.setBytes(hmbOffset, bytes, 0, bytes.length);
+  }
+
+  public long writeTo(HostMemoryBuffer hmb, long hmbOffset) {
+    setInt(hmb, hmbOffset, SER_FORMAT_MAGIC_NUMBER);
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeader.java
user: login: liurenjie1024
body: nit: Writing to an  onheap byte array and then copy it to native memory at once maybe a little faster.
created_at: 2025-05-06T07:08:55Z
author_association: COLLABORATOR
diff_hunk: @@ -41,65 +42,75 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   private final byte[] bitset;
   private long validityBufferLen;
   private long offsetBufferLen;
-  private long totalDataLen;
+  private long dataOnlyLen;
   private int nextColIdx;
 
   private Deque<SliceInfo> sliceInfos = new ArrayDeque<>();
 
   KudoTableHeaderCalc(int rowOffset, int numRows, int numFlattenedCols) {
     this.root = new SliceInfo(rowOffset, numRows);
-    this.totalDataLen = 0;
+    this.dataOnlyLen = 0;
     sliceInfos.addLast(this.root);
     this.bitset = new byte[(numFlattenedCols + 7) / 8];
     this.numFlattenedCols = numFlattenedCols;
     this.nextColIdx = 0;
   }
 
   public KudoTableHeader getHeader() {
+    int headerSize = KudoTableHeader.getSerializedSize(bitset.length);
+    // The validity is a bit odd because we want to pad it for 4 byte alignment
+    // But that is relative to the header, not the payload buffer
+    long paddedValiditySize = padForValidityAlignment(validityBufferLen, headerSize);
     return new KudoTableHeader(toIntExact(root.offset),
         toIntExact(root.rowCount),
-        toIntExact(validityBufferLen),
-        toIntExact(offsetBufferLen),
-        toIntExact(totalDataLen),
+        toIntExact(paddedValiditySize),
+        toIntExact(padForHostAlignment(offsetBufferLen)),
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
user: login: liurenjie1024
body: nit: Repeated calculation of `offsetBufferLen`.
created_at: 2025-05-06T07:10:29Z
author_association: COLLABORATOR
diff_hunk: @@ -58,7 +58,7 @@ static SlicedValidityBufferInfo calc(int rowOffset, int numRows) {
     }
     int bufferOffset = rowOffset / 8;
     int beginBit = rowOffset % 8;
-    int bufferLength = 0;
+    int bufferLength = 1;
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/SlicedValidityBufferInfo.java
user: login: liurenjie1024
body: Why this change? When `nowRows == 0`, it should be 0.
created_at: 2025-05-06T07:16:39Z
author_association: COLLABORATOR
diff_hunk: @@ -219,17 +246,11 @@ private long copySlicedData(HostColumnVectorCore column, SliceInfo sliceInfo) th
     }
   }
 
-  private long copyBufferAndPadForHost(HostMemoryBuffer buffer, long offset, long length)
+  private long copyBuffer(HostMemoryBuffer buffer, long offset, long length)
       throws IOException {
-    if (addCopyBufferTime) {
-      long now = System.nanoTime();
-      writer.copyDataFrom(buffer, offset, length);
-      long ret = padForHostAlignment(writer, length);
-      metrics.addCopyBufferTime(System.nanoTime() - now);
-      return ret;
-    } else {
-      writer.copyDataFrom(buffer, offset, length);
-      return padForHostAlignment(writer, length);
-    }
+    long now = System.nanoTime();
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/SlicedBufferSerializer.java
user: login: liurenjie1024
body: This always enable buffer time measure, which may hurt performance. The default behavior is disabling this measurement.
created_at: 2025-05-06T07:41:53Z
author_association: COLLABORATOR
diff_hunk: @@ -41,65 +42,75 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   private final byte[] bitset;
   private long validityBufferLen;
   private long offsetBufferLen;
-  private long totalDataLen;
+  private long dataOnlyLen;
   private int nextColIdx;
 
   private Deque<SliceInfo> sliceInfos = new ArrayDeque<>();
 
   KudoTableHeaderCalc(int rowOffset, int numRows, int numFlattenedCols) {
     this.root = new SliceInfo(rowOffset, numRows);
-    this.totalDataLen = 0;
+    this.dataOnlyLen = 0;
     sliceInfos.addLast(this.root);
     this.bitset = new byte[(numFlattenedCols + 7) / 8];
     this.numFlattenedCols = numFlattenedCols;
     this.nextColIdx = 0;
   }
 
   public KudoTableHeader getHeader() {
+    int headerSize = KudoTableHeader.getSerializedSize(bitset.length);
+    // The validity is a bit odd because we want to pad it for 4 byte alignment
+    // But that is relative to the header, not the payload buffer
+    long paddedValiditySize = padForValidityAlignment(validityBufferLen, headerSize);
     return new KudoTableHeader(toIntExact(root.offset),
         toIntExact(root.rowCount),
-        toIntExact(validityBufferLen),
-        toIntExact(offsetBufferLen),
-        toIntExact(totalDataLen),
+        toIntExact(paddedValiditySize),
+        toIntExact(padForHostAlignment(offsetBufferLen)),
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
user: login: liurenjie1024
body: When merging validity buffer, we assume that validity buffer is at least 4 bytes, see https://github.com/NVIDIA/spark-rapids-jni/blob/44f2c7a8c6702b83226045b792e54ef0ad8bb2ca/src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableMerger.java#L475 

Changes in this pr breaks this assumption. I would suggest to add padding to header so that header is 4 bytes aligned. Other changes in this pr could remain unchanged.
created_at: 2025-05-06T07:55:05Z
author_association: COLLABORATOR
diff_hunk: @@ -41,65 +42,75 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   private final byte[] bitset;
   private long validityBufferLen;
   private long offsetBufferLen;
-  private long totalDataLen;
+  private long dataOnlyLen;
   private int nextColIdx;
 
   private Deque<SliceInfo> sliceInfos = new ArrayDeque<>();
 
   KudoTableHeaderCalc(int rowOffset, int numRows, int numFlattenedCols) {
     this.root = new SliceInfo(rowOffset, numRows);
-    this.totalDataLen = 0;
+    this.dataOnlyLen = 0;
     sliceInfos.addLast(this.root);
     this.bitset = new byte[(numFlattenedCols + 7) / 8];
     this.numFlattenedCols = numFlattenedCols;
     this.nextColIdx = 0;
   }
 
   public KudoTableHeader getHeader() {
+    int headerSize = KudoTableHeader.getSerializedSize(bitset.length);
+    // The validity is a bit odd because we want to pad it for 4 byte alignment
+    // But that is relative to the header, not the payload buffer
+    long paddedValiditySize = padForValidityAlignment(validityBufferLen, headerSize);
     return new KudoTableHeader(toIntExact(root.offset),
         toIntExact(root.rowCount),
-        toIntExact(validityBufferLen),
-        toIntExact(offsetBufferLen),
-        toIntExact(totalDataLen),
+        toIntExact(paddedValiditySize),
+        toIntExact(padForHostAlignment(offsetBufferLen)),
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
user: login: liurenjie1024
body: Please also modify [this part](https://github.com/NVIDIA/spark-rapids-jni/blob/b48b0c61cd14231738e5f0aa5fe750b1d1e4c400/src/test/java/com/nvidia/spark/rapids/jni/kudo/KudoConcatValidityTest.java#L35) to adjust tests.
created_at: 2025-05-06T07:55:54Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,106 @@
+/*
+ * Copyright (c) 2025, 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"
+#include "shuffle_split.hpp"
+
+extern "C" {
+
+JNIEXPORT jlongArray JNICALL
+Java_com_nvidia_spark_rapids_jni_kudo_KudoGpuSerializer_splitAndSerializeToDevice(
+  JNIEnv* env, jclass, jlong j_table_view, jintArray j_splits)
+{
+  JNI_NULL_CHECK(env, j_table_view, "table is null", NULL);
+  JNI_NULL_CHECK(env, j_splits, "splits is null", NULL);
+  try {
+    cudf::jni::auto_set_device(env);
+
+    auto table = reinterpret_cast<cudf::table_view const*>(j_table_view);
+    const cudf::jni::native_jintArray n_splits(env, j_splits);
+    std::vector<cudf::size_type> splits = n_splits.to_vector<int>();
+
+    auto [split_result, split_meta] = spark_rapids_jni::shuffle_split(
+      *table, splits, cudf::get_default_stream(), cudf::get_current_device_resource());
+    // We need to copy the offsets back to a host memory buffer
+
+    cudf::jni::native_jlongArray result(env, 6);
+    result[0]    = reinterpret_cast<jlong>(split_result.partitions->data());
+    result[1]    = static_cast<jlong>(split_result.partitions->size());
+    result[2]    = reinterpret_cast<jlong>(split_result.partitions.get());
+    auto offsets = std::make_unique<rmm::device_buffer>(std::move(split_result.offsets.release()));
+    result[3]    = reinterpret_cast<jlong>(offsets->data());
+    result[4]    = static_cast<jlong>(offsets->size());
+    result[5]    = reinterpret_cast<jlong>(offsets.get());
+
+    split_result.partitions.release();
+    offsets.release();
path: src/main/cpp/src/KudoGpuSerializerJni.cpp
user: login: hyperbolic2346
body: Am I reading this correctly that you're returning a pointer into offsets, but then releasing offsets here? Shouldn't the release wait until after the data has been copied out of that pointer?
created_at: 2025-05-06T15:53:00Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,106 @@
+/*
+ * Copyright (c) 2025, 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"
+#include "shuffle_split.hpp"
+
+extern "C" {
+
+JNIEXPORT jlongArray JNICALL
+Java_com_nvidia_spark_rapids_jni_kudo_KudoGpuSerializer_splitAndSerializeToDevice(
+  JNIEnv* env, jclass, jlong j_table_view, jintArray j_splits)
+{
+  JNI_NULL_CHECK(env, j_table_view, "table is null", NULL);
+  JNI_NULL_CHECK(env, j_splits, "splits is null", NULL);
+  try {
+    cudf::jni::auto_set_device(env);
+
+    auto table = reinterpret_cast<cudf::table_view const*>(j_table_view);
+    const cudf::jni::native_jintArray n_splits(env, j_splits);
+    std::vector<cudf::size_type> splits = n_splits.to_vector<int>();
+
+    auto [split_result, split_meta] = spark_rapids_jni::shuffle_split(
+      *table, splits, cudf::get_default_stream(), cudf::get_current_device_resource());
+    // We need to copy the offsets back to a host memory buffer
+
+    cudf::jni::native_jlongArray result(env, 6);
+    result[0]    = reinterpret_cast<jlong>(split_result.partitions->data());
+    result[1]    = static_cast<jlong>(split_result.partitions->size());
+    result[2]    = reinterpret_cast<jlong>(split_result.partitions.get());
+    auto offsets = std::make_unique<rmm::device_buffer>(std::move(split_result.offsets.release()));
+    result[3]    = reinterpret_cast<jlong>(offsets->data());
+    result[4]    = static_cast<jlong>(offsets->size());
+    result[5]    = reinterpret_cast<jlong>(offsets.get());
+
+    split_result.partitions.release();
+    offsets.release();
path: src/main/cpp/src/KudoGpuSerializerJni.cpp
user: login: revans2
body: What was I thinking???
created_at: 2025-05-06T18:17:23Z
author_association: COLLABORATOR
diff_hunk: @@ -58,7 +58,7 @@ static SlicedValidityBufferInfo calc(int rowOffset, int numRows) {
     }
     int bufferOffset = rowOffset / 8;
     int beginBit = rowOffset % 8;
-    int bufferLength = 0;
+    int bufferLength = 1;
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/SlicedValidityBufferInfo.java
user: login: revans2
body: I changed the code so that if numRows is 0, then hasValidity returns false so no buffers are added instead of having hasValidity say that they are there, but nothing is included.
created_at: 2025-05-06T20:10:17Z
author_association: COLLABORATOR
diff_hunk: @@ -58,7 +58,7 @@ static SlicedValidityBufferInfo calc(int rowOffset, int numRows) {
     }
     int bufferOffset = rowOffset / 8;
     int beginBit = rowOffset % 8;
-    int bufferLength = 0;
+    int bufferLength = 1;
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/SlicedValidityBufferInfo.java
user: login: revans2
body: I should add that I did this to work around a compatibility issue with the GPU Kudo implementation that was sometimes outputting a byte even if there were 0 rows. That has since been fixed so I can probably revert this, if you want me to.
created_at: 2025-05-06T20:29:34Z
author_association: COLLABORATOR
diff_hunk: @@ -58,7 +58,7 @@ static SlicedValidityBufferInfo calc(int rowOffset, int numRows) {
     }
     int bufferOffset = rowOffset / 8;
     int beginBit = rowOffset % 8;
-    int bufferLength = 0;
+    int bufferLength = 1;
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/SlicedValidityBufferInfo.java
user: login: liurenjie1024
body: I would prefer to revert this as it's quite confusing.
created_at: 2025-05-07T01:38:11Z
author_association: COLLABORATOR
diff_hunk: @@ -41,65 +42,75 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   private final byte[] bitset;
   private long validityBufferLen;
   private long offsetBufferLen;
-  private long totalDataLen;
+  private long dataOnlyLen;
   private int nextColIdx;
 
   private Deque<SliceInfo> sliceInfos = new ArrayDeque<>();
 
   KudoTableHeaderCalc(int rowOffset, int numRows, int numFlattenedCols) {
     this.root = new SliceInfo(rowOffset, numRows);
-    this.totalDataLen = 0;
+    this.dataOnlyLen = 0;
     sliceInfos.addLast(this.root);
     this.bitset = new byte[(numFlattenedCols + 7) / 8];
     this.numFlattenedCols = numFlattenedCols;
     this.nextColIdx = 0;
   }
 
   public KudoTableHeader getHeader() {
+    int headerSize = KudoTableHeader.getSerializedSize(bitset.length);
+    // The validity is a bit odd because we want to pad it for 4 byte alignment
+    // But that is relative to the header, not the payload buffer
+    long paddedValiditySize = padForValidityAlignment(validityBufferLen, headerSize);
     return new KudoTableHeader(toIntExact(root.offset),
         toIntExact(root.rowCount),
-        toIntExact(validityBufferLen),
-        toIntExact(offsetBufferLen),
-        toIntExact(totalDataLen),
+        toIntExact(paddedValiditySize),
+        toIntExact(padForHostAlignment(offsetBufferLen)),
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
user: login: revans2
body: I don't think we have to change that code. Unless we are shuffling something with no data an no offsets, but it does have validity. i.e. a struct of nothing. I will see what I can do to test it. In all other cases it still works. It is a little ugly, but it still works.  Say we have 2 1-byte validity buffers.

The new code will pad that out to

|a|b|?|?|

instead of 

|a|?|?|?|b|?|?|?|

The code for the old way would read in |a|?|?|?| and ignore the three bytes at the end. Then it would read |b|?|?|?| and ignore the three bytes at the end.

The code for the new layout is going to read for a - |a|b|?|?| and ignore the three bytes at the end. Then for b it would read |b|?|?|  plus one extra byte from the section after it, and ignore the three bytes at the end. So the only time we would be in trouble is if there is no extra byte to read after the validity.
created_at: 2025-05-08T19:47:25Z
author_association: COLLABORATOR
diff_hunk: @@ -41,65 +42,75 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   private final byte[] bitset;
   private long validityBufferLen;
   private long offsetBufferLen;
-  private long totalDataLen;
+  private long dataOnlyLen;
   private int nextColIdx;
 
   private Deque<SliceInfo> sliceInfos = new ArrayDeque<>();
 
   KudoTableHeaderCalc(int rowOffset, int numRows, int numFlattenedCols) {
     this.root = new SliceInfo(rowOffset, numRows);
-    this.totalDataLen = 0;
+    this.dataOnlyLen = 0;
     sliceInfos.addLast(this.root);
     this.bitset = new byte[(numFlattenedCols + 7) / 8];
     this.numFlattenedCols = numFlattenedCols;
     this.nextColIdx = 0;
   }
 
   public KudoTableHeader getHeader() {
+    int headerSize = KudoTableHeader.getSerializedSize(bitset.length);
+    // The validity is a bit odd because we want to pad it for 4 byte alignment
+    // But that is relative to the header, not the payload buffer
+    long paddedValiditySize = padForValidityAlignment(validityBufferLen, headerSize);
     return new KudoTableHeader(toIntExact(root.offset),
         toIntExact(root.rowCount),
-        toIntExact(validityBufferLen),
-        toIntExact(offsetBufferLen),
-        toIntExact(totalDataLen),
+        toIntExact(paddedValiditySize),
+        toIntExact(padForHostAlignment(offsetBufferLen)),
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
user: login: revans2
body: As for KudoConcatValidityTest.scala, it is building the validity buffer on the fly each time before it concatenates it. Do you want me to remove the extra padding so that we never walk off past the end of the known good data?
created_at: 2025-05-08T19:53:04Z
author_association: COLLABORATOR
diff_hunk: @@ -41,65 +42,75 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   private final byte[] bitset;
   private long validityBufferLen;
   private long offsetBufferLen;
-  private long totalDataLen;
+  private long dataOnlyLen;
   private int nextColIdx;
 
   private Deque<SliceInfo> sliceInfos = new ArrayDeque<>();
 
   KudoTableHeaderCalc(int rowOffset, int numRows, int numFlattenedCols) {
     this.root = new SliceInfo(rowOffset, numRows);
-    this.totalDataLen = 0;
+    this.dataOnlyLen = 0;
     sliceInfos.addLast(this.root);
     this.bitset = new byte[(numFlattenedCols + 7) / 8];
     this.numFlattenedCols = numFlattenedCols;
     this.nextColIdx = 0;
   }
 
   public KudoTableHeader getHeader() {
+    int headerSize = KudoTableHeader.getSerializedSize(bitset.length);
+    // The validity is a bit odd because we want to pad it for 4 byte alignment
+    // But that is relative to the header, not the payload buffer
+    long paddedValiditySize = padForValidityAlignment(validityBufferLen, headerSize);
     return new KudoTableHeader(toIntExact(root.offset),
         toIntExact(root.rowCount),
-        toIntExact(validityBufferLen),
-        toIntExact(offsetBufferLen),
-        toIntExact(totalDataLen),
+        toIntExact(paddedValiditySize),
+        toIntExact(padForHostAlignment(offsetBufferLen)),
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
user: login: liurenjie1024
body: > The code for the new layout is going to read for a - |a|b|?|?| and ignore the three bytes at the end. Then for b it would read |b|?|?| plus one extra byte from the section after it, and ignore the three bytes at the end. So the only time we would be in trouble is if there is no extra byte to read after the validity.

Yes, this case is what I'm worrying about. I think we still need to change that code a little. Layout in this pr is like 
|header|validity 1|validity 2|padding|

We can adjust a little so that for example when we read validity 2, we could adjust the starting bit, and read more bytes before validity 2 so that it's at least 4 bytes.
created_at: 2025-05-09T08:34:06Z
author_association: COLLABORATOR
diff_hunk: @@ -41,65 +42,75 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   private final byte[] bitset;
   private long validityBufferLen;
   private long offsetBufferLen;
-  private long totalDataLen;
+  private long dataOnlyLen;
   private int nextColIdx;
 
   private Deque<SliceInfo> sliceInfos = new ArrayDeque<>();
 
   KudoTableHeaderCalc(int rowOffset, int numRows, int numFlattenedCols) {
     this.root = new SliceInfo(rowOffset, numRows);
-    this.totalDataLen = 0;
+    this.dataOnlyLen = 0;
     sliceInfos.addLast(this.root);
     this.bitset = new byte[(numFlattenedCols + 7) / 8];
     this.numFlattenedCols = numFlattenedCols;
     this.nextColIdx = 0;
   }
 
   public KudoTableHeader getHeader() {
+    int headerSize = KudoTableHeader.getSerializedSize(bitset.length);
+    // The validity is a bit odd because we want to pad it for 4 byte alignment
+    // But that is relative to the header, not the payload buffer
+    long paddedValiditySize = padForValidityAlignment(validityBufferLen, headerSize);
     return new KudoTableHeader(toIntExact(root.offset),
         toIntExact(root.rowCount),
-        toIntExact(validityBufferLen),
-        toIntExact(offsetBufferLen),
-        toIntExact(totalDataLen),
+        toIntExact(paddedValiditySize),
+        toIntExact(padForHostAlignment(offsetBufferLen)),
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
user: login: liurenjie1024
body: > As for KudoConcatValidityTest.scala, it is building the validity buffer on the fly each time before it concatenates it. Do you want me to remove the extra padding so that we never walk off past the end of the known good data?

Yes, I think it would be better if the test codes aligns with actual serialization logic.
created_at: 2025-05-09T08:35:18Z
author_association: COLLABORATOR
diff_hunk: @@ -41,65 +42,75 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   private final byte[] bitset;
   private long validityBufferLen;
   private long offsetBufferLen;
-  private long totalDataLen;
+  private long dataOnlyLen;
   private int nextColIdx;
 
   private Deque<SliceInfo> sliceInfos = new ArrayDeque<>();
 
   KudoTableHeaderCalc(int rowOffset, int numRows, int numFlattenedCols) {
     this.root = new SliceInfo(rowOffset, numRows);
-    this.totalDataLen = 0;
+    this.dataOnlyLen = 0;
     sliceInfos.addLast(this.root);
     this.bitset = new byte[(numFlattenedCols + 7) / 8];
     this.numFlattenedCols = numFlattenedCols;
     this.nextColIdx = 0;
   }
 
   public KudoTableHeader getHeader() {
+    int headerSize = KudoTableHeader.getSerializedSize(bitset.length);
+    // The validity is a bit odd because we want to pad it for 4 byte alignment
+    // But that is relative to the header, not the payload buffer
+    long paddedValiditySize = padForValidityAlignment(validityBufferLen, headerSize);
     return new KudoTableHeader(toIntExact(root.offset),
         toIntExact(root.rowCount),
-        toIntExact(validityBufferLen),
-        toIntExact(offsetBufferLen),
-        toIntExact(totalDataLen),
+        toIntExact(paddedValiditySize),
+        toIntExact(padForHostAlignment(offsetBufferLen)),
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
user: login: revans2
body: @liurenjie1024 I fixed the issue. It would only show up if we are shuffling columns of only empty structs, but the code now works. I don't think there is much if any performance impact to it, but please let me know if you do run into any and I can try and optimize it more.
created_at: 2025-05-09T16:38:48Z
author_association: COLLABORATOR
diff_hunk: @@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/SlicedValidityBufferInfo.java
user: login: liurenjie1024
body: nit: Revert this change?
created_at: 2025-05-12T09:45:15Z
author_association: COLLABORATOR
diff_hunk: @@ -601,6 +605,40 @@ private int copySourceCaseTwo(HostMemoryBuffer src, int srcOffset,
       return nullCount;
     }
 
+    private int getIntAsBytes(HostMemoryBuffer src, long offset, long length) {
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableMerger.java
user: login: liurenjie1024
body: ```suggestion
    private static int getIntAsBytes(HostMemoryBuffer src, long offset, long length) {
```
created_at: 2025-05-12T09:50:56Z
author_association: COLLABORATOR
diff_hunk: @@ -601,6 +605,40 @@ private int copySourceCaseTwo(HostMemoryBuffer src, int srcOffset,
       return nullCount;
     }
 
+    private int getIntAsBytes(HostMemoryBuffer src, long offset, long length) {
+      // We need to build up the int ourselves and pad it as needed
+      int ret = 0;
+      for (int at = 0; at < 4 && (offset + at) < length ; at++) {
+        int b = src.getByte(offset + at) & 0xFF;
+        ret |= b << (at * 8);
+      }
+      return ret;
+    }
+
+    private int getIntSafe(HostMemoryBuffer src, long offset) {
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableMerger.java
user: login: liurenjie1024
body: ```suggestion
    private static int getIntSafe(HostMemoryBuffer src, long offset) {
```
created_at: 2025-05-12T09:51:06Z
author_association: COLLABORATOR
diff_hunk: @@ -601,6 +605,40 @@ private int copySourceCaseTwo(HostMemoryBuffer src, int srcOffset,
       return nullCount;
     }
 
+    private int getIntAsBytes(HostMemoryBuffer src, long offset, long length) {
+      // We need to build up the int ourselves and pad it as needed
+      int ret = 0;
+      for (int at = 0; at < 4 && (offset + at) < length ; at++) {
+        int b = src.getByte(offset + at) & 0xFF;
+        ret |= b << (at * 8);
+      }
+      return ret;
+    }
+
+    private int getIntSafe(HostMemoryBuffer src, long offset) {
+      long length = src.getLength();
+      if (offset + 4 < length) {
+        return src.getInt(offset);
+      } else {
+        return getIntAsBytes(src, offset, length);
+      }
+    }
+
+    private void getIntsSafe(HostMemoryBuffer src, int[] dst, long dstIndex, long srcOffset, int count) {
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableMerger.java
user: login: liurenjie1024
body: ```suggestion
    private static void getIntsSafe(HostMemoryBuffer src, int[] dst, long dstIndex, long srcOffset, int count) {
```
created_at: 2025-05-12T09:51:15Z
author_association: COLLABORATOR
diff_hunk: @@ -41,65 +42,75 @@ class KudoTableHeaderCalc implements HostColumnsVisitor {
   private final byte[] bitset;
   private long validityBufferLen;
   private long offsetBufferLen;
-  private long totalDataLen;
+  private long dataOnlyLen;
   private int nextColIdx;
 
   private Deque<SliceInfo> sliceInfos = new ArrayDeque<>();
 
   KudoTableHeaderCalc(int rowOffset, int numRows, int numFlattenedCols) {
     this.root = new SliceInfo(rowOffset, numRows);
-    this.totalDataLen = 0;
+    this.dataOnlyLen = 0;
     sliceInfos.addLast(this.root);
     this.bitset = new byte[(numFlattenedCols + 7) / 8];
     this.numFlattenedCols = numFlattenedCols;
     this.nextColIdx = 0;
   }
 
   public KudoTableHeader getHeader() {
+    int headerSize = KudoTableHeader.getSerializedSize(bitset.length);
+    // The validity is a bit odd because we want to pad it for 4 byte alignment
+    // But that is relative to the header, not the payload buffer
+    long paddedValiditySize = padForValidityAlignment(validityBufferLen, headerSize);
     return new KudoTableHeader(toIntExact(root.offset),
         toIntExact(root.rowCount),
-        toIntExact(validityBufferLen),
-        toIntExact(offsetBufferLen),
-        toIntExact(totalDataLen),
+        toIntExact(paddedValiditySize),
+        toIntExact(padForHostAlignment(offsetBufferLen)),
path: src/main/java/com/nvidia/spark/rapids/jni/kudo/KudoTableHeaderCalc.java
user: login: liurenjie1024
body: I think the changes looks good.
created_at: 2025-05-12T10:05:57Z
author_association: COLLABORATOR

PR comments:
user: login: revans2
created_at: 2025-05-05T20:14:10Z
body: build
author_association: COLLABORATOR
user: login: revans2
created_at: 2025-05-05T20:16:04Z
body: build
author_association: COLLABORATOR
user: login: revans2
created_at: 2025-05-09T16:37:19Z
body: build
author_association: COLLABORATOR
user: login: revans2
created_at: 2025-05-12T14:51:33Z
body: @liurenjie1024 and @hyperbolic2346 I think I have addressed all of your comments. Please take another look
author_association: COLLABORATOR
user: login: revans2
created_at: 2025-05-12T14:52:10Z
body: Oops I see some more from @liurenjie1024 will get on those right away
author_association: COLLABORATOR
user: login: revans2
created_at: 2025-05-12T14:56:44Z
body: build
author_association: COLLABORATOR
user: login: revans2
created_at: 2025-05-12T14:57:07Z
body: @liurenjie1024 and @hyperbolic2346 Okay I think I have NOW addressed all of your comments. Please take another look.
author_association: COLLABORATOR