Add Spark-safe UTF-8 string reverse for truncated trailing bytes

firestarman · firestarman · commit 43ef77bcacaa · 2026-08-03T15:48:46.000+08:00
Libcudf reverse assumes well-formed UTF-8 and can over-read into the next
row when Spark StringType holds truncated trailing multi-byte sequences.
Add a JNI reverse that matches UTF8String.numBytesForFirstByte with
SPARK-57507 remaining-byte clamping.

Signed-off-by: Liangcai Li &lt;firestarmanllc@gmail.com&gt;
diff --git a/src/main/cpp/CMakeLists.txt b/src/main/cpp/CMakeLists.txt
@@ -288,6 +288,7 @@ add_library(
   src/protobuf/protobuf_builders.cu
   src/protobuf/protobuf_kernels.cu
   src/regex_rewrite_utils.cu
+  src/reverse_strings.cu
   src/row_conversion.cu
   src/round_float.cu
   src/shuffle_assemble.cu
diff --git a/src/main/cpp/src/StringUtilsJni.cpp b/src/main/cpp/src/StringUtilsJni.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2025, NVIDIA CORPORATION.
+ * Copyright (c) 2025-2026, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,6 +15,7 @@
  */
 
 #include "cudf_jni_apis.hpp"
+#include "reverse_strings.hpp"
 #include "uuid.hpp"
 
 extern "C" {
@@ -31,4 +32,18 @@ JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_StringUtils_randomUUIDs
   }
   JNI_CATCH(env, 0);
 }
+
+JNIEXPORT jlong JNICALL
+Java_com_nvidia_spark_rapids_jni_StringUtils_reverseStrings(JNIEnv* env, jclass, jlong input_handle)
+{
+  JNI_NULL_CHECK(env, input_handle, "input column is null", 0);
+  JNI_TRY
+  {
+    cudf::jni::auto_set_device(env);
+    auto const input = reinterpret_cast<cudf::column_view const*>(input_handle);
+    return cudf::jni::release_as_jlong(
+      spark_rapids_jni::reverse_strings(cudf::strings_column_view{*input}));
+  }
+  JNI_CATCH(env, 0);
+}
 }
