From ff2db7e0b7f623e381fe937f019fefd68501251e Mon Sep 17 00:00:00 2001
From: Chong Gao <res_life@163.com>
Date: Tue, 9 Apr 2024 19:13:50 +0800
Subject: [PATCH] Fix nvbug: get-json-obj get incorrect result on some
 GPUs(H100, V100)

---
 src/main/cpp/src/get_json_object.cu  | 13 ++++++++-----
 src/main/cpp/src/get_json_object.hpp |  1 +
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/main/cpp/src/get_json_object.cu b/src/main/cpp/src/get_json_object.cu
index 51a9cfb67e..3844f42310 100644
--- a/src/main/cpp/src/get_json_object.cu
+++ b/src/main/cpp/src/get_json_object.cu
@@ -1157,7 +1157,7 @@ rmm::device_uvector<path_instruction> construct_path_commands(
  * @param out_buf_size Size of the output buffer
  * @returns A pair containing the result code and the output buffer.
  */
-__device__ thrust::pair<bool, size_t> get_json_object_single(
+__device__ thrust::pair<bool, json_generator> get_json_object_single(
   char const* input,
   cudf::size_type input_len,
   path_instruction const* path_commands_ptr,
@@ -1180,7 +1180,7 @@ __device__ thrust::pair<bool, size_t> get_json_object_single(
 
   j_parser.next_token();
   // JSON validation check
-  if (json_token::ERROR == j_parser.get_current_token()) { return {false, 0}; }
+  if (json_token::ERROR == j_parser.get_current_token()) { return {false, generator}; }
 
   if (nullptr == out_buf) {
     // First step: preprocess sizes
@@ -1193,12 +1193,12 @@ __device__ thrust::pair<bool, size_t> get_json_object_single(
       // set output as zero to tell second step
       generator.set_output_len_zero();
     }
-    return {success, generator.get_output_len()};
+    return {success, generator};
   } else {
     // Second step: writes output
     bool success = evaluate_path(
       j_parser, generator, write_style::raw_style, path_commands_ptr, path_commands_size);
-    return {success, generator.get_output_len()};
+    return {success, generator};
   }
 }
 
@@ -1245,8 +1245,11 @@ __launch_bounds__(block_size) CUDF_KERNEL
         out_buf != nullptr ? output_offsets[tid + 1] - output_offsets[tid] : 0;
 
       // process one single row
-      auto [result, output_size] = get_json_object_single(
+      bool result;
+      json_generator generator;
+      thrust::tie(result, generator) = get_json_object_single(
         str.data(), str.size_bytes(), path_commands_ptr, path_commands_size, dst, dst_size);
+      output_size = generator.get_output_len();
       if (result) { is_valid = true; }
     }
 
diff --git a/src/main/cpp/src/get_json_object.hpp b/src/main/cpp/src/get_json_object.hpp
index 628d8aa749..2d7dffb78d 100644
--- a/src/main/cpp/src/get_json_object.hpp
+++ b/src/main/cpp/src/get_json_object.hpp
@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#pragma once
 
 #include "json_parser.cuh"