From 9c5ff536f497d849b3373e920243fcf3f0e92656 Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Sat, 15 Aug 2026 16:01:58 +0530 Subject: [PATCH 01/12] 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 02/12] 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 03/12] 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 04/12] 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 05/12] 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); } From e6dd561ab3fe47c03918f38b9296b763fd02043e Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Fri, 21 Aug 2026 11:19:17 +0530 Subject: [PATCH 06/12] remove unknown --- cpp/src/arrow/extension/fixed_shape_tensor.cc | 12 ------------ .../arrow/extension/tensor_extension_array_test.cc | 2 +- cpp/src/arrow/extension/variable_shape_tensor.cc | 12 ------------ 3 files changed, 1 insertion(+), 25 deletions(-) diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index fa90c0a7c940..a1d3f52f3a59 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -152,10 +152,6 @@ Result> FixedShapeTensorType::Deserialize( 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); - } - if (type != simdjson::ondemand::json_type::array) { return Status::Invalid("shape must be an array, got ", internal::JsonTypeName(type)); @@ -200,10 +196,6 @@ Result> FixedShapeTensorType::Deserialize( 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); - } - if (type != simdjson::ondemand::json_type::array) { return Status::Invalid("permutation must be an array, got ", internal::JsonTypeName(type)); @@ -248,10 +240,6 @@ Result> FixedShapeTensorType::Deserialize( 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); - } - if (type != simdjson::ondemand::json_type::array) { return Status::Invalid("dim_names must be an array, got ", internal::JsonTypeName(type)); diff --git a/cpp/src/arrow/extension/tensor_extension_array_test.cc b/cpp/src/arrow/extension/tensor_extension_array_test.cc index 31578924cf0a..dba48e0cc0c8 100644 --- a/cpp/src/arrow/extension/tensor_extension_array_test.cc +++ b/cpp/src/arrow/extension/tensor_extension_array_test.cc @@ -212,7 +212,7 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) { CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":["x","y"]})", "Invalid serialized JSON data"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":(3,4)})", - "Invalid serialized JSON data"); + "shape must be an array, got unknown"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"permutation":[1,0,2]})", "Invalid permutation"); diff --git a/cpp/src/arrow/extension/variable_shape_tensor.cc b/cpp/src/arrow/extension/variable_shape_tensor.cc index 17324fed2f64..d06ea4fcfc6d 100644 --- a/cpp/src/arrow/extension/variable_shape_tensor.cc +++ b/cpp/src/arrow/extension/variable_shape_tensor.cc @@ -183,10 +183,6 @@ Result> VariableShapeTensorType::Deserialize( 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); - } - if (type != simdjson::ondemand::json_type::array) { return Status::Invalid("permutation must be an array, got ", internal::JsonTypeName(type)); @@ -235,10 +231,6 @@ Result> VariableShapeTensorType::Deserialize( 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); - } - if (type != simdjson::ondemand::json_type::array) { return Status::Invalid("dim_names must be an array, got ", internal::JsonTypeName(type)); @@ -280,10 +272,6 @@ Result> VariableShapeTensorType::Deserialize( 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); - } - if (type != simdjson::ondemand::json_type::array) { return Status::Invalid("uniform_shape must be an array, got ", internal::JsonTypeName(type)); From e234967e02263417d210c11fbde99577a2de3420 Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Fri, 21 Aug 2026 14:27:50 +0530 Subject: [PATCH 07/12] Add JsonTypeName --- cpp/src/arrow/extension/fixed_shape_tensor.cc | 6 ++++-- .../extension/tensor_extension_array_test.cc | 19 ++++++++++--------- .../arrow/extension/variable_shape_tensor.cc | 7 ++++--- cpp/src/arrow/util/simdjson_internal.h | 14 ++++++++++++++ 4 files changed, 32 insertions(+), 14 deletions(-) diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index a1d3f52f3a59..f4370afb553a 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -182,7 +182,8 @@ Result> FixedShapeTensorType::Deserialize( "Failed to determine shape number type")); if (number_type != simdjson::ondemand::number_type::signed_integer) { - return Status::Invalid("shape must contain integers, got number"); + return Status::Invalid("shape must contain integers, got ", + internal::JsonNumberTypeName(number_type)); } ARROW_ASSIGN_OR_RAISE( @@ -226,7 +227,8 @@ Result> FixedShapeTensorType::Deserialize( "Failed to determine permutation number type")); if (number_type != simdjson::ondemand::number_type::signed_integer) { - return Status::Invalid("permutation must contain integers, got number"); + return Status::Invalid("permutation must contain integers, got ", + internal::JsonNumberTypeName(number_type)); } ARROW_ASSIGN_OR_RAISE( diff --git a/cpp/src/arrow/extension/tensor_extension_array_test.cc b/cpp/src/arrow/extension/tensor_extension_array_test.cc index dba48e0cc0c8..78c2ddeeeeff 100644 --- a/cpp/src/arrow/extension/tensor_extension_array_test.cc +++ b/cpp/src/arrow/extension/tensor_extension_array_test.cc @@ -223,7 +223,7 @@ 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 floating-point number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":["3","4"]})", "shape must contain integers, got string"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[null]})", @@ -248,9 +248,9 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) { CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"permutation":{"a":1}})", "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"); + CheckDeserializationRaises( + ext_type_, storage_type, R"({"shape":[3,4],"permutation":[1.5,0.5]})", + "permutation must contain integers, got floating-point number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"permutation":["a","b"]})", "permutation must contain integers, got string"); @@ -866,8 +866,9 @@ TEST_F(TestVariableShapeTensorType, MetadataSerializationRoundtrip) { // 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"); - CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[1.5,0.5,2.5]})", - "permutation must contain integers, got number"); + CheckDeserializationRaises( + ext_type_, storage_type, R"({"permutation":[1.5,0.5,2.5]})", + "permutation must contain integers, got floating-point number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[null,null,null]})", "permutation must contain integers, got null"); @@ -881,9 +882,9 @@ TEST_F(TestVariableShapeTensorType, MetadataSerializationRoundtrip) { // 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"); - CheckDeserializationRaises(ext_type_, storage_type, - R"({"uniform_shape":[1.5,null,null]})", - "uniform_shape must contain integers or nulls, got number"); + CheckDeserializationRaises( + ext_type_, storage_type, R"({"uniform_shape":[1.5,null,null]})", + "uniform_shape must contain integers or nulls, got floating-point number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"uniform_shape":["x",null,null]})", "uniform_shape must contain integers or nulls, got string"); diff --git a/cpp/src/arrow/extension/variable_shape_tensor.cc b/cpp/src/arrow/extension/variable_shape_tensor.cc index d06ea4fcfc6d..132c29193f79 100644 --- a/cpp/src/arrow/extension/variable_shape_tensor.cc +++ b/cpp/src/arrow/extension/variable_shape_tensor.cc @@ -215,7 +215,8 @@ Result> VariableShapeTensorType::Deserialize( "Failed to determine permutation number type")); if (number_type != simdjson::ondemand::number_type::signed_integer) { - return Status::Invalid("permutation must contain integers, got number"); + return Status::Invalid("permutation must contain integers, got ", + internal::JsonNumberTypeName(number_type)); } ARROW_ASSIGN_OR_RAISE( @@ -309,8 +310,8 @@ Result> VariableShapeTensorType::Deserialize( "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"); + return Status::Invalid("uniform_shape must contain integers or nulls, got ", + internal::JsonNumberTypeName(number_type)); } ARROW_ASSIGN_OR_RAISE( diff --git a/cpp/src/arrow/util/simdjson_internal.h b/cpp/src/arrow/util/simdjson_internal.h index 1badd99938b1..f31dd1e35d0d 100644 --- a/cpp/src/arrow/util/simdjson_internal.h +++ b/cpp/src/arrow/util/simdjson_internal.h @@ -328,5 +328,19 @@ inline Status ValidateJsonDocument(simdjson::ondemand::parser& parser, return ConsumeJsonValue(value); } +inline const char* JsonNumberTypeName(simdjson::ondemand::number_type type) { + switch (type) { + case simdjson::ondemand::number_type::signed_integer: + return "signed integer"; + case simdjson::ondemand::number_type::unsigned_integer: + return "unsigned integer"; + case simdjson::ondemand::number_type::floating_point_number: + return "floating-point number"; + case simdjson::ondemand::number_type::big_integer: + return "big integer"; + } + return "unknown"; +} + } // namespace internal } // namespace arrow From aeed4a6057708f4c6790decbca19c874d4eadbc2 Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Mon, 24 Aug 2026 18:30:42 +0530 Subject: [PATCH 08/12] Address Feedback --- cpp/src/arrow/extension/fixed_shape_tensor.cc | 184 +++-------------- .../extension/tensor_extension_array_test.cc | 37 ++-- cpp/src/arrow/extension/tensor_internal.cc | 118 +++++++++++ cpp/src/arrow/extension/tensor_internal.h | 18 ++ .../arrow/extension/variable_shape_tensor.cc | 190 +++--------------- 5 files changed, 213 insertions(+), 334 deletions(-) diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index f4370afb553a..ac3531fd200d 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -119,183 +119,45 @@ Result> FixedShapeTensorType::Deserialize( auto fsl_type = internal::checked_pointer_cast(storage_type); auto value_type = fsl_type->value_type(); - simdjson::padded_string padded_json(serialized_data); - simdjson::ondemand::parser parser; + simdjson::dom::parser parser; + ARROW_ASSIGN_OR_RAISE(auto object, internal::ParseJsonObject(parser, serialized_data)); - ARROW_ASSIGN_OR_RAISE(auto document, - internal::ResolveSimdjsonResult(parser.iterate(padded_json), + 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")); - ARROW_ASSIGN_OR_RAISE(auto object, - internal::ResolveSimdjsonResult(document.get_object(), - "Invalid serialized JSON data")); - - std::vector shape; std::vector permutation; - 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; - - ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( - value.type(), "Invalid serialized JSON 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 ", - internal::JsonNumberTypeName(number_type)); - } - - ARROW_ASSIGN_OR_RAISE( - auto number, internal::ResolveSimdjsonResult(element.get_int64(), - "Failed to get shape integer")); - - shape.emplace_back(number); - } - - } else if (key == "permutation") { - ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( - value.type(), "Invalid serialized JSON 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 ", - internal::JsonNumberTypeName(number_type)); - } - - 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") { - ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( - value.type(), "Invalid serialized JSON 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); - } - } - } + if (permutation_value.has_value()) { + ARROW_ASSIGN_OR_RAISE(permutation, + internal::GetJsonIntArray(*permutation_value, "permutation")); - 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)); } - if (!dim_names.empty() && shape.size() != dim_names.size()) { - return Status::Invalid("Invalid dim_names"); + ARROW_ASSIGN_OR_RAISE(auto dim_names_value, + internal::GetOptionalJsonField(object, "dim_names")); + + std::vector dim_names; + 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"); + } } - // Validate product of shape dimensions matches storage type list_size. - // This check is intentionally after field parsing so that metadata-level errors - // (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(*ext_type); + ARROW_ASSIGN_OR_RAISE(const int64_t expected_size, internal::ComputeShapeProduct(fst_type.shape())); diff --git a/cpp/src/arrow/extension/tensor_extension_array_test.cc b/cpp/src/arrow/extension/tensor_extension_array_test.cc index 78c2ddeeeeff..29147df6ba6a 100644 --- a/cpp/src/arrow/extension/tensor_extension_array_test.cc +++ b/cpp/src/arrow/extension/tensor_extension_array_test.cc @@ -212,7 +212,7 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) { CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":["x","y"]})", "Invalid serialized JSON data"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":(3,4)})", - "shape must be an array, got unknown"); + "Invalid serialized JSON data"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"permutation":[1,0,2]})", "Invalid permutation"); @@ -223,7 +223,7 @@ 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 floating-point number"); + "shape must contain integers, got number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":["3","4"]})", "shape must contain integers, got string"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[null]})", @@ -248,12 +248,16 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) { CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"permutation":{"a":1}})", "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 floating-point number"); + CheckDeserializationRaises(ext_type_, storage_type, + R"({"shape":[3,4],"permutation":[1.5,0.5]})", + "permutation must contain integers, got number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"permutation":["a","b"]})", "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, @@ -276,6 +280,8 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) { CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"dim_names":[null,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,36 +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"); - CheckDeserializationRaises( - ext_type_, storage_type, R"({"permutation":[1.5,0.5,2.5]})", - "permutation must contain integers, got floating-point number"); + CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[1.5,0.5,2.5]})", + "permutation must contain integers, got number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[null,null,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"); CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":[1,2,3]})", "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"); - CheckDeserializationRaises( - ext_type_, storage_type, R"({"uniform_shape":[1.5,null,null]})", - "uniform_shape must contain integers or nulls, got floating-point number"); + CheckDeserializationRaises(ext_type_, storage_type, + R"({"uniform_shape":[1.5,null,null]})", + "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"); + CheckDeserializationRaises(ext_type_, storage_type, R"({"uniform_shape":[]})", + "Invalid uniform_shape"); } TEST_F(TestVariableShapeTensorType, RoundtripBatch) { diff --git a/cpp/src/arrow/extension/tensor_internal.cc b/cpp/src/arrow/extension/tensor_internal.cc index d2965bc578f7..8168c85622a8 100644 --- a/cpp/src/arrow/extension/tensor_internal.cc +++ b/cpp/src/arrow/extension/tensor_internal.cc @@ -30,6 +30,124 @@ namespace arrow::internal { +namespace { + +const char* JsonTypeName(simdjson::dom::element_type type) { + switch (type) { + case simdjson::dom::element_type::ARRAY: + return "array"; + case simdjson::dom::element_type::OBJECT: + return "object"; + case simdjson::dom::element_type::INT64: + case simdjson::dom::element_type::UINT64: + case simdjson::dom::element_type::DOUBLE: + case simdjson::dom::element_type::BIGINT: + return "number"; + case simdjson::dom::element_type::STRING: + return "string"; + case simdjson::dom::element_type::BOOL: + return "boolean"; + case simdjson::dom::element_type::NULL_VALUE: + return "null"; + } + return "unknown"; +} + +Result GetJsonArray(simdjson::dom::element value, + std::string_view name) { + if (!value.is_array()) { + return Status::Invalid(name, " must be an array, got ", JsonTypeName(value.type())); + } + return ResolveSimdjsonResult(value.get_array(), "Failed to get JSON array"); +} + +Result GetJsonInt(simdjson::dom::element value, std::string_view name, + std::string_view expected) { + if (!value.is_int64()) { + return Status::Invalid(name, " must contain ", expected, ", got ", + JsonTypeName(value.type())); + } + return ResolveSimdjsonResult(value.get_int64(), "Failed to get JSON integer"); +} + +} // namespace + +Result ParseJsonObject(simdjson::dom::parser& parser, + const std::string& json) { + return ResolveSimdjsonResult(parser.parse(json).get_object(), + "Invalid serialized JSON data"); +} + +Result> GetOptionalJsonField( + const simdjson::dom::object& object, std::string_view key) { + auto field = object.at_key(key); + if (field.error() == simdjson::NO_SUCH_FIELD) { + return std::nullopt; + } + + ARROW_ASSIGN_OR_RAISE( + auto value, + ResolveSimdjsonResult(std::move(field), "Failed to get JSON object field")); + + return std::optional(value); +} + +Result> GetJsonIntArray(simdjson::dom::element value, + std::string_view name) { + ARROW_ASSIGN_OR_RAISE(auto array, GetJsonArray(value, name)); + + std::vector result; + result.reserve(array.size()); + + for (auto element : array) { + ARROW_ASSIGN_OR_RAISE(auto number, GetJsonInt(element, name, "integers")); + result.push_back(number); + } + + return result; +} + +Result>> GetJsonNullableIntArray( + simdjson::dom::element value, std::string_view name) { + ARROW_ASSIGN_OR_RAISE(auto array, GetJsonArray(value, name)); + + std::vector> result; + result.reserve(array.size()); + + for (auto element : array) { + if (element.is_null()) { + result.emplace_back(std::nullopt); + } else { + ARROW_ASSIGN_OR_RAISE(auto number, GetJsonInt(element, name, "integers or nulls")); + result.emplace_back(number); + } + } + + return result; +} + +Result> GetJsonStringArray(simdjson::dom::element value, + std::string_view name) { + ARROW_ASSIGN_OR_RAISE(auto array, GetJsonArray(value, name)); + + std::vector result; + result.reserve(array.size()); + + for (auto element : array) { + if (!element.is_string()) { + return Status::Invalid(name, " must contain strings, got ", + JsonTypeName(element.type())); + } + + ARROW_ASSIGN_OR_RAISE( + auto string, + ResolveSimdjsonResult(element.get_string(), "Failed to get JSON string")); + result.emplace_back(string); + } + + return result; +} + 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 b54945ad50ab..080c1a2ca9ed 100644 --- a/cpp/src/arrow/extension/tensor_internal.h +++ b/cpp/src/arrow/extension/tensor_internal.h @@ -21,8 +21,11 @@ #include #include +#include + #include "arrow/result.h" #include "arrow/type_fwd.h" +#include "arrow/util/simdjson_internal.h" namespace arrow::internal { @@ -49,4 +52,19 @@ Result> SliceTensorBuffer(const Array& data_array, const DataType& value_type, std::span shape); +Result ParseJsonObject(simdjson::dom::parser& parser, + const std::string& json); + +Result> GetOptionalJsonField( + const simdjson::dom::object& object, std::string_view key); + +Result> GetJsonIntArray(simdjson::dom::element value, + std::string_view name); + +Result>> GetJsonNullableIntArray( + simdjson::dom::element value, std::string_view name); + +Result> GetJsonStringArray(simdjson::dom::element value, + std::string_view name); + } // namespace arrow::internal diff --git a/cpp/src/arrow/extension/variable_shape_tensor.cc b/cpp/src/arrow/extension/variable_shape_tensor.cc index 132c29193f79..784cd334b1d6 100644 --- a/cpp/src/arrow/extension/variable_shape_tensor.cc +++ b/cpp/src/arrow/extension/variable_shape_tensor.cc @@ -154,176 +154,46 @@ Result> VariableShapeTensorType::Deserialize( internal::checked_cast(*storage_type->field(1)->type()) .list_size(); - simdjson::padded_string padded_json(serialized_data); - simdjson::ondemand::parser parser; + simdjson::dom::parser parser; + ARROW_ASSIGN_OR_RAISE(auto object, internal::ParseJsonObject(parser, serialized_data)); - ARROW_ASSIGN_OR_RAISE(auto document, - internal::ResolveSimdjsonResult(parser.iterate(padded_json), - "Invalid serialized JSON data")); - - ARROW_ASSIGN_OR_RAISE(auto object, - internal::ResolveSimdjsonResult(document.get_object(), - "Invalid serialized JSON data")); + ARROW_ASSIGN_OR_RAISE(auto permutation_value, + internal::GetOptionalJsonField(object, "permutation")); std::vector permutation; - std::vector dim_names; - std::vector> uniform_shape; - - for (auto field_result : object) { - ARROW_ASSIGN_OR_RAISE(auto field, internal::ResolveSimdjsonResult( - field_result, "Invalid serialized JSON data")); - - ARROW_ASSIGN_OR_RAISE( - auto key, internal::ResolveSimdjsonResult(field.unescaped_key(), - "Failed to get JSON object key")); - - auto value = field.value(); - - if (key == "permutation") { - ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( - value.type(), "Invalid serialized JSON 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")); - - 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 ", - internal::JsonNumberTypeName(number_type)); - } - - 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") { - ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( - value.type(), "Invalid serialized JSON 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); + if (permutation_value.has_value()) { + ARROW_ASSIGN_OR_RAISE(permutation, + internal::GetJsonIntArray(*permutation_value, "permutation")); - 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") { - ARROW_ASSIGN_OR_RAISE(auto type, internal::ResolveSimdjsonResult( - value.type(), "Invalid serialized JSON 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 (permutation.size() != static_cast(ndim)) { + return Status::Invalid("Invalid permutation"); + } + RETURN_NOT_OK(internal::IsPermutationValid(permutation)); + } - 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 dim_names_value, + internal::GetOptionalJsonField(object, "dim_names")); - ARROW_ASSIGN_OR_RAISE(auto number_type, - internal::ResolveSimdjsonResult( - element.get_number_type(), - "Failed to determine uniform_shape number type")); + std::vector dim_names; + if (dim_names_value.has_value()) { + ARROW_ASSIGN_OR_RAISE(dim_names, + internal::GetJsonStringArray(*dim_names_value, "dim_names")); - if (number_type != simdjson::ondemand::number_type::signed_integer) { - return Status::Invalid("uniform_shape must contain integers or nulls, got ", - internal::JsonNumberTypeName(number_type)); - } + if (dim_names.size() != static_cast(ndim)) { + return Status::Invalid("Invalid dim_names"); + } + } - ARROW_ASSIGN_OR_RAISE( - auto number, internal::ResolveSimdjsonResult( - element.get_int64(), "Failed to get uniform_shape integer")); + ARROW_ASSIGN_OR_RAISE(auto uniform_shape_value, + internal::GetOptionalJsonField(object, "uniform_shape")); - uniform_shape.emplace_back(number); - } + std::vector> uniform_shape; + if (uniform_shape_value.has_value()) { + ARROW_ASSIGN_OR_RAISE(uniform_shape, internal::GetJsonNullableIntArray( + *uniform_shape_value, "uniform_shape")); - if (uniform_shape.size() != static_cast(ndim)) { - return Status::Invalid("Invalid: uniform_shape"); - } + if (uniform_shape.size() != static_cast(ndim)) { + return Status::Invalid("Invalid uniform_shape"); } } From 921c5e72c880bc54a5bc648ea094b9ec7bdfaa21 Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Mon, 24 Aug 2026 23:35:05 +0530 Subject: [PATCH 09/12] Migrate Helpers to simdjson_internal and restore comment --- cpp/src/arrow/extension/fixed_shape_tensor.cc | 3 + cpp/src/arrow/extension/tensor_internal.cc | 118 ----------------- cpp/src/arrow/extension/tensor_internal.h | 18 --- cpp/src/arrow/util/simdjson_internal.h | 119 ++++++++++++++++++ 4 files changed, 122 insertions(+), 136 deletions(-) diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index ac3531fd200d..6a86d6a7a66f 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -153,6 +153,9 @@ Result> FixedShapeTensorType::Deserialize( } } + // Validate product of shape dimensions matches storage type list_size. + // This check is intentionally after field parsing so that metadata-level errors + // (type mismatches, size mismatches) are reported first. ARROW_ASSIGN_OR_RAISE(auto ext_type, FixedShapeTensorType::Make( value_type, shape, permutation, dim_names)); diff --git a/cpp/src/arrow/extension/tensor_internal.cc b/cpp/src/arrow/extension/tensor_internal.cc index 8168c85622a8..d2965bc578f7 100644 --- a/cpp/src/arrow/extension/tensor_internal.cc +++ b/cpp/src/arrow/extension/tensor_internal.cc @@ -30,124 +30,6 @@ namespace arrow::internal { -namespace { - -const char* JsonTypeName(simdjson::dom::element_type type) { - switch (type) { - case simdjson::dom::element_type::ARRAY: - return "array"; - case simdjson::dom::element_type::OBJECT: - return "object"; - case simdjson::dom::element_type::INT64: - case simdjson::dom::element_type::UINT64: - case simdjson::dom::element_type::DOUBLE: - case simdjson::dom::element_type::BIGINT: - return "number"; - case simdjson::dom::element_type::STRING: - return "string"; - case simdjson::dom::element_type::BOOL: - return "boolean"; - case simdjson::dom::element_type::NULL_VALUE: - return "null"; - } - return "unknown"; -} - -Result GetJsonArray(simdjson::dom::element value, - std::string_view name) { - if (!value.is_array()) { - return Status::Invalid(name, " must be an array, got ", JsonTypeName(value.type())); - } - return ResolveSimdjsonResult(value.get_array(), "Failed to get JSON array"); -} - -Result GetJsonInt(simdjson::dom::element value, std::string_view name, - std::string_view expected) { - if (!value.is_int64()) { - return Status::Invalid(name, " must contain ", expected, ", got ", - JsonTypeName(value.type())); - } - return ResolveSimdjsonResult(value.get_int64(), "Failed to get JSON integer"); -} - -} // namespace - -Result ParseJsonObject(simdjson::dom::parser& parser, - const std::string& json) { - return ResolveSimdjsonResult(parser.parse(json).get_object(), - "Invalid serialized JSON data"); -} - -Result> GetOptionalJsonField( - const simdjson::dom::object& object, std::string_view key) { - auto field = object.at_key(key); - if (field.error() == simdjson::NO_SUCH_FIELD) { - return std::nullopt; - } - - ARROW_ASSIGN_OR_RAISE( - auto value, - ResolveSimdjsonResult(std::move(field), "Failed to get JSON object field")); - - return std::optional(value); -} - -Result> GetJsonIntArray(simdjson::dom::element value, - std::string_view name) { - ARROW_ASSIGN_OR_RAISE(auto array, GetJsonArray(value, name)); - - std::vector result; - result.reserve(array.size()); - - for (auto element : array) { - ARROW_ASSIGN_OR_RAISE(auto number, GetJsonInt(element, name, "integers")); - result.push_back(number); - } - - return result; -} - -Result>> GetJsonNullableIntArray( - simdjson::dom::element value, std::string_view name) { - ARROW_ASSIGN_OR_RAISE(auto array, GetJsonArray(value, name)); - - std::vector> result; - result.reserve(array.size()); - - for (auto element : array) { - if (element.is_null()) { - result.emplace_back(std::nullopt); - } else { - ARROW_ASSIGN_OR_RAISE(auto number, GetJsonInt(element, name, "integers or nulls")); - result.emplace_back(number); - } - } - - return result; -} - -Result> GetJsonStringArray(simdjson::dom::element value, - std::string_view name) { - ARROW_ASSIGN_OR_RAISE(auto array, GetJsonArray(value, name)); - - std::vector result; - result.reserve(array.size()); - - for (auto element : array) { - if (!element.is_string()) { - return Status::Invalid(name, " must contain strings, got ", - JsonTypeName(element.type())); - } - - ARROW_ASSIGN_OR_RAISE( - auto string, - ResolveSimdjsonResult(element.get_string(), "Failed to get JSON string")); - result.emplace_back(string); - } - - return result; -} - 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 080c1a2ca9ed..b54945ad50ab 100644 --- a/cpp/src/arrow/extension/tensor_internal.h +++ b/cpp/src/arrow/extension/tensor_internal.h @@ -21,11 +21,8 @@ #include #include -#include - #include "arrow/result.h" #include "arrow/type_fwd.h" -#include "arrow/util/simdjson_internal.h" namespace arrow::internal { @@ -52,19 +49,4 @@ Result> SliceTensorBuffer(const Array& data_array, const DataType& value_type, std::span shape); -Result ParseJsonObject(simdjson::dom::parser& parser, - const std::string& json); - -Result> GetOptionalJsonField( - const simdjson::dom::object& object, std::string_view key); - -Result> GetJsonIntArray(simdjson::dom::element value, - std::string_view name); - -Result>> GetJsonNullableIntArray( - simdjson::dom::element value, std::string_view name); - -Result> GetJsonStringArray(simdjson::dom::element value, - std::string_view name); - } // namespace arrow::internal diff --git a/cpp/src/arrow/util/simdjson_internal.h b/cpp/src/arrow/util/simdjson_internal.h index f31dd1e35d0d..f76834dd47e4 100644 --- a/cpp/src/arrow/util/simdjson_internal.h +++ b/cpp/src/arrow/util/simdjson_internal.h @@ -19,8 +19,11 @@ #include #include +#include +#include #include #include +#include #include @@ -91,6 +94,122 @@ Result ResolveSimdjsonResult(simdjson::simdjson_result result, return value; } +inline const char* JsonTypeName(simdjson::dom::element_type type) { + switch (type) { + case simdjson::dom::element_type::ARRAY: + return "array"; + case simdjson::dom::element_type::OBJECT: + return "object"; + case simdjson::dom::element_type::INT64: + case simdjson::dom::element_type::UINT64: + case simdjson::dom::element_type::DOUBLE: + case simdjson::dom::element_type::BIGINT: + return "number"; + case simdjson::dom::element_type::STRING: + return "string"; + case simdjson::dom::element_type::BOOL: + return "boolean"; + case simdjson::dom::element_type::NULL_VALUE: + return "null"; + } + return "unknown"; +} + +inline Result GetJsonArray(simdjson::dom::element value, + std::string_view name) { + if (!value.is_array()) { + return Status::Invalid(name, " must be an array, got ", JsonTypeName(value.type())); + } + return ResolveSimdjsonResult(value.get_array(), "Failed to get JSON array"); +} + +inline Result GetJsonInt(simdjson::dom::element value, std::string_view name, + std::string_view expected) { + if (!value.is_int64()) { + return Status::Invalid(name, " must contain ", expected, ", got ", + JsonTypeName(value.type())); + } + return ResolveSimdjsonResult(value.get_int64(), "Failed to get JSON integer"); +} + +inline Result ParseJsonObject(simdjson::dom::parser& parser, + const std::string& json) { + return ResolveSimdjsonResult(parser.parse(json).get_object(), + "Invalid serialized JSON data"); +} + +// object.at_key() performs a linear search. This is acceptable here +// since these objects are expected to contain only a small number of fields. +inline Result> GetOptionalJsonField( + const simdjson::dom::object& object, std::string_view key) { + auto field = object.at_key(key); + if (field.error() == simdjson::NO_SUCH_FIELD) { + return std::nullopt; + } + + ARROW_ASSIGN_OR_RAISE( + auto value, + ResolveSimdjsonResult(std::move(field), "Failed to get JSON object field")); + + return std::optional(value); +} + +inline Result> GetJsonIntArray(simdjson::dom::element value, + std::string_view name) { + ARROW_ASSIGN_OR_RAISE(auto array, GetJsonArray(value, name)); + + std::vector result; + result.reserve(array.size()); + + for (auto element : array) { + ARROW_ASSIGN_OR_RAISE(auto number, GetJsonInt(element, name, "integers")); + result.push_back(number); + } + + return result; +} + +inline Result>> GetJsonNullableIntArray( + simdjson::dom::element value, std::string_view name) { + ARROW_ASSIGN_OR_RAISE(auto array, GetJsonArray(value, name)); + + std::vector> result; + result.reserve(array.size()); + + for (auto element : array) { + if (element.is_null()) { + result.emplace_back(std::nullopt); + } else { + ARROW_ASSIGN_OR_RAISE(auto number, GetJsonInt(element, name, "integers or nulls")); + result.emplace_back(number); + } + } + + return result; +} + +inline Result> GetJsonStringArray(simdjson::dom::element value, + std::string_view name) { + ARROW_ASSIGN_OR_RAISE(auto array, GetJsonArray(value, name)); + + std::vector result; + result.reserve(array.size()); + + for (auto element : array) { + if (!element.is_string()) { + return Status::Invalid(name, " must contain strings, got ", + JsonTypeName(element.type())); + } + + ARROW_ASSIGN_OR_RAISE( + auto string, + ResolveSimdjsonResult(element.get_string(), "Failed to get JSON string")); + result.emplace_back(string); + } + + return result; +} + template From 9f4cf1a89cb88801a48a4b61ef9714724e4d9e57 Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Mon, 24 Aug 2026 23:44:21 +0530 Subject: [PATCH 10/12] Remove unused helper --- cpp/src/arrow/util/simdjson_internal.h | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/cpp/src/arrow/util/simdjson_internal.h b/cpp/src/arrow/util/simdjson_internal.h index f76834dd47e4..94c0866685ec 100644 --- a/cpp/src/arrow/util/simdjson_internal.h +++ b/cpp/src/arrow/util/simdjson_internal.h @@ -447,19 +447,5 @@ inline Status ValidateJsonDocument(simdjson::ondemand::parser& parser, return ConsumeJsonValue(value); } -inline const char* JsonNumberTypeName(simdjson::ondemand::number_type type) { - switch (type) { - case simdjson::ondemand::number_type::signed_integer: - return "signed integer"; - case simdjson::ondemand::number_type::unsigned_integer: - return "unsigned integer"; - case simdjson::ondemand::number_type::floating_point_number: - return "floating-point number"; - case simdjson::ondemand::number_type::big_integer: - return "big integer"; - } - return "unknown"; -} - } // namespace internal } // namespace arrow From 661471652d8b6b007f4ebb80021e2407d6678ffc Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan <156181482+Reranko05@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:53:03 +0530 Subject: [PATCH 11/12] Update cpp/src/arrow/util/simdjson_internal.h Co-authored-by: Antoine Pitrou --- cpp/src/arrow/util/simdjson_internal.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/util/simdjson_internal.h b/cpp/src/arrow/util/simdjson_internal.h index 94c0866685ec..2f21b1a2431b 100644 --- a/cpp/src/arrow/util/simdjson_internal.h +++ b/cpp/src/arrow/util/simdjson_internal.h @@ -151,7 +151,7 @@ inline Result> GetOptionalJsonField( auto value, ResolveSimdjsonResult(std::move(field), "Failed to get JSON object field")); - return std::optional(value); + return std::optional(std::move(value)); } inline Result> GetJsonIntArray(simdjson::dom::element value, From 2a290504c3614ac4e587699c679f4197df563559 Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Tue, 25 Aug 2026 18:31:59 +0530 Subject: [PATCH 12/12] fallback BIGINT to return unknown --- cpp/src/arrow/util/simdjson_internal.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/util/simdjson_internal.h b/cpp/src/arrow/util/simdjson_internal.h index 2f21b1a2431b..799ebfc5ea62 100644 --- a/cpp/src/arrow/util/simdjson_internal.h +++ b/cpp/src/arrow/util/simdjson_internal.h @@ -103,7 +103,6 @@ inline const char* JsonTypeName(simdjson::dom::element_type type) { case simdjson::dom::element_type::INT64: case simdjson::dom::element_type::UINT64: case simdjson::dom::element_type::DOUBLE: - case simdjson::dom::element_type::BIGINT: return "number"; case simdjson::dom::element_type::STRING: return "string"; @@ -111,8 +110,9 @@ inline const char* JsonTypeName(simdjson::dom::element_type type) { return "boolean"; case simdjson::dom::element_type::NULL_VALUE: return "null"; + default: + return "unknown"; } - return "unknown"; } inline Result GetJsonArray(simdjson::dom::element value,