Function main

bigquery/read/arrow/arrow_read.cc:118–182

}  // namespace

int main(int argc, char* argv[]) try {
  if (argc != 4) {
    std::cerr << "Usage: " << argv[0]
              << " <project-id> <dataset-name> <table-name>\n";
    return 1;
  }

  std::string const project_id = argv[1];
  std::string const dataset_name = argv[2];
  std::string const table_name = argv[3];

  std::string const table_id = "projects/" + project_id + "/datasets/" +
                               dataset_name + "/tables/" + table_name;

  // Create a namespace alias to make the code easier to read.
  namespace bigquery_storage = ::google::cloud::bigquery_storage_v1;
  constexpr int kMaxReadStreams = 1;
  // Create the ReadSession.
  auto client = bigquery_storage::BigQueryReadClient(
      bigquery_storage::MakeBigQueryReadConnection());
  ::google::cloud::bigquery::storage::v1::ReadSession read_session;
  read_session.set_data_format(
      google::cloud::bigquery::storage::v1::DataFormat::ARROW);
  read_session.set_table(table_id);
  auto session =
      client.CreateReadSession(google::cloud::Project(project_id).FullName(),
                               read_session, kMaxReadStreams);
  if (!session) throw std::move(session).status();

  // Get schema.
  std::shared_ptr<arrow::Schema> schema =
      GetArrowSchema(session->arrow_schema());

  // Read rows from the ReadSession.
  constexpr int kRowOffset = 0;
  auto read_rows = client.ReadRows(session->streams(0).name(), kRowOffset);

  std::int64_t num_rows = 0;
  std::int64_t record_batch_count = 0;
  for (auto const& read_rows_response : read_rows) {
    if (read_rows_response.ok()) {
      std::shared_ptr<arrow::RecordBatch> record_batch =
          GetArrowRecordBatch(read_rows_response->arrow_record_batch(), schema);

      if (record_batch_count == 0) {
        PrintColumnNames(record_batch);
      }

      ProcessRecordBatch(schema, record_batch, num_rows);
      num_rows += read_rows_response->row_count();
      ++record_batch_count;
    }
  }

  std::cout << std::format(
      "Read {} record batch(es) and {} total row(s) from table: {}\n",
      record_batch_count, num_rows, table_id);
  return 0;