Function TestGetRecordBatchReader

cpp/src/parquet/arrow/arrow_reader_writer_test.cc:2557–2606

}

void TestGetRecordBatchReader(
    ArrowReaderProperties properties = default_arrow_reader_properties()) {
  const int num_columns = 20;
  const int num_rows = 1000;
  const int batch_size = 100;

  std::shared_ptr<Table> table;
  ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table));

  std::shared_ptr<Buffer> buffer;
  ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, num_rows / 2,
                                             default_arrow_writer_properties(), &buffer));

  properties.set_batch_size(batch_size);

  std::unique_ptr<FileReader> reader;
  FileReaderBuilder builder;
  ASSERT_OK(builder.Open(std::make_shared<BufferReader>(buffer)));
  ASSERT_OK(builder.properties(properties)->Build(&reader));

  // Read the whole file, one batch at a time.
  ASSERT_OK_AND_ASSIGN(auto rb_reader, reader->GetRecordBatchReader({0, 1}));
  std::shared_ptr<::arrow::RecordBatch> actual_batch, expected_batch;
  ::arrow::TableBatchReader table_reader(*table);
  table_reader.set_chunksize(batch_size);

  for (int i = 0; i < 10; ++i) {
    ASSERT_OK(rb_reader->ReadNext(&actual_batch));
    ASSERT_OK(table_reader.ReadNext(&expected_batch));
    ASSERT_NO_FATAL_FAILURE(::arrow::AssertBatchesEqual(*expected_batch, *actual_batch));
  }

  ASSERT_OK(rb_reader->ReadNext(&actual_batch));
  ASSERT_EQ(nullptr, actual_batch);

  // ARROW-6005: Read just the second row group
  ASSERT_OK_AND_ASSIGN(rb_reader, reader->GetRecordBatchReader({1}));
  std::shared_ptr<Table> second_rowgroup = table->Slice(num_rows / 2);
  ::arrow::TableBatchReader second_table_reader(*second_rowgroup);
  second_table_reader.set_chunksize(batch_size);

  for (int i = 0; i < 5; ++i) {
    ASSERT_OK(rb_reader->ReadNext(&actual_batch));
    ASSERT_OK(second_table_reader.ReadNext(&expected_batch));
    ASSERT_NO_FATAL_FAILURE(::arrow::AssertBatchesEqual(*expected_batch, *actual_batch));
  }

  ASSERT_OK(rb_reader->ReadNext(&actual_batch));
  ASSERT_EQ(nullptr, actual_batch);
}

TEST(TestArrowReadWrite, GetRecordBatchReader) { TestGetRecordBatchReader(); }