From 9c5ff536f497d849b3373e920243fcf3f0e92656 Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Sat, 15 Aug 2026 16:01:58 +0530 Subject: [PATCH] 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) {