POC: Implement from_json_to_structs using get_json_object kernel

src/main/cpp/CMakeLists.txt

@@ -212,6 +212,7 @@ add_library(
   src/datetime_rebase.cu
   src/decimal_utils.cu
   src/from_json_to_raw_map.cu
+  src/from_json_to_structs.cu
   src/get_json_object.cu
   src/histogram.cu
   src/murmur_hash.cu

src/main/cpp/src/JSONUtilsJni.cpp

@@ -18,12 +18,58 @@
 #include "from_json.hpp"
 #include "get_json_object.hpp"
 
+#include <cudf/io/json.hpp>
 #include <cudf/strings/strings_column_view.hpp>
 
+#include <map>
 #include <vector>
 
 using path_instruction_type = spark_rapids_jni::path_instruction_type;
 
+namespace spark_rapids_jni {
+json_schema_element read_schema_element(int& index,
+                                        cudf::jni::native_jstringArray const& names,
+                                        cudf::jni::native_jintArray const& children,
+                                        cudf::jni::native_jintArray const& types,
+                                        cudf::jni::native_jintArray const& scales)
+{
+  // printf("JNI line %d\n", __LINE__);
+  // fflush(stdout);
+
+  auto d_type = cudf::data_type{static_cast<cudf::type_id>(types[index]), scales[index]};
+  if (d_type.id() == cudf::type_id::STRUCT || d_type.id() == cudf::type_id::LIST) {
+    // printf("JNI line %d\n", __LINE__);
+    // fflush(stdout);
+
+    std::vector<std::pair<std::string, json_schema_element>> child_elems;
+    int num_children = children[index];
+    // go to the next entry, so recursion can parse it.
+    index++;
+    for (int i = 0; i < num_children; i++) {
+      // printf("JNI line %d\n", __LINE__);
+      // fflush(stdout);
+
+      auto const name = std::string{names.get(index).get()};
+      child_elems.emplace_back(name, read_schema_element(index, names, children, types, scales));
+    }
+    return json_schema_element{d_type, std::move(child_elems)};
+  } else {
+    // printf("JNI line %d\n", __LINE__);
+    // printf("children size: %d, idx = %d\n", children.size(), index);
+    // fflush(stdout);
+
+    if (children[index] != 0) {
+      throw std::invalid_argument("found children for a type that should have none");
+    }
+    // go to the next entry before returning...
+    index++;
+    // printf("JNI line %d\n", __LINE__);
+    // fflush(stdout);
+    return json_schema_element{d_type, {}};
+  }
+}
+}  // namespace spark_rapids_jni
+
 extern "C" {
 
 JNIEXPORT jint JNICALL Java_com_nvidia_spark_rapids_jni_JSONUtils_getMaxJSONPathDepth(JNIEnv* env,
@@ -154,4 +200,99 @@ JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_JSONUtils_extractRawMap
   }
   CATCH_STD(env, 0);
 }
+
+JNIEXPORT jlongArray JNICALL
+Java_com_nvidia_spark_rapids_jni_JSONUtils_fromJsonToStructs(JNIEnv* env,
+                                                             jclass,
+                                                             jlong j_input,
+                                                             jobjectArray j_col_names,
+                                                             jintArray j_num_children,
+                                                             jintArray j_types,
+                                                             jintArray j_scales,
+                                                             jboolean allow_leading_zero_numbers,
+                                                             jboolean allow_non_numeric_numbers,
+                                                             jboolean allow_unquoted_control_chars)
+{
+  JNI_NULL_CHECK(env, j_input, "j_input is null", 0);
+  JNI_NULL_CHECK(env, j_col_names, "j_col_names is null", 0);
+  JNI_NULL_CHECK(env, j_num_children, "j_num_children is null", 0);
+  JNI_NULL_CHECK(env, j_types, "j_types is null", 0);
+  JNI_NULL_CHECK(env, j_scales, "j_scales is null", 0);
+
+  try {
+    cudf::jni::auto_set_device(env);
+    cudf::jni::native_jstringArray n_col_names(env, j_col_names);
+    cudf::jni::native_jintArray n_types(env, j_types);
+    cudf::jni::native_jintArray n_scales(env, j_scales);
+    cudf::jni::native_jintArray n_children(env, j_num_children);
+
+    if (n_types.size() != n_scales.size()) {
+      JNI_THROW_NEW(env, cudf::jni::ILLEGAL_ARG_CLASS, "types and scales must match size", 0);
+    }
+    if (n_col_names.size() != n_types.size()) {
+      JNI_THROW_NEW(env, cudf::jni::ILLEGAL_ARG_CLASS, "types and column names must match size", 0);
+    }
+    if (n_children.size() != n_types.size()) {
+      JNI_THROW_NEW(env, cudf::jni::ILLEGAL_ARG_CLASS, "types and num children must match size", 0);
+    }
+
+    // printf("JNI line %d, size = %d\n", __LINE__, (int)n_types.size());
+    // fflush(stdout);
+
+    std::vector<std::pair<std::string, spark_rapids_jni::json_schema_element>> schema;
+    int idx = 0;
+    while (idx < n_types.size()) {
+      // printf("JNI line %d\n", __LINE__);
+      // fflush(stdout);
+
+      auto const name = std::string{n_col_names.get(idx).get()};
+      schema.emplace_back(
+        name,
+        spark_rapids_jni::read_schema_element(idx, n_col_names, n_children, n_types, n_scales));
+
+      // auto const name = n_col_names.get(at).get();
+      // printf("JNI line %d\n", __LINE__);
+      // fflush(stdout);
+
+      // auto child = cudf::jni::read_schema_element(at, n_children, n_col_names, n_types,
+      // n_scales); printf("JNI line %d\n", __LINE__); fflush(stdout);
+
+      // schema.emplace(name, std::move(child));
+    }
+    // printf("JNI line %d\n", __LINE__);
+    // fflush(stdout);
+
+    auto const input_cv = reinterpret_cast<cudf::column_view const*>(j_input);
+    auto output = spark_rapids_jni::from_json_to_structs(cudf::strings_column_view{*input_cv},
+                                                         schema,
+                                                         allow_leading_zero_numbers,
+                                                         allow_non_numeric_numbers,
+                                                         allow_unquoted_control_chars);
+
+    // printf("JNI line %d\n", __LINE__);
+    // fflush(stdout);
+
+    auto out_handles = cudf::jni::native_jlongArray(env, output.size());
+    std::transform(output.begin(), output.end(), out_handles.begin(), [](auto& col) {
+      return cudf::jni::release_as_jlong(col);
+    });
+    return out_handles.get_jArray();
+  }
+  CATCH_STD(env, 0);
+}
+
+JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_JSONUtils_isNullOrEmpty(JNIEnv* env,
+                                                                                 jclass,
+                                                                                 jlong j_input)
+{
+  JNI_NULL_CHECK(env, j_input, "j_input is null", 0);
+
+  try {
+    cudf::jni::auto_set_device(env);
+    auto const input_cv = reinterpret_cast<cudf::column_view const*>(j_input);
+    return cudf::jni::release_as_jlong(
+      spark_rapids_jni::is_null_or_empty(cudf::strings_column_view{*input_cv}));
+  }
+  CATCH_STD(env, 0);
+}
 }

src/main/cpp/src/from_json.hpp

@@ -16,19 +16,42 @@
 
 #pragma once
 
+#include <cudf/io/json.hpp>
 #include <cudf/strings/strings_column_view.hpp>
 #include <cudf/utilities/default_stream.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/resource_ref.hpp>
 
 #include <memory>
+#include <string>
+#include <vector>
 
 namespace spark_rapids_jni {
 
+struct json_schema_element {
+  cudf::data_type type;
+
+  std::vector<std::pair<std::string, json_schema_element>> child_types;
+};
+
 std::unique_ptr<cudf::column> from_json_to_raw_map(
   cudf::strings_column_view const& input,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
   rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
+std::vector<std::unique_ptr<cudf::column>> from_json_to_structs(
+  cudf::strings_column_view const& input,
+  std::vector<std::pair<std::string, json_schema_element>> const& schema,
+  bool allow_leading_zero_numbers,
+  bool allow_non_numeric_numbers,
+  bool allow_unquoted_control_chars,
+  rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+
+std::unique_ptr<cudf::column> is_null_or_empty(
+  cudf::strings_column_view const& input,
+  rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+
 }  // namespace spark_rapids_jni