GH-50901: [C++] Replace RapidJSON with simdjson in tensor extension types

cpp/src/arrow/extension/fixed_shape_tensor.cc

@@ -19,23 +19,22 @@
 #include <numeric>
 #include <sstream>
 
+#include <simdjson.h>
+
 #include "arrow/extension/fixed_shape_tensor.h"
 #include "arrow/extension/tensor_internal.h"
 #include "arrow/scalar.h"
 
 #include "arrow/array/array_nested.h"
 #include "arrow/array/array_primitive.h"
 #include "arrow/json/json_writer_internal.h"
-#include "arrow/json/rapidjson_defs.h"  // IWYU pragma: keep
 #include "arrow/tensor.h"
 #include "arrow/util/logging_internal.h"
 #include "arrow/util/print_internal.h"
+#include "arrow/util/simdjson_internal.h"
 #include "arrow/util/sort_internal.h"
 #include "arrow/util/string.h"
 
-#include <rapidjson/document.h>
-
-namespace rj = arrow::rapidjson;
 using ::arrow::json::JsonWriter;
 
 namespace arrow::extension {
@@ -116,57 +115,39 @@ Result<std::shared_ptr<DataType>> FixedShapeTensorType::Deserialize(
     return Status::Invalid("Expected FixedSizeList storage type, got ",
                            storage_type->ToString());
   }
+
   auto fsl_type = internal::checked_pointer_cast<FixedSizeListType>(storage_type);
   auto value_type = fsl_type->value_type();
-  rj::Document document;
-  if (document.Parse(serialized_data.data(), serialized_data.length()).HasParseError() ||
-      !document.IsObject() || !document.HasMember("shape") ||
-      !document["shape"].IsArray()) {
-    return Status::Invalid("Invalid serialized JSON data: ", serialized_data);
-  }
 
-  std::vector<int64_t> shape;
-  for (const auto& x : document["shape"].GetArray()) {
-    if (!x.IsInt64()) {
-      return Status::Invalid("shape must contain integers, got ",
-                             internal::JsonTypeName(x));
-    }
-    shape.emplace_back(x.GetInt64());
-  }
+  simdjson::dom::parser parser;
+  ARROW_ASSIGN_OR_RAISE(auto object, internal::ParseJsonObject(parser, serialized_data));
+
+  ARROW_ASSIGN_OR_RAISE(auto shape_value,
+                        internal::ResolveSimdjsonResult(object.at_key("shape"),
+                                                        "Invalid serialized JSON data"));
+  ARROW_ASSIGN_OR_RAISE(auto shape, internal::GetJsonIntArray(shape_value, "shape"));
+  ARROW_ASSIGN_OR_RAISE(auto permutation_value,
+                        internal::GetOptionalJsonField(object, "permutation"));
 
   std::vector<int64_t> permutation;
-  if (document.HasMember("permutation")) {
-    const auto& json_permutation = document["permutation"];
-    if (!json_permutation.IsArray()) {
-      return Status::Invalid("permutation must be an array, got ",
-                             internal::JsonTypeName(json_permutation));
-    }
-    for (const auto& x : json_permutation.GetArray()) {
-      if (!x.IsInt64()) {
-        return Status::Invalid("permutation must contain integers, got ",
-                               internal::JsonTypeName(x));
-      }
-      permutation.emplace_back(x.GetInt64());
-    }
+  if (permutation_value.has_value()) {
+    ARROW_ASSIGN_OR_RAISE(permutation,
+                          internal::GetJsonIntArray(*permutation_value, "permutation"));
+
     if (shape.size() != permutation.size()) {
       return Status::Invalid("Invalid permutation");
     }
     RETURN_NOT_OK(internal::IsPermutationValid(permutation));
   }
+
+  ARROW_ASSIGN_OR_RAISE(auto dim_names_value,
+                        internal::GetOptionalJsonField(object, "dim_names"));
+
   std::vector<std::string> dim_names;
-  if (document.HasMember("dim_names")) {
-    const auto& json_dim_names = document["dim_names"];
-    if (!json_dim_names.IsArray()) {
-      return Status::Invalid("dim_names must be an array, got ",
-                             internal::JsonTypeName(json_dim_names));
-    }
-    for (const auto& x : json_dim_names.GetArray()) {
-      if (!x.IsString()) {
-        return Status::Invalid("dim_names must contain strings, got ",
-                               internal::JsonTypeName(x));
-      }
-      dim_names.emplace_back(x.GetString());
-    }
+  if (dim_names_value.has_value()) {
+    ARROW_ASSIGN_OR_RAISE(dim_names,
+                          internal::GetJsonStringArray(*dim_names_value, "dim_names"));
+
     if (shape.size() != dim_names.size()) {
       return Status::Invalid("Invalid dim_names");
     }
@@ -177,14 +158,18 @@ Result<std::shared_ptr<DataType>> FixedShapeTensorType::Deserialize(
   // (type mismatches, size mismatches) are reported first.
   ARROW_ASSIGN_OR_RAISE(auto ext_type, FixedShapeTensorType::Make(
                                            value_type, shape, permutation, dim_names));
+
   const auto& fst_type = internal::checked_cast<const FixedShapeTensorType&>(*ext_type);
+
   ARROW_ASSIGN_OR_RAISE(const int64_t expected_size,
                         internal::ComputeShapeProduct(fst_type.shape()));
+
   if (expected_size != fsl_type->list_size()) {
     return Status::Invalid("Product of shape dimensions (", expected_size,
                            ") does not match FixedSizeList size (", fsl_type->list_size(),
                            ")");
   }
+
   return ext_type;
 }
 

cpp/src/arrow/extension/tensor_extension_array_test.cc

@@ -223,15 +223,15 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) {
   // Validate shape values must be integers. Error message should include the
   // JSON type name of the offending value.
   CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3.5,4]})",
-                             "shape must contain integers, got Number");
+                             "shape must contain integers, got number");
   CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":["3","4"]})",
