From 9c5ff536f497d849b3373e920243fcf3f0e92656 Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Sat, 15 Aug 2026 16:01:58 +0530 Subject: [PATCH 1/5] Migrate FixedShapeTensorType deserialization to simdjson --- cpp/src/arrow/extension/fixed_shape_tensor.cc | 216 ++++++++++++++---- .../extension/tensor_extension_array_test.cc | 24 +- 2 files changed, 185 insertions(+), 55 deletions(-) diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index cd3d783479d6..6a9cbef9b70f 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -19,6 +19,8 @@ #include #include +#include + #include "arrow/extension/fixed_shape_tensor.h" #include "arrow/extension/tensor_internal.h" #include "arrow/scalar.h" @@ -26,16 +28,13 @@ #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 - -namespace rj = arrow::rapidjson; using ::arrow::json::JsonWriter; namespace arrow::extension { @@ -116,60 +115,189 @@ Result> FixedShapeTensorType::Deserialize( return Status::Invalid("Expected FixedSizeList storage type, got ", storage_type->ToString()); } + auto fsl_type = internal::checked_pointer_cast(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()) { + + simdjson::padded_string padded_json(serialized_data); + simdjson::ondemand::parser parser; + simdjson::ondemand::document document; + + if (auto error = parser.iterate(padded_json).get(document); + error != simdjson::SUCCESS) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } - std::vector 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::ondemand::object object; + if (auto error = document.get_object().get(object); error != simdjson::SUCCESS) { + return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } + std::vector shape; std::vector 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)); + std::vector dim_names; + + bool has_shape = false; + + for (auto field_result : object) { + ARROW_ASSIGN_OR_RAISE(auto field, internal::ResolveSimdjsonResult( + field_result, "Failed to iterate JSON object")); + + ARROW_ASSIGN_OR_RAISE( + auto key, internal::ResolveSimdjsonResult(field.unescaped_key(), + "Failed to get JSON object key")); + + auto value = field.value(); + + if (key == "shape") { + has_shape = true; + + simdjson::ondemand::json_type type; + if (auto error = value.type().get(type); + error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + return Status::Invalid("Invalid serialized JSON data: ", serialized_data); + } + + if (type != simdjson::ondemand::json_type::array) { + return Status::Invalid("shape must be an array, got ", + internal::JsonTypeName(type)); + } + + ARROW_ASSIGN_OR_RAISE(auto array, + internal::ResolveSimdjsonResult(value.get_array(), + "Failed to get shape array")); + + for (auto element_result : array) { + ARROW_ASSIGN_OR_RAISE(auto element, + internal::ResolveSimdjsonResult( + element_result, "Failed to iterate shape array")); + + ARROW_ASSIGN_OR_RAISE( + auto element_type, + internal::ResolveSimdjsonResult( + element.type(), "Failed to determine shape element JSON type")); + + if (element_type != simdjson::ondemand::json_type::number) { + return Status::Invalid("shape must contain integers, got ", + internal::JsonTypeName(element_type)); + } + + ARROW_ASSIGN_OR_RAISE( + auto number_type, + internal::ResolveSimdjsonResult(element.get_number_type(), + "Failed to determine shape number type")); + + if (number_type != simdjson::ondemand::number_type::signed_integer) { + return Status::Invalid("shape must contain integers, got number"); + } + + ARROW_ASSIGN_OR_RAISE( + auto number, internal::ResolveSimdjsonResult(element.get_int64(), + "Failed to get shape integer")); + + shape.emplace_back(number); + } + + } else if (key == "permutation") { + simdjson::ondemand::json_type type; + if (auto error = value.type().get(type); + error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + return Status::Invalid("Invalid serialized JSON data: ", serialized_data); + } + + if (type != simdjson::ondemand::json_type::array) { + return Status::Invalid("permutation must be an array, got ", + internal::JsonTypeName(type)); + } + + ARROW_ASSIGN_OR_RAISE( + auto array, internal::ResolveSimdjsonResult(value.get_array(), + "Failed to get permutation array")); + + for (auto element_result : array) { + ARROW_ASSIGN_OR_RAISE(auto element, + internal::ResolveSimdjsonResult( + element_result, "Failed to iterate permutation array")); + + ARROW_ASSIGN_OR_RAISE( + auto element_type, + internal::ResolveSimdjsonResult( + element.type(), "Failed to determine permutation element JSON type")); + + if (element_type != simdjson::ondemand::json_type::number) { + return Status::Invalid("permutation must contain integers, got ", + internal::JsonTypeName(element_type)); + } + + ARROW_ASSIGN_OR_RAISE(auto number_type, + internal::ResolveSimdjsonResult( + element.get_number_type(), + "Failed to determine permutation number type")); + + if (number_type != simdjson::ondemand::number_type::signed_integer) { + return Status::Invalid("permutation must contain integers, got number"); + } + + ARROW_ASSIGN_OR_RAISE( + auto number, internal::ResolveSimdjsonResult( + element.get_int64(), "Failed to get permutation integer")); + + permutation.emplace_back(number); + } + + } else if (key == "dim_names") { + simdjson::ondemand::json_type type; + if (auto error = value.type().get(type); + error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + return Status::Invalid("Invalid serialized JSON data: ", serialized_data); + } + + if (type != simdjson::ondemand::json_type::array) { + return Status::Invalid("dim_names must be an array, got ", + internal::JsonTypeName(type)); + } + + ARROW_ASSIGN_OR_RAISE( + auto array, internal::ResolveSimdjsonResult(value.get_array(), + "Failed to get dim_names array")); + + for (auto element_result : array) { + ARROW_ASSIGN_OR_RAISE(auto element, + internal::ResolveSimdjsonResult( + element_result, "Failed to iterate dim_names array")); + + ARROW_ASSIGN_OR_RAISE( + auto element_type, + internal::ResolveSimdjsonResult( + element.type(), "Failed to determine dim_names element JSON type")); + + if (element_type != simdjson::ondemand::json_type::string) { + return Status::Invalid("dim_names must contain strings, got ", + internal::JsonTypeName(element_type)); + } + + ARROW_ASSIGN_OR_RAISE(auto name, + internal::ResolveSimdjsonResult(element.get_string(), + "Failed to get dim_name")); + + dim_names.emplace_back(name); } - permutation.emplace_back(x.GetInt64()); } + } + + if (!has_shape) { + return Status::Invalid("Invalid serialized JSON data: ", serialized_data); + } + + if (!permutation.empty()) { if (shape.size() != permutation.size()) { return Status::Invalid("Invalid permutation"); } RETURN_NOT_OK(internal::IsPermutationValid(permutation)); } - std::vector 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 (shape.size() != dim_names.size()) { - return Status::Invalid("Invalid dim_names"); - } + + if (!dim_names.empty() && shape.size() != dim_names.size()) { + return Status::Invalid("Invalid dim_names"); } // Validate product of shape dimensions matches storage type list_size. @@ -180,11 +308,13 @@ Result> FixedShapeTensorType::Deserialize( const auto& fst_type = internal::checked_cast(*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; } diff --git a/cpp/src/arrow/extension/tensor_extension_array_test.cc b/cpp/src/arrow/extension/tensor_extension_array_test.cc index 531fc3c01cf5..797a2165b8b0 100644 --- a/cpp/src/arrow/extension/tensor_extension_array_test.cc +++ b/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,16 @@ 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 values must be unique integers in [0, N-1] CheckDeserializationRaises(ext_type_, storage_type, @@ -269,13 +269,13 @@ 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"); } TEST_F(TestFixedShapeTensorType, MakeValidatesShape) { From 3332bfbe640d350960c27c6c202c5544fda058f6 Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Tue, 18 Aug 2026 15:53:45 +0530 Subject: [PATCH 2/5] Replace RapidJSON with simdjson in VariableShapeTensor --- .../extension/tensor_extension_array_test.cc | 16 +- .../arrow/extension/variable_shape_tensor.cc | 231 ++++++++++++++---- 2 files changed, 187 insertions(+), 60 deletions(-) diff --git a/cpp/src/arrow/extension/tensor_extension_array_test.cc b/cpp/src/arrow/extension/tensor_extension_array_test.cc index 797a2165b8b0..31578924cf0a 100644 --- a/cpp/src/arrow/extension/tensor_extension_array_test.cc +++ b/cpp/src/arrow/extension/tensor_extension_array_test.cc @@ -865,28 +865,28 @@ TEST_F(TestVariableShapeTensorType, MetadataSerializationRoundtrip) { // 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"); // 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"); // 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"); } TEST_F(TestVariableShapeTensorType, RoundtripBatch) { diff --git a/cpp/src/arrow/extension/variable_shape_tensor.cc b/cpp/src/arrow/extension/variable_shape_tensor.cc index 40171f909a9d..c697d14e5a3c 100644 --- a/cpp/src/arrow/extension/variable_shape_tensor.cc +++ b/cpp/src/arrow/extension/variable_shape_tensor.cc @@ -17,22 +17,21 @@ #include +#include + #include "arrow/extension/tensor_internal.h" #include "arrow/extension/variable_shape_tensor.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/scalar.h" #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 - -namespace rj = arrow::rapidjson; using ::arrow::json::JsonWriter; namespace arrow::extension { @@ -155,62 +154,190 @@ Result> VariableShapeTensorType::Deserialize( internal::checked_cast(*storage_type->field(1)->type()) .list_size(); - rj::Document document; - if (document.Parse(serialized_data.data(), serialized_data.length()).HasParseError() || - !document.IsObject()) { + simdjson::padded_string padded_json(serialized_data); + simdjson::ondemand::parser parser; + simdjson::ondemand::document document; + + if (auto error = parser.iterate(padded_json).get(document); + error != simdjson::SUCCESS) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } - std::vector 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)); - } - permutation.reserve(ndim); - 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()); - } - RETURN_NOT_OK(internal::IsPermutationValid(permutation)); + simdjson::ondemand::object object; + if (auto error = document.get_object().get(object); error != simdjson::SUCCESS) { + return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } + + std::vector permutation; std::vector 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)); + std::vector> uniform_shape; + + for (auto field_result : object) { + if (field_result.error() != simdjson::SUCCESS) { + return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } - dim_names.reserve(ndim); - for (const auto& x : json_dim_names.GetArray()) { - if (!x.IsString()) { - return Status::Invalid("dim_names must contain strings, got ", - internal::JsonTypeName(x)); + + ARROW_ASSIGN_OR_RAISE(auto field, internal::ResolveSimdjsonResult( + field_result, "Failed to iterate JSON object")); + + ARROW_ASSIGN_OR_RAISE( + auto key, internal::ResolveSimdjsonResult(field.unescaped_key(), + "Failed to get JSON object key")); + + auto value = field.value(); + + if (key == "permutation") { + simdjson::ondemand::json_type type; + if (auto error = value.type().get(type); + error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } - dim_names.emplace_back(x.GetString()); - } - } - std::vector> uniform_shape; - if (document.HasMember("uniform_shape")) { - const auto& json_uniform_shape = document["uniform_shape"]; - if (!json_uniform_shape.IsArray()) { - return Status::Invalid("uniform_shape must be an array, got ", - internal::JsonTypeName(json_uniform_shape)); - } - uniform_shape.reserve(ndim); - for (const auto& x : json_uniform_shape.GetArray()) { - if (x.IsNull()) { - uniform_shape.emplace_back(std::nullopt); - } else if (x.IsInt64()) { - uniform_shape.emplace_back(x.GetInt64()); - } else { - return Status::Invalid("uniform_shape must contain integers or nulls, got ", - internal::JsonTypeName(x)); + if (type != simdjson::ondemand::json_type::array) { + return Status::Invalid("permutation must be an array, got ", + internal::JsonTypeName(type)); + } + + ARROW_ASSIGN_OR_RAISE( + auto array, internal::ResolveSimdjsonResult(value.get_array(), + "Failed to get permutation array")); + + permutation.reserve(ndim); + + for (auto element_result : array) { + ARROW_ASSIGN_OR_RAISE(auto element, + internal::ResolveSimdjsonResult( + element_result, "Failed to iterate permutation array")); + + ARROW_ASSIGN_OR_RAISE( + auto element_type, + internal::ResolveSimdjsonResult( + element.type(), "Failed to determine permutation element JSON type")); + + if (element_type != simdjson::ondemand::json_type::number) { + return Status::Invalid("permutation must contain integers, got ", + internal::JsonTypeName(element_type)); + } + + ARROW_ASSIGN_OR_RAISE(auto number_type, + internal::ResolveSimdjsonResult( + element.get_number_type(), + "Failed to determine permutation number type")); + + if (number_type != simdjson::ondemand::number_type::signed_integer) { + return Status::Invalid("permutation must contain integers, got number"); + } + + ARROW_ASSIGN_OR_RAISE( + auto number, internal::ResolveSimdjsonResult( + element.get_int64(), "Failed to get permutation integer")); + + permutation.emplace_back(number); + } + + RETURN_NOT_OK(internal::IsPermutationValid(permutation)); + + } else if (key == "dim_names") { + simdjson::ondemand::json_type type; + if (auto error = value.type().get(type); + error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + return Status::Invalid("Invalid serialized JSON data: ", serialized_data); + } + + if (type != simdjson::ondemand::json_type::array) { + return Status::Invalid("dim_names must be an array, got ", + internal::JsonTypeName(type)); + } + + ARROW_ASSIGN_OR_RAISE( + auto array, internal::ResolveSimdjsonResult(value.get_array(), + "Failed to get dim_names array")); + + dim_names.reserve(ndim); + + for (auto element_result : array) { + ARROW_ASSIGN_OR_RAISE(auto element, + internal::ResolveSimdjsonResult( + element_result, "Failed to iterate dim_names array")); + + ARROW_ASSIGN_OR_RAISE( + auto element_type, + internal::ResolveSimdjsonResult( + element.type(), "Failed to determine dim_names element JSON type")); + + if (element_type != simdjson::ondemand::json_type::string) { + return Status::Invalid("dim_names must contain strings, got ", + internal::JsonTypeName(element_type)); + } + + ARROW_ASSIGN_OR_RAISE(auto name, + internal::ResolveSimdjsonResult(element.get_string(), + "Failed to get dim_name")); + + dim_names.emplace_back(name); + } + + if (dim_names.size() != static_cast(ndim)) { + return Status::Invalid("Invalid: dim_names"); + } + + } else if (key == "uniform_shape") { + simdjson::ondemand::json_type type; + if (auto error = value.type().get(type); + error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + return Status::Invalid("Invalid serialized JSON data: ", serialized_data); + } + + if (type != simdjson::ondemand::json_type::array) { + return Status::Invalid("uniform_shape must be an array, got ", + internal::JsonTypeName(type)); + } + + ARROW_ASSIGN_OR_RAISE(auto array, + internal::ResolveSimdjsonResult( + value.get_array(), "Failed to get uniform_shape array")); + + uniform_shape.reserve(ndim); + + for (auto element_result : array) { + ARROW_ASSIGN_OR_RAISE( + auto element, internal::ResolveSimdjsonResult( + element_result, "Failed to iterate uniform_shape array")); + + ARROW_ASSIGN_OR_RAISE( + auto element_type, + internal::ResolveSimdjsonResult( + element.type(), "Failed to determine uniform_shape element JSON type")); + + if (element_type == simdjson::ondemand::json_type::null) { + uniform_shape.emplace_back(std::nullopt); + continue; + } + + if (element_type != simdjson::ondemand::json_type::number) { + return Status::Invalid("uniform_shape must contain integers or nulls, got ", + internal::JsonTypeName(element_type)); + } + + ARROW_ASSIGN_OR_RAISE(auto number_type, + internal::ResolveSimdjsonResult( + element.get_number_type(), + "Failed to determine uniform_shape number type")); + + if (number_type != simdjson::ondemand::number_type::signed_integer) { + return Status::Invalid( + "uniform_shape must contain integers or nulls, got number"); + } + + ARROW_ASSIGN_OR_RAISE( + auto number, internal::ResolveSimdjsonResult( + element.get_int64(), "Failed to get uniform_shape integer")); + + uniform_shape.emplace_back(number); + } + + if (uniform_shape.size() != static_cast(ndim)) { + return Status::Invalid("Invalid: uniform_shape"); } } } From bbcaaeca969b8db411897c3d0045704982d1127d Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Wed, 19 Aug 2026 00:18:42 +0530 Subject: [PATCH 3/5] Remove obsolete RapidJSON tensor helper --- cpp/src/arrow/extension/tensor_internal.cc | 14 -------------- cpp/src/arrow/extension/tensor_internal.h | 7 ------- 2 files changed, 21 deletions(-) diff --git a/cpp/src/arrow/extension/tensor_internal.cc b/cpp/src/arrow/extension/tensor_internal.cc index e94ea9a1d181..d2965bc578f7 100644 --- a/cpp/src/arrow/extension/tensor_internal.cc +++ b/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 ComputeShapeProduct(std::span shape) { int64_t product = 1; for (const auto dim : shape) { diff --git a/cpp/src/arrow/extension/tensor_internal.h b/cpp/src/arrow/extension/tensor_internal.h index 19665bf2cd4c..b54945ad50ab 100644 --- a/cpp/src/arrow/extension/tensor_internal.h +++ b/cpp/src/arrow/extension/tensor_internal.h @@ -21,18 +21,11 @@ #include #include -#include "arrow/json/rapidjson_defs.h" // IWYU pragma: keep #include "arrow/result.h" #include "arrow/type_fwd.h" -#include - 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. From ec82dd0cdf2b9dab15be3e3f4073f25de546ec1e Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Wed, 19 Aug 2026 13:17:50 +0530 Subject: [PATCH 4/5] Address Feedback --- cpp/src/arrow/extension/fixed_shape_tensor.cc | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index 6a9cbef9b70f..5da2d6bf02c9 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -123,13 +123,14 @@ Result> FixedShapeTensorType::Deserialize( simdjson::ondemand::parser parser; simdjson::ondemand::document document; - if (auto error = parser.iterate(padded_json).get(document); - error != simdjson::SUCCESS) { + auto error = parser.iterate(padded_json).get(document); + if (error != simdjson::SUCCESS) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } simdjson::ondemand::object object; - if (auto error = document.get_object().get(object); error != simdjson::SUCCESS) { + error = document.get_object().get(object); + if (error != simdjson::SUCCESS) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } @@ -153,8 +154,8 @@ Result> FixedShapeTensorType::Deserialize( has_shape = true; simdjson::ondemand::json_type type; - if (auto error = value.type().get(type); - error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + auto error = value.type().get(type); + if (error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } @@ -200,8 +201,8 @@ Result> FixedShapeTensorType::Deserialize( } else if (key == "permutation") { simdjson::ondemand::json_type type; - if (auto error = value.type().get(type); - error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + auto error = value.type().get(type); + if (error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } @@ -247,8 +248,8 @@ Result> FixedShapeTensorType::Deserialize( } else if (key == "dim_names") { simdjson::ondemand::json_type type; - if (auto error = value.type().get(type); - error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + auto error = value.type().get(type); + if (error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } From 069be12ade0f7e59d16b5c0a782bfbb705e0e173 Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Thu, 20 Aug 2026 14:02:02 +0530 Subject: [PATCH 5/5] use ResolveSimdjsonResult --- cpp/src/arrow/extension/fixed_shape_tensor.cc | 37 ++++++++-------- .../arrow/extension/variable_shape_tensor.cc | 42 +++++++++---------- 2 files changed, 37 insertions(+), 42 deletions(-) diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index 5da2d6bf02c9..fa90c0a7c940 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -121,18 +121,14 @@ Result> FixedShapeTensorType::Deserialize( simdjson::padded_string padded_json(serialized_data); simdjson::ondemand::parser parser; - simdjson::ondemand::document document; - auto error = parser.iterate(padded_json).get(document); - if (error != simdjson::SUCCESS) { - return Status::Invalid("Invalid serialized JSON data: ", serialized_data); - } + ARROW_ASSIGN_OR_RAISE(auto document, + internal::ResolveSimdjsonResult(parser.iterate(padded_json), + "Invalid serialized JSON data")); - simdjson::ondemand::object object; - error = document.get_object().get(object); - if (error != simdjson::SUCCESS) { - return Status::Invalid("Invalid serialized JSON data: ", serialized_data); - } + ARROW_ASSIGN_OR_RAISE(auto object, + internal::ResolveSimdjsonResult(document.get_object(), + "Invalid serialized JSON data")); std::vector shape; std::vector permutation; @@ -153,9 +149,10 @@ Result> FixedShapeTensorType::Deserialize( if (key == "shape") { has_shape = true; - simdjson::ondemand::json_type type; - auto error = value.type().get(type); - if (error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( + value.type(), "Invalid serialized JSON data")); + + if (type == simdjson::ondemand::json_type::unknown) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } @@ -200,9 +197,10 @@ Result> FixedShapeTensorType::Deserialize( } } else if (key == "permutation") { - simdjson::ondemand::json_type type; - auto error = value.type().get(type); - if (error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( + value.type(), "Invalid serialized JSON data")); + + if (type == simdjson::ondemand::json_type::unknown) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } @@ -247,9 +245,10 @@ Result> FixedShapeTensorType::Deserialize( } } else if (key == "dim_names") { - simdjson::ondemand::json_type type; - auto error = value.type().get(type); - if (error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( + value.type(), "Invalid serialized JSON data")); + + if (type == simdjson::ondemand::json_type::unknown) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } diff --git a/cpp/src/arrow/extension/variable_shape_tensor.cc b/cpp/src/arrow/extension/variable_shape_tensor.cc index c697d14e5a3c..17324fed2f64 100644 --- a/cpp/src/arrow/extension/variable_shape_tensor.cc +++ b/cpp/src/arrow/extension/variable_shape_tensor.cc @@ -156,29 +156,22 @@ Result> VariableShapeTensorType::Deserialize( simdjson::padded_string padded_json(serialized_data); simdjson::ondemand::parser parser; - simdjson::ondemand::document document; - if (auto error = parser.iterate(padded_json).get(document); - error != simdjson::SUCCESS) { - return Status::Invalid("Invalid serialized JSON data: ", serialized_data); - } + ARROW_ASSIGN_OR_RAISE(auto document, + internal::ResolveSimdjsonResult(parser.iterate(padded_json), + "Invalid serialized JSON data")); - simdjson::ondemand::object object; - if (auto error = document.get_object().get(object); error != simdjson::SUCCESS) { - return Status::Invalid("Invalid serialized JSON data: ", serialized_data); - } + ARROW_ASSIGN_OR_RAISE(auto object, + internal::ResolveSimdjsonResult(document.get_object(), + "Invalid serialized JSON data")); std::vector permutation; std::vector dim_names; std::vector> uniform_shape; for (auto field_result : object) { - if (field_result.error() != simdjson::SUCCESS) { - return Status::Invalid("Invalid serialized JSON data: ", serialized_data); - } - ARROW_ASSIGN_OR_RAISE(auto field, internal::ResolveSimdjsonResult( - field_result, "Failed to iterate JSON object")); + field_result, "Invalid serialized JSON data")); ARROW_ASSIGN_OR_RAISE( auto key, internal::ResolveSimdjsonResult(field.unescaped_key(), @@ -187,9 +180,10 @@ Result> VariableShapeTensorType::Deserialize( auto value = field.value(); if (key == "permutation") { - simdjson::ondemand::json_type type; - if (auto error = value.type().get(type); - error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( + value.type(), "Invalid serialized JSON data")); + + if (type == simdjson::ondemand::json_type::unknown) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } @@ -238,9 +232,10 @@ Result> VariableShapeTensorType::Deserialize( RETURN_NOT_OK(internal::IsPermutationValid(permutation)); } else if (key == "dim_names") { - simdjson::ondemand::json_type type; - if (auto error = value.type().get(type); - error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( + value.type(), "Invalid serialized JSON data")); + + if (type == simdjson::ondemand::json_type::unknown) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); } @@ -282,9 +277,10 @@ Result> VariableShapeTensorType::Deserialize( } } else if (key == "uniform_shape") { - simdjson::ondemand::json_type type; - if (auto error = value.type().get(type); - error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) { + ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( + value.type(), "Invalid serialized JSON data")); + + if (type == simdjson::ondemand::json_type::unknown) { return Status::Invalid("Invalid serialized JSON data: ", serialized_data); }