Skip to content

Commit 9c5ff53

Browse files
Reranko05Reranko05
committed
Migrate FixedShapeTensorType deserialization to simdjson
1 parent c44eca6 commit 9c5ff53

2 files changed

Lines changed: 185 additions & 55 deletions

File tree

‎cpp/src/arrow/extension/fixed_shape_tensor.cc‎

Lines changed: 173 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -19,23 +19,22 @@
1919
#include <numeric>
2020
#include <sstream>
2121

22+
#include <simdjson.h>
23+
2224
#include "arrow/extension/fixed_shape_tensor.h"
2325
#include "arrow/extension/tensor_internal.h"
2426
#include "arrow/scalar.h"
2527

2628
#include "arrow/array/array_nested.h"
2729
#include "arrow/array/array_primitive.h"
2830
#include "arrow/json/json_writer_internal.h"
29-
#include "arrow/json/rapidjson_defs.h" // IWYU pragma: keep
3031
#include "arrow/tensor.h"
3132
#include "arrow/util/logging_internal.h"
3233
#include "arrow/util/print_internal.h"
34+
#include "arrow/util/simdjson_internal.h"
3335
#include "arrow/util/sort_internal.h"
3436
#include "arrow/util/string.h"
3537

36-
#include <rapidjson/document.h>
37-
38-
namespace rj = arrow::rapidjson;
3938
using ::arrow::json::JsonWriter;
4039

