Method WriteRecordBatch

cpp/src/parquet/arrow/writer.cc:451–511

  }

  Status WriteRecordBatch(const RecordBatch& batch) override {
    RETURN_NOT_OK(CheckClosed());
    if (batch.num_rows() == 0) {
      return Status::OK();
    }

    // Max number of rows allowed in a row group.
    const int64_t max_row_group_length = this->properties().max_row_group_length();

    // Initialize a new buffered row group writer if necessary.
    if (row_group_writer_ == nullptr || !row_group_writer_->buffered() ||
        row_group_writer_->num_rows() >= max_row_group_length) {
      RETURN_NOT_OK(NewBufferedRowGroup());
    }

    auto WriteBatch = [&](int64_t offset, int64_t size) {
      std::vector<std::unique_ptr<ArrowColumnWriterV2>> writers;
      int column_index_start = 0;

      for (int i = 0; i < batch.num_columns(); i++) {
        ChunkedArray chunked_array{batch.column(i)};
        ARROW_ASSIGN_OR_RAISE(
            std::unique_ptr<ArrowColumnWriterV2> writer,
            ArrowColumnWriterV2::Make(chunked_array, offset, size, schema_manifest_,
                                      row_group_writer_, column_index_start));
        column_index_start += writer->leaf_count();
        if (arrow_properties_->use_threads()) {
          writers.emplace_back(std::move(writer));
        } else {
          RETURN_NOT_OK(writer->Write(&column_write_context_));
        }
      }

      if (arrow_properties_->use_threads()) {
        DCHECK_EQ(parallel_column_write_contexts_.size(), writers.size());
        RETURN_NOT_OK(::arrow::internal::ParallelFor(
            static_cast<int>(writers.size()),
            [&](int i) { return writers[i]->Write(&parallel_column_write_contexts_[i]); },
            arrow_properties_->executor()));
      }

      return Status::OK();
    };

    int64_t offset = 0;
    while (offset < batch.num_rows()) {
      const int64_t batch_size =
          std::min(max_row_group_length - row_group_writer_->num_rows(),
                   batch.num_rows() - offset);
      RETURN_NOT_OK(WriteBatch(offset, batch_size));
      offset += batch_size;

      // Flush current row group writer and create a new writer if it is full.
      if (row_group_writer_->num_rows() >= max_row_group_length &&
          offset < batch.num_rows()) {
        RETURN_NOT_OK(NewBufferedRowGroup());
      }
    }