From b220eab6f7bd11799242271263d6c68c5532812f Mon Sep 17 00:00:00 2001 From: Aaditya Srinivasan Date: Sun, 26 Jul 2026 19:35:29 +0530 Subject: [PATCH] Migrate extension type serialization to JsonWriter --- cpp/src/arrow/extension/fixed_shape_tensor.cc | 39 +++++++++-------- cpp/src/arrow/extension/opaque.cc | 27 ++++++------ .../arrow/extension/variable_shape_tensor.cc | 43 ++++++++++--------- 3 files changed, 57 insertions(+), 52 deletions(-) diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index 544616988746..697bab92b339 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -25,6 +25,7 @@ #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" @@ -33,9 +34,9 @@ #include "arrow/util/string.h" #include -#include namespace rj = arrow::rapidjson; +using ::arrow::json::JsonWriter; namespace arrow::extension { @@ -72,36 +73,38 @@ std::string FixedShapeTensorType::ToString(bool show_metadata) const { } std::string FixedShapeTensorType::Serialize() const { - rj::Document document; - document.SetObject(); - rj::Document::AllocatorType& allocator = document.GetAllocator(); + JsonWriter writer; + + writer.StartObject(); - rj::Value shape(rj::kArrayType); + writer.Key("shape"); + writer.StartArray(); for (auto v : shape_) { - shape.PushBack(v, allocator); + writer.Int64(v); } - document.AddMember(rj::Value("shape", allocator), shape, allocator); + writer.EndArray(); if (!permutation_.empty()) { - rj::Value permutation(rj::kArrayType); + writer.Key("permutation"); + writer.StartArray(); for (auto v : permutation_) { - permutation.PushBack(v, allocator); + writer.Int64(v); } - document.AddMember(rj::Value("permutation", allocator), permutation, allocator); + writer.EndArray(); } if (!dim_names_.empty()) { - rj::Value dim_names(rj::kArrayType); - for (const std::string& v : dim_names_) { - dim_names.PushBack(rj::Value{}.SetString(v.c_str(), allocator), allocator); + writer.Key("dim_names"); + writer.StartArray(); + for (const auto& v : dim_names_) { + writer.String(v); } - document.AddMember(rj::Value("dim_names", allocator), dim_names, allocator); + writer.EndArray(); } - rj::StringBuffer buffer; - rj::Writer writer(buffer); - document.Accept(writer); - return buffer.GetString(); + writer.EndObject(); + + return std::string(writer.GetString()); } Result> FixedShapeTensorType::Deserialize( diff --git a/cpp/src/arrow/extension/opaque.cc b/cpp/src/arrow/extension/opaque.cc index e206007ebe44..5fe904fce69c 100644 --- a/cpp/src/arrow/extension/opaque.cc +++ b/cpp/src/arrow/extension/opaque.cc @@ -19,12 +19,14 @@ #include +#include "arrow/json/json_writer_internal.h" #include "arrow/json/rapidjson_defs.h" // IWYU pragma: keep #include "arrow/util/logging_internal.h" #include #include -#include + +using ::arrow::json::JsonWriter; namespace arrow::extension { @@ -46,19 +48,16 @@ bool OpaqueType::ExtensionEquals(const ExtensionType& other) const { } std::string OpaqueType::Serialize() const { - rapidjson::Document document; - document.SetObject(); - rapidjson::Document::AllocatorType& allocator = document.GetAllocator(); - - rapidjson::Value type_name(rapidjson::StringRef(type_name_)); - document.AddMember(rapidjson::Value("type_name", allocator), type_name, allocator); - rapidjson::Value vendor_name(rapidjson::StringRef(vendor_name_)); - document.AddMember(rapidjson::Value("vendor_name", allocator), vendor_name, allocator); - - rapidjson::StringBuffer buffer; - rapidjson::Writer writer(buffer); - document.Accept(writer); - return buffer.GetString(); + JsonWriter writer; + + writer.StartObject(); + + writer.StringField("type_name", type_name_); + writer.StringField("vendor_name", vendor_name_); + + writer.EndObject(); + + return std::string(writer.GetString()); } Result> OpaqueType::Deserialize( diff --git a/cpp/src/arrow/extension/variable_shape_tensor.cc b/cpp/src/arrow/extension/variable_shape_tensor.cc index b1b12583d7fe..a67bd6dea88c 100644 --- a/cpp/src/arrow/extension/variable_shape_tensor.cc +++ b/cpp/src/arrow/extension/variable_shape_tensor.cc @@ -21,6 +21,7 @@ #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" @@ -30,9 +31,9 @@ #include "arrow/util/string.h" #include -#include namespace rj = arrow::rapidjson; +using ::arrow::json::JsonWriter; namespace arrow::extension { @@ -82,42 +83,44 @@ std::string VariableShapeTensorType::ToString(bool show_metadata) const { } std::string VariableShapeTensorType::Serialize() const { - rj::Document document; - document.SetObject(); - rj::Document::AllocatorType& allocator = document.GetAllocator(); + JsonWriter writer; + + writer.StartObject(); if (!permutation_.empty()) { - rj::Value permutation(rj::kArrayType); + writer.Key("permutation"); + writer.StartArray(); for (auto v : permutation_) { - permutation.PushBack(v, allocator); + writer.Int64(v); } - document.AddMember(rj::Value("permutation", allocator), permutation, allocator); + writer.EndArray(); } if (!dim_names_.empty()) { - rj::Value dim_names(rj::kArrayType); - for (const std::string& v : dim_names_) { - dim_names.PushBack(rj::Value{}.SetString(v.c_str(), allocator), allocator); + writer.Key("dim_names"); + writer.StartArray(); + for (const auto& v : dim_names_) { + writer.String(v); } - document.AddMember(rj::Value("dim_names", allocator), dim_names, allocator); + writer.EndArray(); } if (!uniform_shape_.empty()) { - rj::Value uniform_shape(rj::kArrayType); - for (auto v : uniform_shape_) { + writer.Key("uniform_shape"); + writer.StartArray(); + for (const auto& v : uniform_shape_) { if (v.has_value()) { - uniform_shape.PushBack(v.value(), allocator); + writer.Int64(*v); } else { - uniform_shape.PushBack(rj::Value{}.SetNull(), allocator); + writer.Null(); } } - document.AddMember(rj::Value("uniform_shape", allocator), uniform_shape, allocator); + writer.EndArray(); } - rj::StringBuffer buffer; - rj::Writer writer(buffer); - document.Accept(writer); - return buffer.GetString(); + writer.EndObject(); + + return std::string(writer.GetString()); } Result> VariableShapeTensorType::Deserialize(