diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index cd3d783479d6..6a86d6a7a66f 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,57 +115,39 @@ 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()) { - 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::dom::parser parser; + ARROW_ASSIGN_OR_RAISE(auto object, internal::ParseJsonObject(parser, serialized_data)); + + ARROW_ASSIGN_OR_RAISE(auto shape_value, + internal::ResolveSimdjsonResult(object.at_key("shape"), + "Invalid serialized JSON data")); + ARROW_ASSIGN_OR_RAISE(auto shape, internal::GetJsonIntArray(shape_value, "shape")); + ARROW_ASSIGN_OR_RAISE(auto permutation_value, + internal::GetOptionalJsonField(object, "permutation")); std::vector permutation; - if (document.HasMember("permutation")) { - const auto& json_permutation = document["permutation"]; - if (!json_permutation.IsArray()) { - return Status::Invalid("permutation must be an array, got ", - internal::JsonTypeName(json_permutation)); - } - for (const auto& x : json_permutation.GetArray()) { - if (!x.IsInt64()) { - return Status::Invalid("permutation must contain integers, got ", - internal::JsonTypeName(x)); - } - permutation.emplace_back(x.GetInt64()); - } + if (permutation_value.has_value()) { + ARROW_ASSIGN_OR_RAISE(permutation, + internal::GetJsonIntArray(*permutation_value, "permutation")); + if (shape.size() != permutation.size()) { return Status::Invalid("Invalid permutation"); } RETURN_NOT_OK(internal::IsPermutationValid(permutation)); } + + ARROW_ASSIGN_OR_RAISE(auto dim_names_value, + internal::GetOptionalJsonField(object, "dim_names")); + std::vector dim_names; - if (document.HasMember("dim_names")) { - const auto& json_dim_names = document["dim_names"]; - if (!json_dim_names.IsArray()) { - return Status::Invalid("dim_names must be an array, got ", - internal::JsonTypeName(json_dim_names)); - } - for (const auto& x : json_dim_names.GetArray()) { - if (!x.IsString()) { - return Status::Invalid("dim_names must contain strings, got ", - internal::JsonTypeName(x)); - } - dim_names.emplace_back(x.GetString()); - } + if (dim_names_value.has_value()) { + ARROW_ASSIGN_OR_RAISE(dim_names, + internal::GetJsonStringArray(*dim_names_value, "dim_names")); + if (shape.size() != dim_names.size()) { return Status::Invalid("Invalid dim_names"); } @@ -177,14 +158,18 @@ Result> FixedShapeTensorType::Deserialize( // (type mismatches, size mismatches) are reported first. ARROW_ASSIGN_OR_RAISE(auto ext_type, FixedShapeTensorType::Make( value_type, shape, permutation, dim_names)); + const auto& fst_type = internal::checked_cast(*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..29147df6ba6a 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,20 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) { // Validate permutation member must be an array with integer values CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"permutation":"invalid"})", - "permutation must be an array, got String"); + "permutation must be an array, got string"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"permutation":{"a":1}})", - "permutation must be an array, got Object"); + "permutation must be an array, got object"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"permutation":[1.5,0.5]})", - "permutation must contain integers, got Number"); + "permutation must contain integers, got number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"permutation":["a","b"]})", - "permutation must contain integers, got String"); + "permutation must contain integers, got string"); + // Validate permutation member must be an array with integer values + CheckDeserializationRaises(ext_type_, storage_type, + R"({"shape":[3,4],"permutation":[]})", + "Invalid permutation"); // Validate permutation values must be unique integers in [0, N-1] CheckDeserializationRaises(ext_type_, storage_type, @@ -269,13 +273,15 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) { // Validate dim_names member must be an array with string values CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"dim_names":"invalid"})", - "dim_names must be an array, got String"); + "dim_names must be an array, got string"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"dim_names":[1,2]})", - "dim_names must contain strings, got Number"); + "dim_names must contain strings, got number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"dim_names":[null,null]})", - "dim_names must contain strings, got Null"); + "dim_names must contain strings, got null"); + CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3,4],"dim_names":[]})", + "Invalid dim_names"); } TEST_F(TestFixedShapeTensorType, MakeValidatesShape) { @@ -858,35 +864,41 @@ TEST_F(TestVariableShapeTensorType, MetadataSerializationRoundtrip) { CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":(3,4)})", "Invalid serialized JSON data"); CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[1,0]})", - "Invalid: permutation"); + "Invalid permutation"); CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":["x","y"]})", - "Invalid: dim_names"); + "Invalid dim_names"); // Validate permutation member must be an array with integer values. Error // message should include the JSON type name of the offending value. CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":"invalid"})", - "permutation must be an array, got String"); + "permutation must be an array, got string"); CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[1.5,0.5,2.5]})", - "permutation must contain integers, got Number"); + "permutation must contain integers, got number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[null,null,null]})", - "permutation must contain integers, got Null"); + "permutation must contain integers, got null"); + CheckDeserializationRaises(ext_type_, storage_type, R"({"permutation":[]})", + "Invalid permutation"); // Validate dim_names member must be an array with string values CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":"invalid"})", - "dim_names must be an array, got String"); + "dim_names must be an array, got string"); CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":[1,2,3]})", - "dim_names must contain strings, got Number"); + "dim_names must contain strings, got number"); + CheckDeserializationRaises(ext_type_, storage_type, R"({"dim_names":[]})", + "Invalid dim_names"); // Validate uniform_shape member must be an array with integer-or-null values CheckDeserializationRaises(ext_type_, storage_type, R"({"uniform_shape":"invalid"})", - "uniform_shape must be an array, got String"); + "uniform_shape must be an array, got string"); CheckDeserializationRaises(ext_type_, storage_type, R"({"uniform_shape":[1.5,null,null]})", - "uniform_shape must contain integers or nulls, got Number"); + "uniform_shape must contain integers or nulls, got number"); CheckDeserializationRaises(ext_type_, storage_type, R"({"uniform_shape":["x",null,null]})", - "uniform_shape must contain integers or nulls, got String"); + "uniform_shape must contain integers or nulls, got string"); + CheckDeserializationRaises(ext_type_, storage_type, R"({"uniform_shape":[]})", + "Invalid uniform_shape"); } TEST_F(TestVariableShapeTensorType, RoundtripBatch) { 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. diff --git a/cpp/src/arrow/extension/variable_shape_tensor.cc b/cpp/src/arrow/extension/variable_shape_tensor.cc index 40171f909a9d..784cd334b1d6 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,63 +154,46 @@ 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()) { - return Status::Invalid("Invalid serialized JSON data: ", serialized_data); - } + simdjson::dom::parser parser; + ARROW_ASSIGN_OR_RAISE(auto object, internal::ParseJsonObject(parser, serialized_data)); + + ARROW_ASSIGN_OR_RAISE(auto permutation_value, + internal::GetOptionalJsonField(object, "permutation")); 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()); + if (permutation_value.has_value()) { + ARROW_ASSIGN_OR_RAISE(permutation, + internal::GetJsonIntArray(*permutation_value, "permutation")); + + if (permutation.size() != static_cast(ndim)) { + return Status::Invalid("Invalid permutation"); } RETURN_NOT_OK(internal::IsPermutationValid(permutation)); } + + ARROW_ASSIGN_OR_RAISE(auto dim_names_value, + internal::GetOptionalJsonField(object, "dim_names")); + std::vector 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)); - } - 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)); - } - dim_names.emplace_back(x.GetString()); + if (dim_names_value.has_value()) { + ARROW_ASSIGN_OR_RAISE(dim_names, + internal::GetJsonStringArray(*dim_names_value, "dim_names")); + + if (dim_names.size() != static_cast(ndim)) { + return Status::Invalid("Invalid dim_names"); } } + ARROW_ASSIGN_OR_RAISE(auto uniform_shape_value, + internal::GetOptionalJsonField(object, "uniform_shape")); + 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 (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"); } } diff --git a/cpp/src/arrow/util/simdjson_internal.h b/cpp/src/arrow/util/simdjson_internal.h index 1badd99938b1..799ebfc5ea62 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: + 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"; + default: + 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(std::move(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