Skip to content

Commit e4bcb83

Browse files
Shally-KatariyaShally-Katariya
committed
[C++] Migrate arrow/ipc/metadata_internal.h to Result<T>
1 parent 12b3eda commit e4bcb83

7 files changed

Lines changed: 65 additions & 63 deletions

File tree

‎cpp/src/arrow/ipc/message.cc‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -102,9 +102,8 @@ class Message::MessageImpl {
102102

103103
if (message_->custom_metadata() != nullptr) {
104104
// Deserialize from Flatbuffers if first time called
105-
std::shared_ptr<KeyValueMetadata> md;
106-
RETURN_NOT_OK(internal::GetKeyValueMetadata(message_->custom_metadata(), &md));
107-
custom_metadata_ = std::move(md); // const-ify
105+
ARROW_ASSIGN_OR_RAISE(custom_metadata_,
106+
internal::GetKeyValueMetadata(message_->custom_metadata()));
108107
}
109108

110109
return Status::OK();

‎cpp/src/arrow/ipc/message_internal_test.cc‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -62,9 +62,8 @@ TEST(TestMessageInternal, TestByteIdentical) {
6262
auto schema = ::arrow::schema({f0}, Endianness::Little, metadata);
6363

6464
// Serialize the Schema to a Buffer
65-
std::shared_ptr<Buffer> out_buffer;
66-
ASSERT_OK(
67-
WriteSchemaMessage(*schema, mapper, IpcWriteOptions::Defaults(), &out_buffer));
65+
ASSERT_OK_AND_ASSIGN(auto out_buffer,
66+
WriteSchemaMessage(*schema, mapper, IpcWriteOptions::Defaults()));
6867

6968
// This is example output from macOS+ARM+LLVM
7069
const uint8_t expected[] = {
@@ -104,9 +103,9 @@ TEST(TestMessageInternal, TestEndiannessRoundtrip) {
104103
auto schema = ::arrow::schema({f0}, endianness, metadata);
105104

106105
// Serialize the Schema to a Buffer
107-
std::shared_ptr<Buffer> out_buffer;
108-
ASSERT_OK(
109-
WriteSchemaMessage(*schema, mapper, IpcWriteOptions::Defaults(), &out_buffer));
106+
ASSERT_OK_AND_ASSIGN(
107+
auto out_buffer,
108+
WriteSchemaMessage(*schema, mapper, IpcWriteOptions::Defaults()));
110109

111110
// Re-open to a new Message and parse Schema
112111
ASSERT_OK_AND_ASSIGN(auto message, Message::Open(out_buffer, /*body=*/nullptr));
@@ -183,9 +182,10 @@ class MessageDecodingTest : public ::testing::Test {
183182
auto buffer_md = std::vector{BufferMetadata{.offset = 64, .length = 1},
184183
BufferMetadata{.offset = 72, .length = 10}};
185184
std::shared_ptr<Buffer> out;
186-
RETURN_NOT_OK(WriteRecordBatchMessage(/*length=*/kNumRows, params.body_length,
187-
params.custom_metadata, field_md, buffer_md,
188-
/*variadic_counts=*/{}, params.options, &out));
185+
ARROW_ASSIGN_OR_RAISE(
186+
out, WriteRecordBatchMessage(/*length=*/kNumRows, params.body_length,
187+
params.custom_metadata, field_md, buffer_md,
188+
/*variadic_counts=*/{}, params.options));
189189
ARROW_ASSIGN_OR_RAISE(out, EncapsulateMetadata(out, params.options));
190190
// Generate a dummy body of the advertised length
191191
ARROW_ASSIGN_OR_RAISE(auto body_bytes, AllocateBuffer(params.body_length));

‎cpp/src/arrow/ipc/metadata_internal.cc‎

Lines changed: 18 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -866,7 +866,8 @@ Status FieldFromFlatbuffer(const flatbuf::Field* field, FieldPosition field_pos,
866866
std::shared_ptr<DataType> type;
867867

868868
std::shared_ptr<KeyValueMetadata> metadata;
869-
RETURN_NOT_OK(internal::GetKeyValueMetadata(field->custom_metadata(), &metadata));
869+
ARROW_ASSIGN_OR_RAISE(metadata,
870+
internal::GetKeyValueMetadata(field->custom_metadata()));
870871

871872
// Reconstruct the data type
872873
// 1. Data type children
@@ -1279,11 +1280,10 @@ Status MakeSparseTensor(FBB& fbb, const SparseTensor& sparse_tensor, int64_t bod
12791280

12801281
} // namespace
12811282

1282-
Status GetKeyValueMetadata(const KVVector* fb_metadata,
1283-
std::shared_ptr<KeyValueMetadata>* out) {
1283+
Result<std::shared_ptr<KeyValueMetadata>> GetKeyValueMetadata(
1284+
const KVVector* fb_metadata) {
12841285
if (fb_metadata == nullptr) {
1285-
*out = nullptr;
1286-
return Status::OK();
1286+
return nullptr;
12871287
}
12881288

12891289
auto metadata = std::make_shared<KeyValueMetadata>();
@@ -1295,36 +1295,33 @@ Status GetKeyValueMetadata(const KVVector* fb_metadata,
12951295
metadata->Append(pair->key()->str(), pair->value()->str());
12961296
}
12971297

1298-
*out = std::move(metadata);
1299-
return Status::OK();
1298+
return metadata;
13001299
}
13011300

1302-
Status WriteSchemaMessage(const Schema& schema, const DictionaryFieldMapper& mapper,
1303-
const IpcWriteOptions& options, std::shared_ptr<Buffer>* out) {
1301+
Result<std::shared_ptr<Buffer>> WriteSchemaMessage(const Schema& schema,
1302+
const DictionaryFieldMapper& mapper,
1303+
const IpcWriteOptions& options) {
13041304
FBB fbb;
13051305
flatbuffers::Offset<flatbuf::Schema> fb_schema;
13061306
RETURN_NOT_OK(SchemaToFlatbuffer(fbb, schema, mapper, &fb_schema));
13071307
return WriteFBMessage(fbb, flatbuf::MessageHeader::MessageHeader_Schema,
13081308
fb_schema.Union(),
13091309
/*body_length=*/0, options.metadata_version,
1310-
/*custom_metadata=*/nullptr, options.memory_pool)
1311-
.Value(out);
1310+
/*custom_metadata=*/nullptr, options.memory_pool);
13121311
}
13131312

1314-
Status WriteRecordBatchMessage(
1313+
Result<std::shared_ptr<Buffer>> WriteRecordBatchMessage(
13151314
int64_t length, int64_t body_length,
13161315
const std::shared_ptr<const KeyValueMetadata>& custom_metadata,
13171316
const std::vector<FieldMetadata>& nodes, const std::vector<BufferMetadata>& buffers,
1318-
const std::vector<int64_t>& variadic_buffer_counts, const IpcWriteOptions& options,
1319-
std::shared_ptr<Buffer>* out) {
1317+
const std::vector<int64_t>& variadic_buffer_counts, const IpcWriteOptions& options) {
13201318
FBB fbb;
13211319
RecordBatchOffset record_batch;
13221320
RETURN_NOT_OK(MakeRecordBatch(fbb, length, body_length, nodes, buffers,
13231321
variadic_buffer_counts, options, &record_batch));
13241322
return WriteFBMessage(fbb, flatbuf::MessageHeader::MessageHeader_RecordBatch,
13251323
record_batch.Union(), body_length, options.metadata_version,
1326-
custom_metadata, options.memory_pool)
1327-
.Value(out);
1324+
custom_metadata, options.memory_pool);
13281325
}
13291326

13301327
Result<std::shared_ptr<Buffer>> WriteTensorMessage(const Tensor& tensor,
@@ -1373,12 +1370,11 @@ Result<std::shared_ptr<Buffer>> WriteSparseTensorMessage(
13731370
/*custom_metadata=*/nullptr, options.memory_pool);
13741371
}
13751372

1376-
Status WriteDictionaryMessage(
1373+
Result<std::shared_ptr<Buffer>> WriteDictionaryMessage(
13771374
int64_t id, bool is_delta, int64_t length, int64_t body_length,
13781375
const std::shared_ptr<const KeyValueMetadata>& custom_metadata,
13791376
const std::vector<FieldMetadata>& nodes, const std::vector<BufferMetadata>& buffers,
1380-
const std::vector<int64_t>& variadic_buffer_counts, const IpcWriteOptions& options,
1381-
std::shared_ptr<Buffer>* out) {
1377+
const std::vector<int64_t>& variadic_buffer_counts, const IpcWriteOptions& options) {
13821378
FBB fbb;
13831379
RecordBatchOffset record_batch;
13841380
RETURN_NOT_OK(MakeRecordBatch(fbb, length, body_length, nodes, buffers,
@@ -1387,8 +1383,7 @@ Status WriteDictionaryMessage(
13871383
flatbuf::CreateDictionaryBatch(fbb, id, record_batch, is_delta).Union();
13881384
return WriteFBMessage(fbb, flatbuf::MessageHeader::MessageHeader_DictionaryBatch,
13891385
dictionary_batch, body_length, options.metadata_version,
1390-
custom_metadata, options.memory_pool)
1391-
.Value(out);
1386+
custom_metadata, options.memory_pool);
13921387
}
13931388

13941389
static flatbuffers::Offset<flatbuffers::Vector<const flatbuf::Block*>>
@@ -1462,7 +1457,8 @@ Status GetSchema(const void* opaque_schema, DictionaryMemo* dictionary_memo,
14621457
}
14631458

14641459
std::shared_ptr<KeyValueMetadata> metadata;
1465-
RETURN_NOT_OK(internal::GetKeyValueMetadata(schema->custom_metadata(), &metadata));
1460+
ARROW_ASSIGN_OR_RAISE(metadata,
1461+
internal::GetKeyValueMetadata(schema->custom_metadata()));
14661462
// set endianness using the value in flatbuf schema
14671463
auto endianness = schema->endianness() == flatbuf::Endianness::Endianness_Little
14681464
? Endianness::Little

‎cpp/src/arrow/ipc/metadata_internal.h‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -162,8 +162,8 @@ Status GetSparseTensorMetadata(const Buffer& metadata, std::shared_ptr<DataType>
162162
SparseTensorFormat::type* sparse_tensor_format_id);
163163

164164
ARROW_EXPORT
165-
Status GetKeyValueMetadata(const KVVector* fb_metadata,
166-
std::shared_ptr<KeyValueMetadata>* out);
165+
Result<std::shared_ptr<KeyValueMetadata>> GetKeyValueMetadata(
166+
const KVVector* fb_metadata);
167167

168168
ARROW_EXPORT
169169
Status ConcreteTypeFromFlatbuffer(flatbuf::Type type, const void* type_data,
@@ -193,17 +193,17 @@ static inline Status VerifyMessage(const uint8_t* data, int64_t size,
193193

194194
// Serialize arrow::Schema as a Flatbuffer
195195
ARROW_EXPORT
196-
Status WriteSchemaMessage(const Schema& schema, const DictionaryFieldMapper& mapper,
197-
const IpcWriteOptions& options, std::shared_ptr<Buffer>* out);
196+
Result<std::shared_ptr<Buffer>> WriteSchemaMessage(const Schema& schema,
197+
const DictionaryFieldMapper& mapper,
198+
const IpcWriteOptions& options);
198199

199200
// This function is used in a unit test
200201
ARROW_EXPORT
201-
Status WriteRecordBatchMessage(
202+
Result<std::shared_ptr<Buffer>> WriteRecordBatchMessage(
202203
const int64_t length, const int64_t body_length,
203204
const std::shared_ptr<const KeyValueMetadata>& custom_metadata,
204205
const std::vector<FieldMetadata>& nodes, const std::vector<BufferMetadata>& buffers,
205-
const std::vector<int64_t>& variadic_counts, const IpcWriteOptions& options,
206-
std::shared_ptr<Buffer>* out);
206+
const std::vector<int64_t>& variadic_counts, const IpcWriteOptions& options);
207207

208208
ARROW_EXPORT
209209
Result<std::shared_ptr<Buffer>> WriteTensorMessage(const Tensor& tensor,
@@ -222,13 +222,12 @@ Status WriteFileFooter(const Schema& schema, const std::vector<FileBlock>& dicti
222222
io::OutputStream* out);
223223

224224
ARROW_EXPORT
225-
Status WriteDictionaryMessage(
225+
Result<std::shared_ptr<Buffer>> WriteDictionaryMessage(
226226
const int64_t id, const bool is_delta, const int64_t length,
227227
const int64_t body_length,
228228
const std::shared_ptr<const KeyValueMetadata>& custom_metadata,
229229
const std::vector<FieldMetadata>& nodes, const std::vector<BufferMetadata>& buffers,
230-
const std::vector<int64_t>& variadic_counts, const IpcWriteOptions& options,
231-
std::shared_ptr<Buffer>* out);
230+
const std::vector<int64_t>& variadic_counts, const IpcWriteOptions& options);
232231

233232
static inline Result<std::shared_ptr<Buffer>> WriteFlatbufferBuilder(
234233
flatbuffers::FlatBufferBuilder& fbb, // NOLINT non-const reference

‎cpp/src/arrow/ipc/read_write_test.cc‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -156,11 +156,11 @@ TEST_P(TestMessage, SerializeCustomMetadata) {
156156
nullptr, key_value_metadata({}, {}),
157157
key_value_metadata({"foo", "bar"}, {"fizz", "buzz"})};
158158
for (auto metadata : cases) {
159-
std::shared_ptr<Buffer> serialized;
160-
ASSERT_OK(internal::WriteRecordBatchMessage(
161-
/*length=*/0, /*body_length=*/0, metadata,
162-
/*nodes=*/{},
163-
/*buffers=*/{}, /*variadic_counts=*/{}, options_, &serialized));
159+
ASSERT_OK_AND_ASSIGN(auto serialized,
160+
internal::WriteRecordBatchMessage(
161+
/*length=*/0, /*body_length=*/0, metadata,
162+
/*nodes=*/{},
163+
/*buffers=*/{}, /*variadic_counts=*/{}, options_));
164164
ASSERT_OK_AND_ASSIGN(std::unique_ptr<Message> message,
165165
Message::Open(serialized, /*body=*/nullptr));
166166

‎cpp/src/arrow/ipc/reader.cc‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -765,7 +765,8 @@ Status GetCompressionExperimental(const flatbuf::Message* message,
765765
if (message->custom_metadata() != nullptr) {
766766
// TODO: Ensure this deserialization only ever happens once
767767
std::shared_ptr<KeyValueMetadata> metadata;
768-
RETURN_NOT_OK(internal::GetKeyValueMetadata(message->custom_metadata(), &metadata));
768+
ARROW_ASSIGN_OR_RAISE(metadata,
769+
internal::GetKeyValueMetadata(message->custom_metadata()));
769770
int index = metadata->FindKey("ARROW:experimental_compression");
770771
if (index != -1) {
771772
// Arrow 0.17 stored string in upper case, internal utils now require lower case
@@ -810,8 +811,8 @@ Result<RecordBatchWithMetadata> ReadRecordBatchInternal(
810811

811812
std::shared_ptr<KeyValueMetadata> custom_metadata;
812813
if (message->custom_metadata() != nullptr) {
813-
RETURN_NOT_OK(
814-
internal::GetKeyValueMetadata(message->custom_metadata(), &custom_metadata));
814+
ARROW_ASSIGN_OR_RAISE(custom_metadata,
815+
internal::GetKeyValueMetadata(message->custom_metadata()));
815816
}
816817
ARROW_ASSIGN_OR_RAISE(auto record_batch,
817818
LoadRecordBatch(batch, schema, inclusion_mask, context, file));
@@ -1426,8 +1427,8 @@ class RecordBatchFileReaderImpl : public RecordBatchFileReader {
14261427
ARROW_ASSIGN_OR_RAISE(auto message, GetFlatbufMessage(message_obj));
14271428
std::shared_ptr<KeyValueMetadata> custom_metadata;
14281429
if (message->custom_metadata() != nullptr) {
1429-
RETURN_NOT_OK(
1430-
internal::GetKeyValueMetadata(message->custom_metadata(), &custom_metadata));
1430+
ARROW_ASSIGN_OR_RAISE(custom_metadata,
1431+
internal::GetKeyValueMetadata(message->custom_metadata()));
14311432
}
14321433
return RecordBatchWithMetadata{std::move(batch), std::move(custom_metadata)};
14331434
}
@@ -1928,8 +1929,8 @@ class RecordBatchFileReaderImpl : public RecordBatchFileReader {
19281929
auto fb_metadata = self->footer_->custom_metadata();
19291930
if (fb_metadata != nullptr) {
19301931
std::shared_ptr<KeyValueMetadata> md;
1931-
RETURN_NOT_OK(internal::GetKeyValueMetadata(fb_metadata, &md));
1932-
self->metadata_ = std::move(md); // const-ify
1932+
ARROW_ASSIGN_OR_RAISE(self->metadata_,
1933+
internal::GetKeyValueMetadata(fb_metadata));
19331934
}
19341935
return Status::OK();
19351936
});

‎cpp/src/arrow/ipc/writer.cc‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -180,9 +180,11 @@ class RecordBatchSerializer {
180180

181181
// Override this for writing dictionary metadata
182182
virtual Status SerializeMetadata(int64_t num_rows) {
183-
return WriteRecordBatchMessage(num_rows, out_->body_length, custom_metadata_,
184-
field_nodes_, buffer_meta_, variadic_counts_, options_,
185-
&out_->metadata);
183+
ARROW_ASSIGN_OR_RAISE(
184+
out_->metadata,
185+
WriteRecordBatchMessage(num_rows, out_->body_length, custom_metadata_,
186+
field_nodes_, buffer_meta_, variadic_counts_, options_));
187+
return Status::OK();
186188
}
187189

188190
bool ShouldCompress(int64_t uncompressed_size, int64_t compressed_size) const {
@@ -746,9 +748,12 @@ class DictionarySerializer : public RecordBatchSerializer {
746748
is_delta_(is_delta) {}
747749

748750
Status SerializeMetadata(int64_t num_rows) override {
749-
return WriteDictionaryMessage(dictionary_id_, is_delta_, num_rows, out_->body_length,
750-
custom_metadata_, field_nodes_, buffer_meta_,
751-
variadic_counts_, options_, &out_->metadata);
751+
ARROW_ASSIGN_OR_RAISE(
752+
out_->metadata,
753+
WriteDictionaryMessage(dictionary_id_, is_delta_, num_rows, out_->body_length,
754+
custom_metadata_, field_nodes_, buffer_meta_,
755+
variadic_counts_, options_));
756+
return Status::OK();
752757
}
753758

754759
Status Assemble(const std::shared_ptr<Array>& dictionary) {
@@ -804,7 +809,9 @@ Status WriteIpcPayload(const IpcPayload& payload, const IpcWriteOptions& options
804809
Status GetSchemaPayload(const Schema& schema, const IpcWriteOptions& options,
805810
const DictionaryFieldMapper& mapper, IpcPayload* out) {
806811
out->type = MessageType::SCHEMA;
807-
return internal::WriteSchemaMessage(schema, mapper, options, &out->metadata);
812+
ARROW_ASSIGN_OR_RAISE(out->metadata,
813+
internal::WriteSchemaMessage(schema, mapper, options));
814+
return Status::OK();
808815
}
809816

810817
Status GetDictionaryPayload(int64_t id, const std::shared_ptr<Array>& dictionary,

0 commit comments

Comments
 (0)