Skip to content

Commit 623a15e

Browse files
zagtopitrou
zagto
authored and
pitrou
committed
ARROW-15984: [C++] Change RecordBatchReader API to use Result<>
Closes apache#12678 from zagto/recordbatchreader-result Authored-by: Tobias Zagorni <tobias@zagorni.eu> Signed-off-by: Antoine Pitrou <antoine@python.org>
1 parent 7711182 commit 623a15e

18 files changed

Lines changed: 89 additions & 44 deletions

File tree

‎c_glib/arrow-glib/reader.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,7 @@ garrow_record_batch_reader_read_all(GArrowRecordBatchReader *reader,
343343
{
344344
auto arrow_reader = garrow_record_batch_reader_get_raw(reader);
345345
std::shared_ptr<arrow::Table> arrow_table;
346-
auto status = arrow_reader->ReadAll(&arrow_table);
346+
auto status = arrow_reader->ToTable().Value(&arrow_table);
347347

348348
if (garrow::check(error,
349349
status,

‎cpp/src/arrow/compute/kernels/vector_sort.cc‎

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -146,10 +146,7 @@ ArrayVector GetPhysicalChunks(const ChunkedArray& chunked_array,
146146
}
147147

148148
Result<RecordBatchVector> BatchesFromTable(const Table& table) {
149-
RecordBatchVector batches;
150-
TableBatchReader reader(table);
151-
RETURN_NOT_OK(reader.ReadAll(&batches));
152-
return batches;
149+
return TableBatchReader(table).ToRecordBatches();
153150
}
154151

155152
// ----------------------------------------------------------------------
@@ -937,11 +934,7 @@ class TableSorter {
937934

938935
Status SortInternal() {
939936
// Sort each batch independently and merge to sorted indices.
940-
RecordBatchVector batches;
941-
{
942-
TableBatchReader reader(table_);
943-
RETURN_NOT_OK(reader.ReadAll(&batches));
944-
}
937+
ARROW_ASSIGN_OR_RAISE(RecordBatchVector batches, BatchesFromTable(table_));
945938
const int64_t num_batches = static_cast<int64_t>(batches.size());
946939
if (num_batches == 0) {
947940
return Status::OK();

‎cpp/src/arrow/csv/reader_test.cc‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,7 @@ class StreamingReaderAsTableReader : public TableReader {
4848
: reader_(std::move(reader)) {}
4949
virtual ~StreamingReaderAsTableReader() = default;
5050
virtual Result<std::shared_ptr<Table>> Read() {
51-
std::shared_ptr<Table> table;
52-
RETURN_NOT_OK(reader_->ReadAll(&table));
51+
ARROW_ASSIGN_OR_RAISE(auto table, reader_->ToTable());
5352
return table;
5453
}
5554
virtual Future<std::shared_ptr<Table>> ReadAsync() {

‎cpp/src/arrow/dataset/file_csv_test.cc‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,7 @@ class CsvFormatHelper {
4343
using FormatType = CsvFileFormat;
4444
static Result<std::shared_ptr<Buffer>> Write(RecordBatchReader* reader) {
4545
ARROW_ASSIGN_OR_RAISE(auto sink, io::BufferOutputStream::Create());
46-
std::shared_ptr<Table> table;
47-
RETURN_NOT_OK(reader->ReadAll(&table));
46+
ARROW_ASSIGN_OR_RAISE(auto table, reader->ToTable());
4847
auto options = csv::WriteOptions::Defaults();
4948
RETURN_NOT_OK(csv::WriteCSV(*table, options, sink.get()));
5049
return sink->Finish();

‎cpp/src/arrow/dataset/file_ipc_test.cc‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,8 +47,7 @@ class IpcFormatHelper {
4747
static Result<std::shared_ptr<Buffer>> Write(RecordBatchReader* reader) {
4848
ARROW_ASSIGN_OR_RAISE(auto sink, io::BufferOutputStream::Create());
4949
ARROW_ASSIGN_OR_RAISE(auto writer, ipc::MakeFileWriter(sink, reader->schema()));
50-
std::vector<std::shared_ptr<RecordBatch>> batches;
51-
RETURN_NOT_OK(reader->ReadAll(&batches));
50+
ARROW_ASSIGN_OR_RAISE(auto batches, reader->ToRecordBatches());
5251
for (auto batch : batches) {
5352
RETURN_NOT_OK(writer->WriteRecordBatch(*batch));
5453
}

‎cpp/src/arrow/dataset/scanner_test.cc‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -577,11 +577,10 @@ TEST_P(TestScanner, ToRecordBatchReader) {
577577

578578
ASSERT_OK_AND_ASSIGN(auto expected, Table::FromRecordBatches(batches));
579579

580-
std::shared_ptr<Table> actual;
581580
auto scanner = MakeScanner(batch);
582581
ASSERT_OK_AND_ASSIGN(auto reader, scanner->ToRecordBatchReader());
583582
scanner.reset();
584-
ASSERT_OK(reader->ReadAll(&actual));
583+
ASSERT_OK_AND_ASSIGN(auto actual, reader->ToTable());
585584
AssertTablesEqual(*expected, *actual, /*same_chunk_layout=*/false);
586585
}
587586

‎cpp/src/arrow/filesystem/s3fs_benchmark.cc‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -314,14 +314,13 @@ static void ParquetRead(benchmark::State& st, S3FileSystem* fs, const std::strin
314314
ASSERT_OK(builder.Open(file, parquet_properties));
315315
ASSERT_OK(builder.properties(properties)->Build(&reader));
316316

317-
std::shared_ptr<Table> table;
318-
319317
if (read_strategy == "ReadTable") {
318+
std::shared_ptr<Table> table;
320319
ASSERT_OK(reader->ReadTable(column_indices, &table));
321320
} else {
322321
std::shared_ptr<RecordBatchReader> rb_reader;
323322
ASSERT_OK(reader->GetRecordBatchReader({0}, column_indices, &rb_reader));
324-
ASSERT_OK(rb_reader->ReadAll(&table));
323+
ASSERT_OK(rb_reader->ToTable());
325324
}
326325

327326
// TODO: actually measure table memory usage

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1156,7 +1156,7 @@ struct StreamWriterHelper {
11561156
ReadStats* out_stats = nullptr) {
11571157
auto buf_reader = std::make_shared<io::BufferReader>(buffer_);
11581158
ARROW_ASSIGN_OR_RAISE(auto reader, RecordBatchStreamReader::Open(buf_reader, options))
1159-
RETURN_NOT_OK(reader->ReadAll(out_batches));
1159+
ARROW_ASSIGN_OR_RAISE(*out_batches, reader->ToRecordBatches());
11601160
if (out_stats) {
11611161
*out_stats = reader->stats();
11621162
}

‎cpp/src/arrow/record_batch.cc‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -330,22 +330,30 @@ Status RecordBatch::ValidateFull() const {
330330
// ----------------------------------------------------------------------
331331
// Base record batch reader
332332

333-
Status RecordBatchReader::ReadAll(std::vector<std::shared_ptr<RecordBatch>>* batches) {
333+
Result<RecordBatchVector> RecordBatchReader::ToRecordBatches() {
334+
RecordBatchVector batches;
334335
while (true) {
335336
std::shared_ptr<RecordBatch> batch;
336337
RETURN_NOT_OK(ReadNext(&batch));
337338
if (!batch) {
338339
break;
339340
}
340-
batches->emplace_back(std::move(batch));
341+
batches.emplace_back(std::move(batch));
341342
}
342-
return Status::OK();
343+
return batches;
344+
}
345+
346+
Status RecordBatchReader::ReadAll(RecordBatchVector* batches) {
347+
return ToRecordBatches().Value(batches);
348+
}
349+
350+
Result<std::shared_ptr<Table>> RecordBatchReader::ToTable() {
351+
ARROW_ASSIGN_OR_RAISE(auto batches, ToRecordBatches());
352+
return Table::FromRecordBatches(schema(), std::move(batches));
343353
}
344354

345355
Status RecordBatchReader::ReadAll(std::shared_ptr<Table>* table) {
346-
std::vector<std::shared_ptr<RecordBatch>> batches;
347-
RETURN_NOT_OK(ReadAll(&batches));
348-
return Table::FromRecordBatches(schema(), std::move(batches)).Value(table);
356+
return ToTable().Value(table);
349357
}
350358

351359
class SimpleRecordBatchReader : public RecordBatchReader {

‎cpp/src/arrow/record_batch.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,9 +297,15 @@ class ARROW_EXPORT RecordBatchReader {
297297
RecordBatchReaderIterator end() { return RecordBatchReaderIterator(); }
298298

299299
/// \brief Consume entire stream as a vector of record batches
300+
Result<RecordBatchVector> ToRecordBatches();
301+
302+
ARROW_DEPRECATED("Deprecated in 8.0.0. Use ToRecordBatches instead.")
300303
Status ReadAll(RecordBatchVector* batches);
301304

302305
/// \brief Read all batches and concatenate as arrow::Table
306+
Result<std::shared_ptr<Table>> ToTable();
307+
308+
ARROW_DEPRECATED("Deprecated in 8.0.0. Use ToTable instead.")
303309
Status ReadAll(std::shared_ptr<Table>* table);
304310

305311
/// \brief Create a RecordBatchReader from a vector of RecordBatch.

0 commit comments

Comments
 (0)