Repository: apache/arrow
File: cpp/src/arrow/csv/writer.cc (C++)

Code before:

  int64_t CalculateHeaderSize() const {
    int64_t header_length = 0;
    for (int col = 0; col < schema_->num_fields(); col++) {
      const std::string& col_name = schema_->field(col)->name();
      header_length += col_name.size();
      header_length += CountQuotes(col_name);
    }
    // header_length + ([quotes + ','] * schema_->num_fields()) + (eol - ',')
    return header_length + (kQuoteDelimiterCount * schema_->num_fields()) +
           (options_.eol.size() - 1);
  }

  Status WriteHeader() {
    // Only called once, as part of initialization
    RETURN_NOT_OK(data_buffer_->Resize(CalculateHeaderSize(), /*shrink_to_fit=*/false));
    char* next = reinterpret_cast<char*>(data_buffer_->mutable_data());
    for (int col = 0; col < schema_->num_fields(); ++col) {
      *next++ = '"';
      next = Escape(schema_->field(col)->name(), next);
      *next++ = '"';
      if (col != schema_->num_fields() - 1) {
        *next++ = options_.delimiter;
      }
    }
    memcpy(next, options_.eol.data(), options_.eol.size());
    next += options_.eol.size();
    DCHECK_EQ(reinterpret_cast<uint8_t*>(next),
              data_buffer_->data() + data_buffer_->size());
    return sink_->Write(data_buffer_);
  }

  Status TranslateMinimalBatch(const RecordBatch& batch) {
    if (batch.num_rows() == 0) {
      return Status::OK();
    }
    offsets_.resize(batch.num_rows());
    std::fill(offsets_.begin(), offsets_.end(), 0);

    // Calculate relative offsets for each row (excluding delimiters)
    for (int32_t col = 0; col < static_cast<int32_t>(column_populators_.size()); col++) {
      RETURN_NOT_OK(
          column_populators_[col]->UpdateRowLengths(*batch.column(col), offsets_.data()));
    }
    // Calculate cumulative offsets for each row (including delimiters).
    // - before conversion: offsets_[i] = length of i-th row
    // - after conversion:  offsets_[i] = offset to the starting of i-th row buffer
    //   - offsets_[0] = 0
    //   - offsets_[i] = offsets_[i-1] + len(i-1-th row) + len(delimiters)
    // Delimiters: ',' * (num_columns - 1) + eol
    const int32_t delimiters_length =

Review comment (suggestion):
Actually I don't know would calculateHeaderSize can bettering handle this without a round of resize. But current also looks ok to me

Code after:
 private:
  Status PrepareForContentsWrite() {
    // Only called once, as part of initialization
    if (data_buffer_ == nullptr) {
      ARROW_ASSIGN_OR_RAISE(
          data_buffer_,
          AllocateResizableBuffer(
              options_.batch_size * schema_->num_fields() * kColumnSizeGuess,
              options_.io_context.pool()));
    }
    return Status::OK();
  }

  int64_t CalculateHeaderSize(QuotingStyle quoting_style) const {
    int64_t header_length = 0;
    for (int col = 0; col < schema_->num_fields(); col++) {
      const std::string& col_name = schema_->field(col)->name();
      header_length += col_name.size();
      switch (quoting_style) {
        case QuotingStyle::None:
          break;
        case QuotingStyle::Needed:
        case QuotingStyle::AllValid:
          header_length += CountQuotes(col_name);
          break;
      }
    }
    header_length += kDelimiterCount * (schema_->num_fields() - 1) + options_.eol.size();
    switch (quoting_style) {
      case QuotingStyle::None:
        break;
      case QuotingStyle::Needed:
      case QuotingStyle::AllValid:
        header_length += kQuoteCount * schema_->num_fields();
        break;
    }
    return header_length;
  }

  Status WriteHeader() {
    // Only called once, as part of initialization
    RETURN_NOT_OK(data_buffer_->Resize(CalculateHeaderSize(options_.quoting_header),
                                       /*shrink_to_fit=*/false));
    char* next = reinterpret_cast<char*>(data_buffer_->mutable_data());
    for (int col = 0; col < schema_->num_fields(); ++col) {
      const std::string& col_name = schema_->field(col)->name();
      switch (options_.quoting_header) {
        case QuotingStyle::None:
          if (StopAtStructuralChar(reinterpret_cast<const uint8_t*>(col_name.c_str()),
                                   col_name.length(), options_.delimiter) !=
              static_cast<int64_t>(col_name.length())) {