4140
namespace arrow::extension {
@@ -116,60 +115,189 @@ Result<std::shared_ptr<DataType>> FixedShapeTensorType::Deserialize(
116115
return Status::Invalid("Expected FixedSizeList storage type, got ",
117116
storage_type->ToString());
118117
}
118+
119119
auto fsl_type = internal::checked_pointer_cast<FixedSizeListType>(storage_type);
120120
auto value_type = fsl_type->value_type();
121-
rj::Document document;
122-
if (document.Parse(serialized_data.data(), serialized_data.length()).HasParseError() ||
123-
!document.IsObject() || !document.HasMember("shape") ||
124-
!document["shape"].IsArray()) {
121+
122+
simdjson::padded_string padded_json(serialized_data);
123+
simdjson::ondemand::parser parser;
124+
simdjson::ondemand::document document;
125+
126+
if (auto error = parser.iterate(padded_json).get(document);
127+
error != simdjson::SUCCESS) {
125128
return Status::Invalid("Invalid serialized JSON data: ", serialized_data);
126129
}
127130

128-
std::vector<int64_t> shape;
129-
for (const auto& x : document["shape"].GetArray()) {
130-
if (!x.IsInt64()) {
131-
return Status::Invalid("shape must contain integers, got ",
132-
internal::JsonTypeName(x));
133-
}
134-
shape.emplace_back(x.GetInt64());
131+
simdjson::ondemand::object object;
132+
if (auto error = document.get_object().get(object); error != simdjson::SUCCESS) {
133+
return Status::Invalid("Invalid serialized JSON data: ", serialized_data);
135134
}
136135

136+
std::vector<int64_t> shape;
137137
std::vector<int64_t> permutation;
138-
if (document.HasMember("permutation")) {
139-
const auto& json_permutation = document["permutation"];
140-
if (!json_permutation.IsArray()) {
141-
return Status::Invalid("permutation must be an array, got ",
142-
internal::JsonTypeName(json_permutation));
143-
}
144-
for (const auto& x : json_permutation.GetArray()) {
145-
if (!x.IsInt64()) {
146-
return Status::Invalid("permutation must contain integers, got ",
147-
internal::JsonTypeName(x));
138+
std::vector<std::string> dim_names;
139+
140+
bool has_shape = false;
141+
142+
for (auto field_result : object) {
143+
ARROW_ASSIGN_OR_RAISE(auto field, internal::ResolveSimdjsonResult(
144+
field_result, "Failed to iterate JSON object"));
145+
146+
ARROW_ASSIGN_OR_RAISE(
147+
auto key, internal::ResolveSimdjsonResult(field.unescaped_key(),
148+
"Failed to get JSON object key"));
149+
150+
auto value = field.value();
151+
152+
if (key == "shape") {
153+
has_shape = true;
154+
155+
simdjson::ondemand::json_type type;
156+
if (auto error = value.type().get(type);
157+
error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) {
158+
return Status::Invalid("Invalid serialized JSON data: ", serialized_data);
159+
}
160+
161+
if (type != simdjson::ondemand::json_type::array) {
162+
return Status::Invalid("shape must be an array, got ",
163+
internal::JsonTypeName(type));
164+
}
165+
166+
ARROW_ASSIGN_OR_RAISE(auto array,
167+
internal::ResolveSimdjsonResult(value.get_array(),
168+
"Failed to get shape array"));
169+
170+
for (auto element_result : array) {
171+
ARROW_ASSIGN_OR_RAISE(auto element,
172+
internal::ResolveSimdjsonResult(
173+
element_result, "Failed to iterate shape array"));
174+
175+
ARROW_ASSIGN_OR_RAISE(
176+
auto element_type,
177+
internal::ResolveSimdjsonResult(
178+
element.type(), "Failed to determine shape element JSON type"));
179+
180+
if (element_type != simdjson::ondemand::json_type::number) {
181+
return Status::Invalid("shape must contain integers, got ",
182+
internal::JsonTypeName(element_type));
183+
}
184+
185+
ARROW_ASSIGN_OR_RAISE(
186+
auto number_type,
187+
internal::ResolveSimdjsonResult(element.get_number_type(),
188+
"Failed to determine shape number type"));
189+
190+
if (number_type != simdjson::ondemand::number_type::signed_integer) {
191+
return Status::Invalid("shape must contain integers, got number");
192+
}
193+
194+
ARROW_ASSIGN_OR_RAISE(
195+
auto number, internal::ResolveSimdjsonResult(element.get_int64(),
196+
"Failed to get shape integer"));
197+
198+
shape.emplace_back(number);
199+
}
200+
201+
} else if (key == "permutation") {
202+
simdjson::ondemand::json_type type;
203+
if (auto error = value.type().get(type);
204+
error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) {
205+
return Status::Invalid("Invalid serialized JSON data: ", serialized_data);
206+
}
207+
208+
if (type != simdjson::ondemand::json_type::array) {
209+
return Status::Invalid("permutation must be an array, got ",
210+
internal::JsonTypeName(type));
211+
}
212+
213+
ARROW_ASSIGN_OR_RAISE(
214+
auto array, internal::ResolveSimdjsonResult(value.get_array(),
215+
"Failed to get permutation array"));
216+
217+
for (auto element_result : array) {
218+
ARROW_ASSIGN_OR_RAISE(auto element,
219+
internal::ResolveSimdjsonResult(
220+
element_result, "Failed to iterate permutation array"));
221+
222+
ARROW_ASSIGN_OR_RAISE(
223+
auto element_type,
224+
internal::ResolveSimdjsonResult(
225+
element.type(), "Failed to determine permutation element JSON type"));
226+
227+
if (element_type != simdjson::ondemand::json_type::number) {
228+
return Status::Invalid("permutation must contain integers, got ",
229+
internal::JsonTypeName(element_type));
230+
}
231+
232+
ARROW_ASSIGN_OR_RAISE(auto number_type,
233+
internal::ResolveSimdjsonResult(
234+
element.get_number_type(),
235+
"Failed to determine permutation number type"));
236+
237+
if (number_type != simdjson::ondemand::number_type::signed_integer) {
238+
return Status::Invalid("permutation must contain integers, got number");
239+
}
240+
241+
ARROW_ASSIGN_OR_RAISE(
242+
auto number, internal::ResolveSimdjsonResult(
243+
element.get_int64(), "Failed to get permutation integer"));
244+
245+
permutation.emplace_back(number);
246+
}
247+
248+
} else if (key == "dim_names") {
249+
simdjson::ondemand::json_type type;
250+
if (auto error = value.type().get(type);
251+
error != simdjson::SUCCESS || type == simdjson::ondemand::json_type::unknown) {
252+
return Status::Invalid("Invalid serialized JSON data: ", serialized_data);
253+
}
254+
255+
if (type != simdjson::ondemand::json_type::array) {
256+
return Status::Invalid("dim_names must be an array, got ",
257+
internal::JsonTypeName(type));
258+
}
259+
260+
ARROW_ASSIGN_OR_RAISE(
261+
auto array, internal::ResolveSimdjsonResult(value.get_array(),
262+
"Failed to get dim_names array"));
263+
264+
for (auto element_result : array) {
265+
ARROW_ASSIGN_OR_RAISE(auto element,
266+
internal::ResolveSimdjsonResult(
267+
element_result, "Failed to iterate dim_names array"));
268+
269+
ARROW_ASSIGN_OR_RAISE(
270+
auto element_type,
271+
internal::ResolveSimdjsonResult(
272+
element.type(), "Failed to determine dim_names element JSON type"));
273+
274+
if (element_type != simdjson::ondemand::json_type::string) {
275+
return Status::Invalid("dim_names must contain strings, got ",
276+
internal::JsonTypeName(element_type));
277+
}
278+
279+
ARROW_ASSIGN_OR_RAISE(auto name,
280+
internal::ResolveSimdjsonResult(element.get_string(),
281+
"Failed to get dim_name"));
282+
283+
dim_names.emplace_back(name);
148284
}
149-
permutation.emplace_back(x.GetInt64());
150285
}
286+
}
287+
288+
if (!has_shape) {
289+
return Status::Invalid("Invalid serialized JSON data: ", serialized_data);
290+
}
291+
292+
if (!permutation.empty()) {
151293
if (shape.size() != permutation.size()) {
152294
return Status::Invalid("Invalid permutation");
153295
}
154296
RETURN_NOT_OK(internal::IsPermutationValid(permutation));
155297
}
156-
std::vector<std::string> dim_names;
157-
if (document.HasMember("dim_names")) {
158-
const auto& json_dim_names = document["dim_names"];
159-
if (!json_dim_names.IsArray()) {
160-
return Status::Invalid("dim_names must be an array, got ",
161-
internal::JsonTypeName(json_dim_names));
162-
}
163-
for (const auto& x : json_dim_names.GetArray()) {
164-
if (!x.IsString()) {
165-
return Status::Invalid("dim_names must contain strings, got ",
166-
internal::JsonTypeName(x));
167-
}
168-
dim_names.emplace_back(x.GetString());
169-
}
170-
if (shape.size() != dim_names.size()) {
171-
return Status::Invalid("Invalid dim_names");
172-
}
298+
299+
if (!dim_names.empty() && shape.size() != dim_names.size()) {
300+
return Status::Invalid("Invalid dim_names");
173301
}
174302

175303
// Validate product of shape dimensions matches storage type list_size.
@@ -180,11 +308,13 @@ Result<std::shared_ptr<DataType>> FixedShapeTensorType::Deserialize(
180308
const auto& fst_type = internal::checked_cast<const FixedShapeTensorType&>(*ext_type);
181309
ARROW_ASSIGN_OR_RAISE(const int64_t expected_size,
182310
internal::ComputeShapeProduct(fst_type.shape()));
311+
183312
if (expected_size != fsl_type->list_size()) {
184313
return Status::Invalid("Product of shape dimensions (", expected_size,
185314
") does not match FixedSizeList size (", fsl_type->list_size(),
186315
")");
187316
}
317+
188318
return ext_type;
189319
}
190320

‎cpp/src/arrow/extension/tensor_extension_array_test.cc‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -223,15 +223,15 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) {
223223
// Validate shape values must be integers. Error message should include the
224224
// JSON type name of the offending value.
225225
CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[3.5,4]})",
226-
"shape must contain integers, got Number");
226+
"shape must contain integers, got number");
227227
CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":["3","4"]})",
228-
"shape must contain integers, got String");
228+
"shape must contain integers, got string");
229229
CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[null]})",
230-
"shape must contain integers, got Null");
230+
"shape must contain integers, got null");
231231
CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[true]})",
232-
"shape must contain integers, got True");
232+
"shape must contain integers, got boolean");
233233
CheckDeserializationRaises(ext_type_, storage_type, R"({"shape":[false]})",
234-
"shape must contain integers, got False");
234+
"shape must contain integers, got boolean");
235235

