Skip to content

[WIP] GH-50873: [C++] Migrate FixedShapeTensorType deserialization to simdjson #50874

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Reranko05 wants to merge 1 commit into
base: main
Choose a base branch
from
+185 −55
Draft

[WIP] GH-50873: [C++] Migrate FixedShapeTensorType deserialization to simdjson #50874

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
216 changes: 173 additions & 43 deletions cpp/src/arrow/extension/fixed_shape_tensor.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,23 +19,22 @@
#include <numeric>
#include <sstream>

#include <simdjson.h>

#include "arrow/extension/fixed_shape_tensor.h"
#include "arrow/extension/tensor_internal.h"
#include "arrow/scalar.h"

#include "arrow/array/array_nested.h"
#include "arrow/array/array_primitive.h"
#include "arrow/json/json_writer_internal.h"
#include "arrow/json/rapidjson_defs.h" // IWYU pragma: keep
#include "arrow/tensor.h"
#include "arrow/util/logging_internal.h"
#include "arrow/util/print_internal.h"
#include "arrow/util/simdjson_internal.h"
#include "arrow/util/sort_internal.h"
#include "arrow/util/string.h"

#include <rapidjson/document.h>

namespace rj = arrow::rapidjson;
using ::arrow::json::JsonWriter;

namespace arrow::extension {
Expand Down Expand Up @@ -116,60 +115,189 @@ Result<std::shared_ptr<DataType>> FixedShapeTensorType::Deserialize(
return Status::Invalid("Expected FixedSizeList storage type, got ",
storage_type->ToString());
}

auto fsl_type = internal::checked_pointer_cast<FixedSizeListType>(storage_type);
auto value_type = fsl_type->value_type();
rj::Document document;
if (document.Parse(serialized_data.data(), serialized_data.length()).HasParseError() ||
!document.IsObject() || !document.HasMember("shape") ||
!document["shape"].IsArray()) {

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<int64_t> shape;
for (const auto& x : document["shape"].GetArray()) {
if (!x.IsInt64()) {
return Status::Invalid("shape must contain integers, got ",
internal::JsonTypeName(x));
}
shape.emplace_back(x.GetInt64());
simdjson::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<int64_t> shape;
std::vector<int64_t> permutation;
if (document.HasMember("permutation")) {
const auto& json_permutation = document["permutation"];
if (!json_permutation.IsArray()) {
return Status::Invalid("permutation must be an array, got ",
internal::JsonTypeName(json_permutation));
}
for (const auto& x : json_permutation.GetArray()) {
if (!x.IsInt64()) {
return Status::Invalid("permutation must contain integers, got ",
internal::JsonTypeName(x));
std::vector<std::string> 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<std::string> dim_names;
if (document.HasMember("dim_names")) {
const auto& json_dim_names = document["dim_names"];
if (!json_dim_names.IsArray()) {
return Status::Invalid("dim_names must be an array, got ",
internal::JsonTypeName(json_dim_names));
}
for (const auto& x : json_dim_names.GetArray()) {
if (!x.IsString()) {
return Status::Invalid("dim_names must contain strings, got ",
internal::JsonTypeName(x));
}
dim_names.emplace_back(x.GetString());
}
if (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.
Expand All @@ -180,11 +308,13 @@ Result<std::shared_ptr<DataType>> FixedShapeTensorType::Deserialize(
const auto& fst_type = internal::checked_cast<const FixedShapeTensorType&>(*ext_type);
ARROW_ASSIGN_OR_RAISE(const int64_t expected_size,
internal::ComputeShapeProduct(fst_type.shape()));

if (expected_size != fsl_type->list_size()) {
return Status::Invalid("Product of shape dimensions (", expected_size,
") does not match FixedSizeList size (", fsl_type->list_size(),
")");
}

return ext_type;
}

Expand Down
24 changes: 12 additions & 12 deletions cpp/src/arrow/extension/tensor_extension_array_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]})",
Expand All @@ -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,
Expand All @@ -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) {
Expand Down
Loading