| 1277 | template <bool kCoalesce> |
| 1278 | struct 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 | |
| 1335 | struct StreamWriterHelper { |
nothing calls this directly
no test coverage detected