Repository: apache/arrow
File: cpp/src/parquet/column_writer.cc (C++)

Code before:
    pager_->Close(has_dictionary_, fallback_);
  }

  return total_bytes_written_;
}

void ColumnWriterImpl::FlushBufferedDataPages() {
  // Write all outstanding data to a new page
  if (num_buffered_values_ > 0) {
    AddDataPage();
  }
  for (const auto& page_ptr : data_pages_) {
    WriteDataPage(*page_ptr);
  }
  data_pages_.clear();
  total_compressed_bytes_ = 0;
}

// ----------------------------------------------------------------------
// TypedColumnWriter

template <typename Action>
inline void DoInBatches(int64_t total, int64_t batch_size, Action&& action) {
  int64_t num_batches = static_cast<int>(total / batch_size);
  for (int round = 0; round < num_batches; round++) {
    action(round * batch_size, batch_size, /*check_page_size=*/true);
  }
  // Write the remaining values
  if (total % batch_size > 0) {
    action(num_batches * batch_size, total % batch_size, /*check_page_size=*/true);
  }
}

template <typename Action>
inline void DoInBatches(const int16_t* def_levels, const int16_t* rep_levels,
                        int64_t num_levels, int64_t batch_size, Action&& action,
                        bool pages_change_on_record_boundaries) {
  if (!pages_change_on_record_boundaries || !rep_levels) {
    // If rep_levels is null, then we are writing a non-repeated column.
    // In this case, every record contains only one level.
    return DoInBatches(num_levels, batch_size, std::forward<Action>(action));
  }

  int64_t offset = 0;
  while (offset < num_levels) {
    int64_t end_offset = std::min(offset + batch_size, num_levels);

    // Find next record boundary (i.e. rep_level = 0)
    while (end_offset < num_levels && rep_levels[end_offset] != 0) {
      end_offset++;
    }

Review comment (bug):
Like `Action`, can we use template function rather than `std::function<>`, which `std::function` is hard to optimize by compiler, and might be `nullptr`?

Code after:
    pager_->Close(has_dictionary_, fallback_);
  }

  return total_bytes_written_;
}

void ColumnWriterImpl::FlushBufferedDataPages() {
  // Write all outstanding data to a new page
  if (num_buffered_values_ > 0) {
    AddDataPage();
  }
  for (const auto& page_ptr : data_pages_) {
    WriteDataPage(*page_ptr);
  }
  data_pages_.clear();
  total_compressed_bytes_ = 0;
}

// ----------------------------------------------------------------------
// TypedColumnWriter

// DoInBatches for non-repeated columns
template <typename Action, typename GetBufferedRows>
inline void DoInBatchesNonRepeated(int64_t num_levels, int64_t batch_size,
                                   int64_t max_rows_per_page, Action&& action,
                                   GetBufferedRows&& curr_page_buffered_rows) {
  int64_t offset = 0;
  while (offset < num_levels) {
    int64_t page_buffered_rows = curr_page_buffered_rows();
    ARROW_DCHECK_LE(page_buffered_rows, max_rows_per_page);

    // Every record contains only one level.
    int64_t max_batch_size = std::min(batch_size, num_levels - offset);
    max_batch_size = std::min(max_batch_size, max_rows_per_page - page_buffered_rows);
    int64_t end_offset = offset + max_batch_size;

    ARROW_DCHECK_LE(offset, end_offset);
    ARROW_DCHECK_LE(end_offset, num_levels);

    // Always check page limit for non-repeated columns.
    action(offset, end_offset - offset, /*check_page_limit=*/true);

    offset = end_offset;
  }
}

// DoInBatches for repeated columns
template <typename Action, typename GetBufferedRows>
inline void DoInBatchesRepeated(const int16_t* def_levels, const int16_t* rep_levels,
                                int64_t num_levels, int64_t batch_size,
                                int64_t max_rows_per_page,