-                             "shape must contain integers, got String");
+                             "shape must contain integers, got string");
   CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[null]})",
-                             "shape must contain integers, got Null");
+                             "shape must contain integers, got null");
   CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[true]})",
-                             "shape must contain integers, got True");
+                             "shape must contain integers, got boolean");
   CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[false]})",
-                             "shape must contain integers, got False");
+                             "shape must contain integers, got boolean");
 
   // Validate shape values must be non-negative
   CheckDeserializationRaises(ext_type_, fixed_size_list(int64(), 1), R"({"shape":[-1]})",
@@ -244,16 +244,20 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) {
   // Validate permutation member must be an array with integer values
   CheckDeserializationRaises(ext_type_, storage_type,
                              R"({"shape":[3,4],"permutation":"invalid"})",
-                             "permutation must be an array, got String");
+                             "permutation must be an array, got string");
   CheckDeserializationRaises(ext_type_, storage_type,
                              R"({"shape":[3,4],"permutation":{"a":1}})",
-                             "permutation must be an array, got Object");
+                             "permutation must be an array, got object");
   CheckDeserializationRaises(ext_type_, storage_type,
                              R"({"shape":[3,4],"permutation":[1.5,0.5]})",
-                             "permutation must contain integers, got Number");
+                             "permutation must contain integers, got number");
   CheckDeserializationRaises(ext_type_, storage_type,
                              R"({"shape":[3,4],"permutation":["a","b"]})",
-                             "permutation must contain integers, got String");
+                             "permutation must contain integers, got string");
+  // Validate permutation member must be an array with integer values
+  CheckDeserializationRaises(ext_type_, storage_type,
+                             R"({"shape":[3,4],"permutation":[]})",
+                             "Invalid permutation");
 
   // Validate permutation values must be unique integers in [0, N-1]
   CheckDeserializationRaises(ext_type_, storage_type,
@@ -269,13 +273,15 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) {
   // Validate dim_names member must be an array with string values
   CheckDeserializationRaises(ext_type_, storage_type,
                              R"({"shape":[3,4],"dim_names":"invalid"})",
-                             "dim_names must be an array, got String");
+                             "dim_names must be an array, got string");
   CheckDeserializationRaises(ext_type_, storage_type,
                              R"({"shape":[3,4],"dim_names":[1,2]})",
-                             "dim_names must contain strings, got Number");
+                             "dim_names must contain strings, got number");
   CheckDeserializationRaises(ext_type_, storage_type,
                              R"({"shape":[3,4],"dim_names":[null,null]})",
-                             "dim_names must contain strings, got Null");
+                             "dim_names must contain strings, got null");
+  CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"dim_names":[]})",
+                             "Invalid dim_names");
 }
 
 TEST_F(TestFixedShapeTensorType, MakeValidatesShape) {
@@ -858,35 +864,41 @@ TEST_F(TestVariableShapeTensorType, MetadataSerializationRoundtrip) {
   CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":(3,4)})",
                              "Invalid serialized JSON data");
   CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[1,0]})",
