GH-45201: [C++][Parquet] Improve performance of generating size statistics

cpp/src/parquet/column_writer.cc

@@ -1468,42 +1468,43 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
  // which case we call back to the dense write path)
  std::shared_ptr<::arrow::Array> preserved_dictionary_;
 
- int64_t WriteLevels(int64_t num_values, const int16_t* def_levels,
+ int64_t WriteLevels(int64_t num_levels, const int16_t* def_levels,
  const int16_t* rep_levels) {
+ // Update histograms now, to maximize cache efficiency.
+ UpdateLevelHistogram(num_levels, def_levels, rep_levels);
+
  int64_t values_to_write = 0;
  // If the field is required and non-repeated, there are no definition levels
  if (descr_->max_definition_level() > 0) {
- for (int64_t i = 0; i < num_values; ++i) {
+ for (int64_t i = 0; i < num_levels; ++i) {
  if (def_levels[i] == descr_->max_definition_level()) {
  ++values_to_write;
  }
  }
 
- WriteDefinitionLevels(num_values, def_levels);
+ WriteDefinitionLevels(num_levels, def_levels);
  } else {
  // Required field, write all values
- values_to_write = num_values;
+ values_to_write = num_levels;
  }
 
  // Not present for non-repeated fields
  if (descr_->max_repetition_level() > 0) {
  // A row could include more than one value
  // Count the occasions where we start a new row
- for (int64_t i = 0; i < num_values; ++i) {
+ for (int64_t i = 0; i < num_levels; ++i) {
  if (rep_levels[i] == 0) {
  rows_written_++;
  num_buffered_rows_++;
  }
  }
 
- WriteRepetitionLevels(num_values, rep_levels);
+ WriteRepetitionLevels(num_levels, rep_levels);
  } else {
  // Each value is exactly one row
- rows_written_ += num_values;
- num_buffered_rows_ += num_values;
+ rows_written_ += num_levels;
+ num_buffered_rows_ += num_levels;
  }
-
- UpdateLevelHistogram(num_values, def_levels, rep_levels);
  return values_to_write;
  }
 
@@ -1575,6 +1576,9 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
 
  void WriteLevelsSpaced(int64_t num_levels, const int16_t* def_levels,
  const int16_t* rep_levels) {
+ // Update histograms now, to maximize cache efficiency.
+ UpdateLevelHistogram(num_levels, def_levels, rep_levels);
+
  // If the field is required and non-repeated, there are no definition levels
  if (descr_->max_definition_level() > 0) {
  WriteDefinitionLevels(num_levels, def_levels);
@@ -1595,8 +1599,6 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
  rows_written_ += num_levels;
  num_buffered_rows_ += num_levels;
  }
-
- UpdateLevelHistogram(num_levels, def_levels, rep_levels);
  }
 
  void UpdateLevelHistogram(int64_t num_levels, const int16_t* def_levels,
@@ -1606,26 +1608,17 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
  }
 
  auto add_levels = [](std::vector<int64_t>& level_histogram,
- ::arrow::util::span<const int16_t> levels) {
- for (int16_t level : levels) {
- ARROW_DCHECK_LT(level, static_cast<int16_t>(level_histogram.size()));
- ++level_histogram[level];
- }
+ ::arrow::util::span<const int16_t> levels, int16_t max_level) {
+ ARROW_DCHECK_EQ(static_cast<size_t>(max_level) + 1, level_histogram.size());
+ ::parquet::UpdateLevelHistogram(levels, level_histogram);
  };
 
- if (descr_->max_definition_level() > 0) {
- add_levels(page_size_statistics_->definition_level_histogram,
- {def_levels, static_cast<size_t>(num_levels)});
- } else {
- page_size_statistics_->definition_level_histogram[0] += num_levels;
- }
-
- if (descr_->max_repetition_level() > 0) {
- add_levels(page_size_statistics_->repetition_level_histogram,
- {rep_levels, static_cast<size_t>(num_levels)});
- } else {
- page_size_statistics_->repetition_level_histogram[0] += num_levels;
- }
+ add_levels(page_size_statistics_->definition_level_histogram,
+ {def_levels, static_cast<size_t>(num_levels)},
+ descr_->max_definition_level());
+ add_levels(page_size_statistics_->repetition_level_histogram,
+ {rep_levels, static_cast<size_t>(num_levels)},
+ descr_->max_repetition_level());
  }
 
  // Update the unencoded data bytes for ByteArray only per the specification.

cpp/src/parquet/size_statistics.cc

@@ -18,13 +18,27 @@
 #include "parquet/size_statistics.h"
 
 #include <algorithm>
+#include <numeric>
+#include <ostream>
+#include <string_view>
 
 #include "arrow/util/logging.h"
 #include "parquet/exception.h"
 #include "parquet/schema.h"
 
 namespace parquet {
 
+namespace {
+
+void MergeLevelHistogram(::arrow::util::span<int64_t> histogram,
+ ::arrow::util::span<const int64_t> other) {
+ ARROW_DCHECK_EQ(histogram.size(), other.size());
+ std::transform(histogram.begin(), histogram.end(), other.begin(), histogram.begin(),
+ std::plus<>());
+}
+
+} // namespace
+
 void SizeStatistics::Merge(const SizeStatistics& other) {
  if (repetition_level_histogram.size() != other.repetition_level_histogram.size()) {
  throw ParquetException("Repetition level histogram size mismatch");
@@ -36,12 +50,8 @@ void SizeStatistics::Merge(const SizeStatistics& other) {
  other.unencoded_byte_array_data_bytes.has_value()) {
  throw ParquetException("Unencoded byte array data bytes are not consistent");
  }
- std::transform(repetition_level_histogram.begin(), repetition_level_histogram.end(),
- other.repetition_level_histogram.begin(),
- repetition_level_histogram.begin(), std::plus<>());
- std::transform(definition_level_histogram.begin(), definition_level_histogram.end(),
- other.definition_level_histogram.begin(),
- definition_level_histogram.begin(), std::plus<>());
+ MergeLevelHistogram(repetition_level_histogram, other.repetition_level_histogram);
+ MergeLevelHistogram(definition_level_histogram, other.definition_level_histogram);
  if (unencoded_byte_array_data_bytes.has_value()) {
  unencoded_byte_array_data_bytes = unencoded_byte_array_data_bytes.value() +
  other.unencoded_byte_array_data_bytes.value();
@@ -91,4 +101,88 @@ std::unique_ptr<SizeStatistics> SizeStatistics::Make(const ColumnDescriptor* des
  return size_stats;
 }
 
+std::ostream& operator<<(std::ostream& os, const SizeStatistics& size_stats) {
+ constexpr std::string_view kComma = ", ";
+ os << "SizeStatistics{";
+ std::string_view sep = "";
+ if (size_stats.unencoded_byte_array_data_bytes.has_value()) {
+ os << "unencoded_byte_array_data_bytes="
+ << *size_stats.unencoded_byte_array_data_bytes;
+ sep = kComma;
+ }
+ auto print_histogram = [&](std::string_view name,
+ const std::vector<int64_t>& histogram) {
+ if (!histogram.empty()) {
+ os << sep << name << "={";
+ sep = kComma;
+ std::string_view value_sep = "";
+ for (int64_t v : histogram) {
+ os << value_sep << v;
+ value_sep = kComma;
+ }
+ os << "}";
+ }
+ };
+ print_histogram("repetition_level_histogram", size_stats.repetition_level_histogram);
+ print_histogram("definition_level_histogram", size_stats.definition_level_histogram);
+ os << "}";
+ return os;
+}
+
+void UpdateLevelHistogram(::arrow::util::span<const int16_t> levels,
+ ::arrow::util::span<int64_t> histogram) {
+ const int64_t num_levels = static_cast<int64_t>(levels.size());
+ DCHECK_GE(histogram.size(), 1);
+ const int16_t max_level = static_cast<int16_t>(histogram.size() - 1);
+ if (max_level == 0) {
+ histogram[0] += num_levels;
+ return;
+ }
+
+#ifndef NDEBUG
+ for (auto level : levels) {
+ ARROW_DCHECK_LE(level, max_level);
+ }
+#endif
+
+ if (max_level == 1) {
+ // Specialize the common case for non-repeated non-nested columns.
+ // Summing the levels gives us the number of 1s, and the number of 0s follows.
+ // We do repeated sums in the int16_t space, which the compiler is likely
+ // to vectorize efficiently.
+ constexpr int64_t kChunkSize = 1 << 14; // to avoid int16_t overflows
+ int64_t hist1 = 0;
+ auto it = levels.begin();
+ while (it != levels.end()) {
+ const auto chunk_size = std::min<int64_t>(levels.end() - it, kChunkSize);
+ hist1 += std::accumulate(levels.begin(), levels.begin() + chunk_size, int16_t{0});
+ it += chunk_size;
+ }
+ histogram[0] += num_levels - hist1;
+ histogram[1] += hist1;
+ return;
+ }
+
+ // The generic implementation issues a series of histogram load-stores.
+ // However, it limits store-to-load dependencies by interleaving partial histogram
+ // updates.
+ constexpr int kUnroll = 4;
+ std::array<std::vector<int64_t>, kUnroll> partial_hist;
+ for (auto& hist : partial_hist) {
+ hist.assign(histogram.size(), 0);
+ }
+ int64_t i = 0;
+ for (; i <= num_levels - kUnroll; i += kUnroll) {
+ for (int j = 0; j < kUnroll; ++j) {
+ ++partial_hist[j][levels[i + j]];
+ }
+ }
+ for (; i < num_levels; ++i) {
+ ++partial_hist[0][levels[i]];
+ }
+ for (const auto& hist : partial_hist) {
+ MergeLevelHistogram(histogram, hist);
+ }
+}
+
 } // namespace parquet

cpp/src/parquet/size_statistics.h

@@ -17,9 +17,12 @@
 
 #pragma once
 
+#include <cstdint>
+#include <iosfwd>
 #include <optional>
 #include <vector>
 
+#include "arrow/util/span.h"
 #include "parquet/platform.h"
 #include "parquet/type_fwd.h"
 
@@ -89,4 +92,11 @@ struct PARQUET_EXPORT SizeStatistics {
  static std::unique_ptr<SizeStatistics> Make(const ColumnDescriptor* descr);
 };
 
+PARQUET_EXPORT
+std::ostream& operator<<(std::ostream&, const SizeStatistics&);
+
+PARQUET_EXPORT
+void UpdateLevelHistogram(::arrow::util::span<const int16_t> levels,
+ ::arrow::util::span<int64_t> histogram);
+
 } // namespace parquet