236236
// Validate shape values must be non-negative
237237
CheckDeserializationRaises(ext_type_, fixed_size_list(int64(), 1), R"({"shape":[-1]})",
@@ -244,16 +244,16 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) {
244244
// Validate permutation member must be an array with integer values
245245
CheckDeserializationRaises(ext_type_, storage_type,
246246
R"({"shape":[3,4],"permutation":"invalid"})",
247-
"permutation must be an array, got String");
247+
"permutation must be an array, got string");
248248
CheckDeserializationRaises(ext_type_, storage_type,
249249
R"({"shape":[3,4],"permutation":{"a":1}})",
250-
"permutation must be an array, got Object");
250+
"permutation must be an array, got object");
251251
CheckDeserializationRaises(ext_type_, storage_type,
252252
R"({"shape":[3,4],"permutation":[1.5,0.5]})",
253-
"permutation must contain integers, got Number");
253+
"permutation must contain integers, got number");
254254
CheckDeserializationRaises(ext_type_, storage_type,
255255
R"({"shape":[3,4],"permutation":["a","b"]})",
256-
"permutation must contain integers, got String");
256+
"permutation must contain integers, got string");
257257

258258
// Validate permutation values must be unique integers in [0, N-1]
259259
CheckDeserializationRaises(ext_type_, storage_type,
@@ -269,13 +269,13 @@ TEST_F(TestFixedShapeTensorType, MetadataSerializationRoundtrip) {
269269
// Validate dim_names member must be an array with string values
270270
CheckDeserializationRaises(ext_type_, storage_type,
271271
R"({"shape":[3,4],"dim_names":"invalid"})",
272-
"dim_names must be an array, got String");
272+
"dim_names must be an array, got string");
273273
CheckDeserializationRaises(ext_type_, storage_type,
274274
R"({"shape":[3,4],"dim_names":[1,2]})",
275-
"dim_names must contain strings, got Number");
275+
"dim_names must contain strings, got number");
276276
CheckDeserializationRaises(ext_type_, storage_type,
277277
R"({"shape":[3,4],"dim_names":[null,null]})",
278-
"dim_names must contain strings, got Null");
278+
"dim_names must contain strings, got null");
279279
}
280280

281281
TEST_F(TestFixedShapeTensorType, MakeValidatesShape) {

0 commit comments

Comments
 (0)