-                             "Invalid: permutation");
+                             "Invalid permutation");
   CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":["x","y"]})",
-                             "Invalid: dim_names");
+                             "Invalid dim_names");
 
   // Validate permutation member must be an array with integer values. Error
   // message should include the JSON type name of the offending value.
   CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":"invalid"})",
-                             "permutation must be an array, got String");
+                             "permutation must be an array, got string");
   CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[1.5,0.5,2.5]})",
-                             "permutation must contain integers, got Number");
+                             "permutation must contain integers, got number");
   CheckDeserializationRaises(ext_type_, storage_type,
                              R"({"permutation":[null,null,null]})",
-                             "permutation must contain integers, got Null");
+                             "permutation must contain integers, got null");
+  CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[]})",
+                             "Invalid permutation");
 
   // Validate dim_names member must be an array with string values
   CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":"invalid"})",
-                             "dim_names must be an array, got String");
+                             "dim_names must be an array, got string");
   CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":[1,2,3]})",
-                             "dim_names must contain strings, got Number");
+                             "dim_names must contain strings, got number");
+  CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":[]})",
+                             "Invalid dim_names");
 
   // Validate uniform_shape member must be an array with integer-or-null values
   CheckDeserializationRaises(ext_type_, storage_type, R"({"uniform_shape":"invalid"})",
-                             "uniform_shape must be an array, got String");
+                             "uniform_shape must be an array, got string");
   CheckDeserializationRaises(ext_type_, storage_type,
                              R"({"uniform_shape":[1.5,null,null]})",
-                             "uniform_shape must contain integers or nulls, got Number");
+                             "uniform_shape must contain integers or nulls, got number");
   CheckDeserializationRaises(ext_type_, storage_type,
                              R"({"uniform_shape":["x",null,null]})",
-                             "uniform_shape must contain integers or nulls, got String");
+                             "uniform_shape must contain integers or nulls, got string");
+  CheckDeserializationRaises(ext_type_, storage_type, R"({"uniform_shape":[]})",
+                             "Invalid uniform_shape");
 }
 
 TEST_F(TestVariableShapeTensorType, RoundtripBatch) {

cpp/src/arrow/extension/tensor_internal.cc

@@ -30,20 +30,6 @@
 
 namespace arrow::internal {
 
-namespace {
-
-// Names indexed by rapidjson::Type enum value:
-// kNullType=0, kFalseType=1, kTrueType=2, kObjectType=3,
-// kArrayType=4, kStringType=5, kNumberType=6.
-constexpr const char* kJsonTypeNames[] = {"Null",  "False",  "True",  "Object",
-                                          "Array", "String", "Number"};
-
-}  // namespace
-
-const char* JsonTypeName(const ::arrow::rapidjson::Value& v) {
-  return kJsonTypeNames[v.GetType()];
-}
-
 Result<int64_t> ComputeShapeProduct(std::span<const int64_t> shape) {
   int64_t product = 1;
   for (const auto dim : shape) {

cpp/src/arrow/extension/tensor_internal.h

@@ -21,18 +21,11 @@
 #include <span>
 #include <vector>
 
-#include "arrow/json/rapidjson_defs.h"  // IWYU pragma: keep
 #include "arrow/result.h"
 #include "arrow/type_fwd.h"
 
-#include <rapidjson/document.h>
-
 namespace arrow::internal {
 
-/// \brief Return the name of a RapidJSON value's type (e.g., "Null", "Array", "Number").
-ARROW_EXPORT
-const char* JsonTypeName(const ::arrow::rapidjson::Value& v);
-
 /// \brief Compute the product of the given shape dimensions.
 ///
 /// Returns Status::Invalid if the product would overflow int64_t.