From 4341dd00e3549cfa3b7f48bf379ed2fbf938a187 Mon Sep 17 00:00:00 2001
From: Nghia Truong <nghiat@nvidia.com>
Date: Mon, 14 Oct 2024 12:10:43 -0700
Subject: [PATCH] Use `make_strings_column_batch`

Signed-off-by: Nghia Truong <nghiat@nvidia.com>
---
 src/main/cpp/src/get_json_object.cu | 35 ++++++++++++++---------------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/src/main/cpp/src/get_json_object.cu b/src/main/cpp/src/get_json_object.cu
index 2ada677ff..8fce46bca 100644
--- a/src/main/cpp/src/get_json_object.cu
+++ b/src/main/cpp/src/get_json_object.cu
@@ -1019,7 +1019,6 @@ std::vector<std::unique_ptr<cudf::column>> get_json_object_batch(
     construct_path_commands(json_paths, stream);
 
   auto const num_outputs = json_paths.size();
-  std::vector<std::unique_ptr<cudf::column>> output;
 
   // The error check array contains markers denoting if there is any out-of-bound write occurs
   // (first `num_outputs` elements), or if the nesting depth exceeded its limits (the last element).
@@ -1052,39 +1051,31 @@ std::vector<std::unique_ptr<cudf::column>> get_json_object_batch(
   auto d_path_data = cudf::detail::make_device_uvector_async(
     h_path_data, stream, rmm::mr::get_current_device_resource());
   thrust::uninitialized_fill(
-    rmm::exec_policy(stream), d_error_check.begin(), d_error_check.end(), 0);
+    rmm::exec_policy_nosync(stream), d_error_check.begin(), d_error_check.end(), 0);
 
   kernel_launcher::exec(input, d_path_data, d_max_path_depth_exceeded, stream);
   auto h_error_check = cudf::detail::make_host_vector_sync(d_error_check, stream);
   auto has_no_oob    = check_error(h_error_check);
 
+  std::vector<cudf::device_span<thrust::pair<char const*, cudf::size_type> const>>
+    batch_stringviews;
+  batch_stringviews.reserve(out_stringviews.size());
+
   // If we didn't see any out-of-bound write, everything is good so far.
   // Just gather the output strings and return.
   if (has_no_oob) {
-    nvtx3::scoped_range_in range("create output");
-#if 0
-    for (auto const& out_sview : out_stringviews) {
-      output.emplace_back(cudf::make_strings_column(out_sview, stream, mr));
-    }
-    return output;
-#else
-    std::vector<cudf::device_span<thrust::pair<char const*, cudf::size_type> const>>
-      batch_stringviews;
-    batch_stringviews.reserve(out_stringviews.size());
     for (auto const& out_sview : out_stringviews) {
       batch_stringviews.emplace_back(out_sview);
     }
     return cudf::make_strings_column_batch(batch_stringviews, stream, mr);
-#endif
   }
   // From here, we had out-of-bound write. Although this is very rare, it may still happen.
 
-  throw std::runtime_error("Not implemented");
-
   std::vector<std::pair<rmm::device_buffer, cudf::size_type>> out_null_masks_and_null_counts;
   std::vector<std::pair<std::unique_ptr<cudf::column>, int64_t>> out_offsets_and_sizes;
   std::vector<rmm::device_uvector<char>> out_char_buffers;
   std::vector<std::size_t> oob_indices;
+  std::vector<std::size_t> no_oob_indices;
 
   // Check validity from the stored char pointers.
   auto const validator = [] __device__(thrust::pair<char const*, cudf::size_type> const item) {
@@ -1098,7 +1089,6 @@ std::vector<std::unique_ptr<cudf::column>> get_json_object_batch(
 
     if (h_error_check[idx]) {
       oob_indices.emplace_back(idx);
-      output.emplace_back(nullptr);  // just placeholder.
 
       out_null_masks_and_null_counts.emplace_back(
         cudf::detail::valid_if(out_sview.begin(), out_sview.end(), validator, stream, mr));
@@ -1124,9 +1114,18 @@ std::vector<std::unique_ptr<cudf::column>> get_json_object_batch(
                                   out_char_buffers.back().data(),
                                   d_error_check.data() + idx});
     } else {
-      output.emplace_back(cudf::make_strings_column(out_sview, stream, mr));
+      no_oob_indices.emplace_back(idx);
+      batch_stringviews.emplace_back(out_sview);
     }
   }
+
+  std::vector<std::unique_ptr<cudf::column>> output(num_outputs);
+  auto no_oob_output = cudf::make_strings_column_batch(batch_stringviews, stream, mr);
+  for (std::size_t idx = 0; idx < no_oob_indices.size(); ++idx) {
+    auto const out_idx = no_oob_indices[idx];
+    output[out_idx]    = std::move(no_oob_output[idx]);
+  }
+
   // These buffers are no longer needed.
   scratch_buffers.clear();
   out_stringviews.clear();
@@ -1135,7 +1134,7 @@ std::vector<std::unique_ptr<cudf::column>> get_json_object_batch(
   d_path_data = cudf::detail::make_device_uvector_async(
     h_path_data, stream, rmm::mr::get_current_device_resource());
   thrust::uninitialized_fill(
-    rmm::exec_policy(stream), d_error_check.begin(), d_error_check.end(), 0);
+    rmm::exec_policy_nosync(stream), d_error_check.begin(), d_error_check.end(), 0);
   kernel_launcher::exec(input, d_path_data, d_max_path_depth_exceeded, stream);
   h_error_check = cudf::detail::make_host_vector_sync(d_error_check, stream);
   has_no_oob    = check_error(h_error_check);