diff --git a/src/main/cpp/src/reverse_strings.cu b/src/main/cpp/src/reverse_strings.cu
@@ -0,0 +1,123 @@
+/*
+ * Copyright (c) 2026, 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 "nvtx_ranges.hpp"
+#include "reverse_strings.hpp"
+
+#include <cudf/column/column_device_view.cuh>
+#include <cudf/column/column_factories.hpp>
+#include <cudf/detail/offsets_iterator_factory.cuh>
+#include <cudf/strings/string_view.cuh>
+#include <cudf/strings/strings_column_view.hpp>
+#include <cudf/utilities/default_stream.hpp>
+#include <cudf/utilities/memory_resource.hpp>
+
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/exec_policy.hpp>
+
+#include <cuda/iterator>
+#include <cuda/std/algorithm>
+#include <thrust/for_each.h>
+
+namespace spark_rapids_jni {
+namespace detail {
+namespace {
+
+/**
+ * @brief Spark `UTF8String.numBytesForFirstByte` character width.
+ *
+ * Continuation bytes and UTF-8-disallowed lead bytes are treated as width 1, matching
+ * Spark's `bytesOfCodePointInUTF8` table with the `(numBytes == 0) ? 1 : numBytes` rule.
+ */
+__device__ __forceinline__ cudf::size_type spark_num_bytes_for_first_byte(uint8_t byte)
+{
+  if (byte <= 0x7F) { return 1; }
+  // 0x80-0xBF continuation and 0xC0-0xC1 disallowed -> Spark width 1
+  if (byte <= 0xC1) { return 1; }
+  if (byte <= 0xDF) { return 2; }
+  if (byte <= 0xEF) { return 3; }
+  // 0xF0-0xF4 valid 4-byte leads; 0xF5-0xFF disallowed -> Spark width 1
+  if (byte <= 0xF4) { return 4; }
+  return 1;
+}
+
+/**
+ * @brief Reverse characters in each string with Spark clamp semantics (SPARK-57507).
+ */
+struct reverse_characters_fn {
+  cudf::column_device_view const d_strings;
+  cudf::detail::input_offsetalator d_offsets;
+  char* d_chars;
+
+  __device__ void operator()(cudf::size_type idx) const
+  {
+    if (d_strings.is_null(idx)) { return; }
+    auto const d_str  = d_strings.element<cudf::string_view>(idx);
+    auto const nbytes = d_str.size_bytes();
+    if (nbytes == 0) { return; }
+
+    auto const* in = reinterpret_cast<uint8_t const*>(d_str.data());
+    // Write character chunks from the end of this row's output region, matching Spark:
+    // keep bytes within a character in order, reverse the order of characters, and clamp
+    // each character width to the bytes remaining in this row.
+    auto* out_end     = d_chars + d_offsets[idx] + nbytes;
+    cudf::size_type i = 0;
+    while (i < nbytes) {
+      auto const declared = spark_num_bytes_for_first_byte(in[i]);
+      auto const len      = cuda::std::min(declared, nbytes - i);
+      out_end -= len;
+      for (cudf::size_type j = 0; j < len; ++j) {
+        out_end[j] = static_cast<char>(in[i + j]);
+      }
+      i += len;
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<cudf::column> reverse_strings(cudf::strings_column_view const& input,
+                                              rmm::cuda_stream_view stream,
+                                              rmm::device_async_resource_ref mr)
+{
+  if (input.is_empty()) { return cudf::make_empty_column(cudf::type_id::STRING); }
+
+  // Preserve offsets/nulls; rewrite only the character bytes.
+  auto result          = std::make_unique<cudf::column>(input.parent(), stream, mr);
+  auto const sv        = cudf::strings_column_view(result->view());
+  auto const d_offsets = cudf::detail::offsetalator_factory::make_input_iterator(sv.offsets());
+  auto* d_chars        = result->mutable_view().head<char>();
+
+  auto const d_column = cudf::column_device_view::create(input.parent(), stream);
+  thrust::for_each_n(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+                     cuda::counting_iterator<cudf::size_type>{0},
+                     input.size(),
+                     reverse_characters_fn{*d_column, d_offsets, d_chars});
+
+  return result;
+}
+
+}  // namespace detail
+
+std::unique_ptr<cudf::column> reverse_strings(cudf::strings_column_view const& input,
+                                              rmm::cuda_stream_view stream,
+                                              rmm::device_async_resource_ref mr)
+{
+  SRJ_FUNC_RANGE();
+  return detail::reverse_strings(input, stream, mr);
+}
+
+}  // namespace spark_rapids_jni
diff --git a/src/main/cpp/src/reverse_strings.hpp b/src/main/cpp/src/reverse_strings.hpp
@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2026, 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.
+ */
+#pragma once
+
+#include <cudf/column/column.hpp>
+#include <cudf/strings/strings_column_view.hpp>
+#include <cudf/utilities/default_stream.hpp>
+#include <cudf/utilities/memory_resource.hpp>
+
+namespace spark_rapids_jni {
+
+/**
+ * @brief Reverse strings using Spark `UTF8String.reverse` character-width semantics.
+ *
+ * Unlike libcudf `cudf::strings::reverse`, character widths follow Spark's
+ * `UTF8String.numBytesForFirstByte` and are clamped to the bytes remaining in each
+ * row (`min(declared_width, remaining)`). This matches SPARK-57507 and avoids
+ * reading past a truncated trailing multi-byte UTF-8 sequence into the next row.
+ *
+ * Well-formed UTF-8 results match libcudf reverse. Null rows remain null. Output
+ * offsets match the input offsets (byte length is preserved per row).
+ *
+ * @param input Strings column
+ * @param stream CUDA stream used for device memory operations
+ * @param mr Device memory resource used to allocate the returned column
+ * @return New strings column with Spark-compatible reversed contents
+ */
+std::unique_ptr<cudf::column> reverse_strings(
+  cudf::strings_column_view const& input,
+  rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
+
+}  // namespace spark_rapids_jni
diff --git a/src/main/cpp/tests/CMakeLists.txt b/src/main/cpp/tests/CMakeLists.txt
@@ -122,6 +122,9 @@ ConfigureTest(MAP_UTILS
 ConfigureTest(SUBSTRING_INDEX
     substring_index.cpp)
 
+ConfigureTest(REVERSE_STRINGS
+    reverse_strings.cpp)
+
 ConfigureTest(SHUFFLE_SPLIT
     shuffle_split.cu)
     
diff --git a/src/main/cpp/tests/reverse_strings.cpp b/src/main/cpp/tests/reverse_strings.cpp
@@ -0,0 +1,95 @@
+/*
+ * Copyright (c) 2026, 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 "reverse_strings.hpp"
+
+#include <cudf_test/base_fixture.hpp>
+#include <cudf_test/column_utilities.hpp>
+#include <cudf_test/column_wrapper.hpp>
+
+#include <cudf/strings/strings_column_view.hpp>
+
+#include <string>
+
+namespace {
+
+std::string bytes_to_string(std::initializer_list<uint8_t> bytes)
+{
+  return std::string(bytes.begin(), bytes.end());
+}
+
+}  // namespace
+
+struct ReverseStringsTests : public cudf::test::BaseFixture {};
+
+TEST_F(ReverseStringsTests, EmptyAndNulls)
+{
+  auto const input  = cudf::test::strings_column_wrapper{};
+  auto const result = spark_rapids_jni::reverse_strings(cudf::strings_column_view{input});
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, input);
+
+  auto const with_nulls =
+    cudf::test::strings_column_wrapper({"abc", "", "世界"}, {true, false, true});
+  auto const expected =
+    cudf::test::strings_column_wrapper({"cba", "", "界世"}, {true, false, true});
+  auto const reversed = spark_rapids_jni::reverse_strings(cudf::strings_column_view{with_nulls});
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*reversed, expected);
+}
+
+TEST_F(ReverseStringsTests, WellFormedUtf8)
+{
+  // ASCII, complete 2-/3-/4-byte characters.
+  auto const input = cudf::test::strings_column_wrapper(
+    {"ABC",
+     bytes_to_string({0x41, 0xC3, 0xA9}),                // Aé
+     bytes_to_string({0x41, 0xE4, 0xB8, 0x96}),          // A世
+     bytes_to_string({0x41, 0xF0, 0x90, 0x8D, 0x88})});  // A + U+10048-ish 4-byte
+  auto const expected =
+    cudf::test::strings_column_wrapper({"CBA",
+                                        bytes_to_string({0xC3, 0xA9, 0x41}),
+                                        bytes_to_string({0xE4, 0xB8, 0x96, 0x41}),
+                                        bytes_to_string({0xF0, 0x90, 0x8D, 0x88, 0x41})});
+  auto const result = spark_rapids_jni::reverse_strings(cudf::strings_column_view{input});
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, expected);
+}
+
+TEST_F(ReverseStringsTests, TruncatedTrailingUtf8NoOverread)
+{
+  // SPARK-57507 cases. Neighbor rows carry sentinel bytes; an over-read would pull them in.
+  auto const row0 = bytes_to_string({0x41, 0xCE});              // A + truncated 2-byte lead
+  auto const row1 = bytes_to_string({0xA9, 0x42});              // must not leak into row0
+  auto const row2 = bytes_to_string({0x41, 0xE4, 0xB8});        // A + truncated 3-byte lead
+  auto const row3 = bytes_to_string({0x96, 0x43});              // must not leak into row2
+  auto const row4 = bytes_to_string({0x41, 0xF0, 0x90});        // A + truncated 4-byte lead
+  auto const row5 = bytes_to_string({0x8D, 0x88, 0x44});        // must not leak into row4
+  auto const row6 = bytes_to_string({0xE4, 0xB8, 0x96, 0xCE});  // 世 + orphan 2-byte lead
+  auto const row7 = bytes_to_string({0x45});
+
+  auto const input =
+    cudf::test::strings_column_wrapper({row0, row1, row2, row3, row4, row5, row6, row7});
+  auto const expected = cudf::test::strings_column_wrapper(
+    {bytes_to_string({0xCE, 0x41}),
+     bytes_to_string({0x42, 0xA9}),  // 0xA9 is continuation -> Spark width 1, then 'B'
+     bytes_to_string({0xE4, 0xB8, 0x41}),
+     bytes_to_string({0x43, 0x96}),
+     bytes_to_string({0xF0, 0x90, 0x41}),
+     bytes_to_string({0x44, 0x88, 0x8D}),
+     bytes_to_string({0xCE, 0xE4, 0xB8, 0x96}),
+     bytes_to_string({0x45})});
+
+  auto const result = spark_rapids_jni::reverse_strings(cudf::strings_column_view{input});
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, expected);
+}
diff --git a/src/main/java/com/nvidia/spark/rapids/jni/StringUtils.java b/src/main/java/com/nvidia/spark/rapids/jni/StringUtils.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2025, NVIDIA CORPORATION.
+ * Copyright (c) 2025-2026, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,13 +17,20 @@
 package com.nvidia.spark.rapids.jni;
 
 import ai.rapids.cudf.ColumnVector;
+import ai.rapids.cudf.ColumnView;
 import ai.rapids.cudf.Cuda;
+import ai.rapids.cudf.CudfException;
+import ai.rapids.cudf.NativeDepsLoader;
 import java.lang.management.ManagementFactory;
 import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicLong;
 
 public class StringUtils {
 
+  static {
+    NativeDepsLoader.loadNativeDeps();
+  }
+
   // Stores the sequence ID of calling generate UUIDs.
   private static AtomicLong sequence = new AtomicLong(0);
 
@@ -88,5 +95,20 @@ public static ColumnVector randomUUIDsWithSeed(int rowCount, long seed) {
     return new ColumnVector(randomUUIDs(rowCount, seed));
   }
 
+  /**
+   * Reverse each string using Spark {@code UTF8String.reverse} semantics.
+   * Character widths follow Spark's {@code numBytesForFirstByte} and are clamped to the
+   * bytes remaining in each row (SPARK-57507), so truncated trailing UTF-8 sequences do
+   * not read past the row boundary.
+   *
+   * @param input strings column
+   * @return new strings column with reversed contents
+   */
+  public static ColumnVector reverseStrings(ColumnView input) {
+    return new ColumnVector(reverseStrings(input.getNativeView()));
+  }
+
   private static native long randomUUIDs(int rowCount, long seed);
+
+  private static native long reverseStrings(long inputHandle) throws CudfException;
 }
diff --git a/src/test/java/com/nvidia/spark/rapids/jni/StringUtilsTest.java b/src/test/java/com/nvidia/spark/rapids/jni/StringUtilsTest.java
