Method RunExchangeTotal

cpp/src/arrow/flight/test_flight_server.cc:260–311

Read int64 batches from the client, each time sending back a batch with a running sum of columns.

// Read int64 batches from the client, each time sending back a
// batch with a running sum of columns.
Status TestFlightServer::RunExchangeTotal(std::unique_ptr<FlightMessageReader> reader,
                                          std::unique_ptr<FlightMessageWriter> writer) {
  FlightStreamChunk chunk{};
  ARROW_ASSIGN_OR_RAISE(auto schema, reader->GetSchema());
  // Ensure the schema contains only int64 columns
  for (const auto& field : schema->fields()) {
    if (field->type()->id() != Type::type::INT64) {
      return Status::Invalid("Field is not INT64: ", field->name());
    }
  }
  std::vector<int64_t> sums(schema->num_fields());
  std::vector<std::shared_ptr<Array>> columns(schema->num_fields());
  RETURN_NOT_OK(writer->Begin(schema));
  while (true) {
    ARROW_ASSIGN_OR_RAISE(chunk, reader->Next());
    if (!chunk.data && !chunk.app_metadata) {
      break;
    }
    if (chunk.data) {
      if (!chunk.data->schema()->Equals(schema, false)) {
        // A compliant client implementation would make this impossible
        return Status::Invalid("Schemas are incompatible");
      }

      // Update the running totals
      auto builder = std::make_shared<Int64Builder>();
      int col_index = 0;
      for (const auto& column : chunk.data->columns()) {
        auto arr = std::dynamic_pointer_cast<Int64Array>(column);
        if (!arr) {
          return MakeFlightError(FlightStatusCode::Internal, "Could not cast array");
        }
        for (int row = 0; row < column->length(); row++) {
          if (!arr->IsNull(row)) {
            sums[col_index] += arr->Value(row);
          }
        }

        builder->Reset();
        RETURN_NOT_OK(builder->Append(sums[col_index]));
        RETURN_NOT_OK(builder->Finish(&columns[col_index]));

        col_index++;
      }

      // Echo the totals to the client
      auto response = RecordBatch::Make(schema, /* num_rows */ 1, columns);
      RETURN_NOT_OK(writer->WriteRecordBatch(*response));
    }
  }
  return Status::OK();
}

// Echo the client's messages back.
Status TestFlightServer::RunExchangeEcho(std::unique_ptr<FlightMessageReader> reader,