MCPcopy Create free account
hub / github.com/apache/arrow / ReadBatches

Method ReadBatches

cpp/src/arrow/ipc/read_write_test.cc:1279–1332  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

1277template <bool kCoalesce>
1278struct FileGeneratorWriterHelper : public FileWriterHelper {
1279 Status ReadBatches(const IpcReadOptions& options, RecordBatchVector* out_batches,
1280 ReadStats* out_stats = nullptr,
1281 MetadataVector* out_metadata_list = nullptr) override {
1282 // The generator doesn't track stats.
1283 EXPECT_EQ(nullptr, out_stats);
1284
1285 auto read_batches = [&](bool pre_buffer) -> Result<RecordBatchVector> {
1286 std::shared_ptr<io::RandomAccessFile> buf_reader;
1287 if (kCoalesce) {
1288 // Use a non-zero-copy enabled BufferReader so we can test paths properly
1289 buf_reader = std::make_shared<NoZeroCopyBufferReader>(buffer_);
1290 } else {
1291 buf_reader = std::make_shared<io::BufferReader>(buffer_);
1292 }
1293 AsyncGenerator<std::shared_ptr<RecordBatch>> generator;
1294
1295 {
1296 auto fut = RecordBatchFileReader::OpenAsync(buf_reader, footer_offset_, options);
1297 ARROW_ASSIGN_OR_RAISE(auto reader, fut.result());
1298 EXPECT_EQ(num_batches_written_, reader->num_record_batches());
1299 if (pre_buffer) {
1300 RETURN_NOT_OK(reader->PreBufferMetadata(/*indices=*/{}));
1301 }
1302 // Generator will keep reader alive internally
1303 ARROW_ASSIGN_OR_RAISE(generator, reader->GetRecordBatchGenerator(kCoalesce));
1304 }
1305
1306 // Generator is async-reentrant
1307 std::vector<Future<std::shared_ptr<RecordBatch>>> futures;
1308 for (int i = 0; i < num_batches_written_; ++i) {
1309 futures.push_back(generator());
1310 }
1311 auto fut = generator();
1312 ARROW_ASSIGN_OR_RAISE(auto final_batch, fut.result());
1313 EXPECT_EQ(nullptr, final_batch);
1314
1315 RecordBatchVector batches;
1316 for (auto& future : futures) {
1317 ARROW_ASSIGN_OR_RAISE(auto batch, future.result());
1318 EXPECT_NE(nullptr, batch);
1319 batches.push_back(batch);
1320 }
1321 return batches;
1322 };
1323
1324 ARROW_ASSIGN_OR_RAISE(*out_batches, read_batches(/*pre_buffer=*/false));
1325 // Also read with pre-buffered metadata, and check the results are equal
1326 ARROW_ASSIGN_OR_RAISE(auto batches_pre_buffered, read_batches(/*pre_buffer=*/true));
1327 for (int i = 0; i < num_batches_written_; ++i) {
1328 AssertBatchesEqual(*batches_pre_buffered[i], *(*out_batches)[i],
1329 /*check_metadata=*/true);
1330 }
1331 return Status::OK();
1332 }
1333};
1334
1335struct StreamWriterHelper {

Callers

nothing calls this directly

Calls 7

AssertBatchesEqualFunction · 0.85
PreBufferMetadataMethod · 0.80
push_backMethod · 0.80
ARROW_ASSIGN_OR_RAISEFunction · 0.70
OKFunction · 0.50
num_record_batchesMethod · 0.45

Tested by

no test coverage detected