GH-55399: [C++][Parquet] Minor optimize level counting with CounfIf16

cpp/src/parquet/column_reader.cc

@@ -47,6 +47,7 @@
 #include "parquet/encryption/encryption_internal.h"
 #include "parquet/encryption/internal_file_decryptor.h"
 #include "parquet/exception.h"
+#include "parquet/level_comparison.h"
 #include "parquet/level_conversion.h"
 #include "parquet/level_decoder_internal.h"
 #include "parquet/properties.h"
@@ -1546,8 +1547,10 @@ class TypedColumnReaderImpl
       // TODO(wesm): this tallying of values-to-decode can be performed with better
       // cache-efficiency if fused with the level decoding.
       ARROW_DCHECK_LE(*num_def_levels, std::numeric_limits<int32_t>::max());
-      *non_null_values_to_read += static_cast<int32_t>(
-          std::count(def_levels, def_levels + *num_def_levels, this->max_def_level()));
+      const int16_t max_def_level = this->max_def_level();
+      *non_null_values_to_read += static_cast<int32_t>(internal::CountIf16(
+          def_levels, *num_def_levels,
+          [max_def_level](int16_t level) { return level == max_def_level; }));
     } else {
       // Required field, read all values
       if (num_def_levels != nullptr) {
@@ -1961,9 +1964,10 @@ int64_t TypedRecordReader<DT, VS>::SkipRecordsInBufferNonRepeated(int64_t num_re
   // We skipped the levels by incrementing 'levels_position_'. For values
   // we do not have a buffer, so we need to read them and throw them away.
   // First we need to figure out how many present/not-null values there are.
-  const auto values_to_read = static_cast<int32_t>(
-      std::count(def_levels() + levels_position_, def_levels() + remaining_levels_pos,
-                 this->max_def_level()));
+  const int16_t max_def_level = this->max_def_level();
+  const auto values_to_read = static_cast<int32_t>(internal::CountIf16(
+      def_levels() + levels_position_, remaining_levels_pos - levels_position_,
+      [max_def_level](int16_t level) { return level == max_def_level; }));
 
   // Now that we have figured out number of values to read, we do not need
   // these levels anymore. We will remove these values from the buffer.
@@ -2128,9 +2132,8 @@ int64_t TypedRecordReader<DT, VS>::DelimitRecords(int64_t num_records,
     const int64_t stride =
         std::min(levels_written() - levels_position_, num_records - records_read);
     const int64_t position_end = levels_position_ + stride;
-    for (int64_t i = levels_position_; i < position_end; ++i) {
-      records_read += rep_levels[i] == 0;
-    }
+    records_read += CountIf16(rep_levels + levels_position_, stride,
+                              [](int16_t rep_level) { return rep_level == 0; });
     levels_position_ = position_end;
     if (records_read == num_records) {
       // Check last rep_level reaches the boundary and
@@ -2144,8 +2147,10 @@ int64_t TypedRecordReader<DT, VS>::DelimitRecords(int64_t num_records,
     }
   }
   // Scan definition levels to find number of physical values
-  *values_seen = std::count(def_levels + level, def_levels + levels_position_,
-                            this->max_def_level());
+  const int16_t max_def_level = this->max_def_level();
+  *values_seen = CountIf16(
+      def_levels + level, levels_position_ - level,
+      [max_def_level](int16_t def_level) { return def_level == max_def_level; });
   return records_read;
 }
 
@@ -2234,8 +2239,10 @@ void TypedRecordReader<DT, VS>::ReadDenseForOptionalInBuffer(
 
   // When reading dense we need to figure out number of values to read.
   const int16_t* def_levels = this->def_levels();
-  *values_to_read += std::count(def_levels + start_levels_position,
-                                def_levels + levels_position_, this->max_def_level());
+  const int16_t max_def_level = this->max_def_level();
+  *values_to_read += CountIf16(
+      def_levels + start_levels_position, levels_position_ - start_levels_position,
+      [max_def_level](int16_t def_level) { return def_level == max_def_level; });
   // This is only reading in the current page so this fits in an int32.
   value_sink_.ReadValuesDense(*this->current_decoder_.get(),
                               clamp_to<int32_t>(*values_to_read));

cpp/src/parquet/column_writer.cc

@@ -55,6 +55,7 @@
 #include "parquet/encoding.h"
 #include "parquet/encryption/encryption_internal.h"
 #include "parquet/encryption/internal_file_encryptor.h"
+#include "parquet/level_comparison.h"
 #include "parquet/level_conversion.h"
 #include "parquet/metadata.h"
 #include "parquet/page_index.h"
@@ -1698,11 +1699,10 @@ class TypedColumnWriterImpl : public ColumnWriterImpl,
     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_levels; ++i) {
-        if (def_levels[i] == descr_->max_definition_level()) {
-          ++values_to_write;
-        }
-      }
+      const int16_t max_def_level = descr_->max_definition_level();
+      values_to_write = internal::CountIf16(
+          def_levels, num_levels,
+          [max_def_level](int16_t def_level) { return def_level == max_def_level; });
 
       WriteDefinitionLevels(num_levels, def_levels);
     } else {
@@ -1714,12 +1714,10 @@ class TypedColumnWriterImpl : public ColumnWriterImpl,
     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_levels; ++i) {
-        if (rep_levels[i] == 0) {
-          rows_written_++;
-          num_buffered_rows_++;
-        }
-      }
+      const int64_t records_written = internal::CountIf16(
+          rep_levels, num_levels, [](int16_t rep_level) { return rep_level == 0; });
+      rows_written_ += records_written;
+      num_buffered_rows_ += records_written;
 
       WriteRepetitionLevels(num_levels, rep_levels);
     } else {
@@ -1749,11 +1747,16 @@ class TypedColumnWriterImpl : public ColumnWriterImpl,
         *out_spaced_values_to_write = batch_size;
         *null_count = 0;
       } else {
-        for (int x = 0; x < batch_size; x++) {
-          *out_values_to_write += def_levels[x] == level_info_.def_level ? 1 : 0;
-          *out_spaced_values_to_write +=
-              def_levels[x] >= level_info_.repeated_ancestor_def_level ? 1 : 0;
-        }
+        const int16_t def_level = level_info_.def_level;
+        const int16_t repeated_ancestor_def_level =
+            level_info_.repeated_ancestor_def_level;
+        *out_values_to_write += internal::CountIf16(
+            def_levels, batch_size,
+            [def_level](int16_t level) { return level == def_level; });
+        *out_spaced_values_to_write += internal::CountIf16(
+            def_levels, batch_size, [repeated_ancestor_def_level](int16_t level) {
+              return level >= repeated_ancestor_def_level;
+            });
         *null_count = batch_size - *out_values_to_write;
       }
       return;
@@ -1809,12 +1812,10 @@ class TypedColumnWriterImpl : public ColumnWriterImpl,
     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_levels; ++i) {
-        if (rep_levels[i] == 0) {
-          rows_written_++;
-          num_buffered_rows_++;
-        }
-      }
+      const int64_t records_written = internal::CountIf16(
+          rep_levels, num_levels, [](int16_t rep_level) { return rep_level == 0; });
+      rows_written_ += records_written;
+      num_buffered_rows_ += records_written;
       WriteRepetitionLevels(num_levels, rep_levels);
     } else {
       // Each value is exactly one row

cpp/src/parquet/file_reader.cc

@@ -44,6 +44,7 @@
 #include "parquet/encryption/internal_file_decryptor.h"
 #include "parquet/exception.h"
 #include "parquet/file_writer.h"
+#include "parquet/level_comparison.h"
 #include "parquet/metadata.h"
 #include "parquet/page_index.h"
 #include "parquet/platform.h"
@@ -977,11 +978,9 @@ int64_t ScanFileContents(std::vector<int> columns, const int32_t column_batch_si
             ScanAllValues(column_batch_size, def_levels.data(), rep_levels.data(),
                           values.data(), &values_read, col_reader.get());
         if (col_reader->descr()->max_repetition_level() > 0) {
-          for (size_t i = 0; i < static_cast<size_t>(levels_read); i++) {
-            if (rep_levels[i] == 0) {
-              total_rows[col]++;
-            }
-          }
+          total_rows[col] +=
+              internal::CountIf16(rep_levels.data(), levels_read,
+                                  [](int16_t rep_level) { return rep_level == 0; });
         } else {
           total_rows[col] += levels_read;
         }

cpp/src/parquet/level_comparison.h

@@ -18,6 +18,7 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <limits>
 
 #include "parquet/platform.h"
 
@@ -35,4 +36,78 @@ struct MinMax {
 
 MinMax PARQUET_EXPORT FindMinMax(const int16_t* levels, int64_t num_levels);
 
+/// Counts the levels that satisfy `pred`, in chunks small enough that a
+/// `uint16_t` accumulator cannot overflow.
+///
+/// This is the fallback for inputs that do not fit in a `uint16_t` accumulator.
+/// It is kept out of line so that the common path (a single Parquet page chunk,
+/// always short enough for one accumulator) is a plain loop without any chunk
+/// bookkeeping in front of it.
+template <typename Predicate>
+ARROW_NOINLINE int64_t CountIf16Chunked(const int16_t* levels, int64_t num_levels,
+                                        Predicate pred) {
+  // The loop below accumulates one match per lane in a `uint16_t`, so a chunk
+  // must not hold more levels than that can count.
+  constexpr int64_t kMaxCount = std::numeric_limits<uint16_t>::max();
+  int64_t total = 0;
+  while (num_levels > 0) {
+    const int64_t chunk_length = std::min(num_levels, kMaxCount);
+    uint16_t chunk_count = 0;
+    for (int64_t i = 0; i < chunk_length; ++i) {
+      chunk_count += pred(levels[i]) ? 1 : 0;
+    }
+    total += chunk_count;
+    levels += chunk_length;
+    num_levels -= chunk_length;
+  }
+  return total;
+}
+
+/// Count the levels that satisfy the given predicate.
+///
+/// The matches are accumulated in a `uint16_t`, which cannot overflow as long
+/// as at most `std::numeric_limits<uint16_t>::max()` levels are counted at a
+/// time, so longer inputs are delegated to `CountIf16Chunked`. This is both
+/// safe and much faster than `std::count`, as the loop can be vectorized with
+/// 16-bit accumulators (one accumulated match per lane) instead of widening
+/// every match to a 64-bit integer.
+///
+/// \param[in] levels The level values to inspect.
+/// \param[in] num_levels The number of levels to inspect, must be non-negative.
+/// \param[in] pred The predicate to apply, must be callable as
+///   `bool(int16_t)`.
+/// \return The number of levels for which `pred` returns true.
+template <typename Predicate>
+inline int64_t CountIf16(const int16_t* levels, int64_t num_levels, Predicate pred) {
+  // Short inputs are counted first, before the wide loop and before the chunk
+  // clamp, because a handful of levels cannot amortize either. Running the
+  // small case first is also how folly::simd::detail::simdForEach orders its
+  // single steps before the unrolled ones, so that short inputs never pay for
+  // the unrolled set up.
+  //
+  // 32 is the widest per-iteration body this reduction gets: a packed int16_t
+  // register holds 8 levels with SSE2/NEON and 16 with AVX2, but clang runs up
+  // to 32 levels per iteration on NEON (four registers), so shorter inputs
+  // never reach a full iteration on the narrow targets. The bound also has to
+  // be a compile-time power of two, which is what turns the loop into an
+  // unrolled ladder of 8/16/24-wide compares instead of scalar code.
+  constexpr int64_t kSmallLevels = 32;
+  constexpr int64_t kMaxCount = std::numeric_limits<uint16_t>::max();
+  if (num_levels < kSmallLevels) {
+    uint16_t count = 0;
+    for (int64_t i = 0; i < num_levels; ++i) {
+      count += pred(levels[i]) ? 1 : 0;
+    }
+    return count;
+  }
+  if (num_levels > kMaxCount) {
+    return CountIf16Chunked(levels, num_levels, pred);
+  }
+  uint16_t count = 0;
+  for (int64_t i = 0; i < num_levels; ++i) {
+    count += pred(levels[i]) ? 1 : 0;
+  }
+  return count;
+}
+
 }  // namespace parquet::internal

cpp/src/parquet/level_conversion_test.cc

@@ -23,6 +23,7 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
+#include <limits>
 #include <string>
 #include <vector>
 
@@ -113,6 +114,53 @@ TEST(GreaterThanBitmap, GeneratesExpectedBitmasks) {
 }
 #endif
 
+TEST(CountIf16, CountsMatchingLevels) {
+  const std::vector<int16_t> levels = {3, 1, 3, 0, 3, 2};
+
+  EXPECT_EQ(CountIf16(levels.data(), /*num_levels=*/0,
+                      [](int16_t level) { return level == 3; }),
+            0);
+  EXPECT_EQ(CountIf16(levels.data(), /*num_levels=*/6,
+                      [](int16_t level) { return level == 3; }),
+            3);
+  EXPECT_EQ(CountIf16(levels.data(), /*num_levels=*/6,
+                      [](int16_t level) { return level == 0; }),
+            1);
+  EXPECT_EQ(
+      CountIf16(levels.data(), /*num_levels=*/6, [](int16_t level) { return level > 0; }),
+      5);
+  EXPECT_EQ(CountIf16(levels.data(), /*num_levels=*/6, [](int16_t) { return false; }), 0);
+  // Cover the single chunk fast path, a partial vector iteration and a plain
+  // scalar loop by cross checking every prefix length of a longer input.
+  std::vector<int16_t> longer_levels(/*count=*/40, 0);
+  for (size_t i = 0; i < longer_levels.size(); i += 3) {
+    longer_levels[i] = 3;
+  }
+  for (int64_t n = 0; n <= static_cast<int64_t>(longer_levels.size()); ++n) {
+    int64_t expected = 0;
+    for (int64_t i = 0; i < n; ++i) {
+      expected += longer_levels[i] == 3 ? 1 : 0;
+    }
+    EXPECT_EQ(
+        CountIf16(longer_levels.data(), n, [](int16_t level) { return level == 3; }),
+        expected)
+        << "num_levels: " << n;
+  }
+}
+
+TEST(CountIf16, HandlesChunkBoundaries) {
+  // A chunk must not overflow its 16-bit accumulator, and chunk boundaries must
+  // neither lose nor double count levels.
+  constexpr int64_t kChunkSize = std::numeric_limits<uint16_t>::max();
+  std::vector<int16_t> levels(/*count=*/5 * kChunkSize, 1);
+  const std::vector<int64_t> num_levels = {kChunkSize - 1, kChunkSize, kChunkSize + 1,
+                                           3 * kChunkSize, 5 * kChunkSize};
+  for (const int64_t n : num_levels) {
+    EXPECT_EQ(CountIf16(levels.data(), n, [](int16_t level) { return level == 1; }), n)
+        << "num_levels: " << n;
+  }
+}
+
 TEST(DefLevelsToBitmap, WithRepetitionLevelFiltersOutEmptyListValues) {
   std::vector<uint8_t> validity_bitmap(/*count*/ 8, 0);
 

cpp/src/parquet/size_statistics.cc

@@ -19,14 +19,14 @@
 
 #include <algorithm>
 #include <array>
-#include <numeric>
 #include <ostream>
 #include <span>
 #include <string_view>
 #include <vector>
 
 #include "arrow/util/logging_internal.h"
 #include "parquet/exception.h"
+#include "parquet/level_comparison.h"
 #include "parquet/schema.h"
 
 namespace parquet {
@@ -160,17 +160,11 @@ void UpdateLevelHistogram(std::span<const int16_t> levels, std::span<int64_t> hi
 
   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(it, it + chunk_size, int16_t{0});
-      it += chunk_size;
-    }
+    // Counting the 1s directly gives us the number of 1s, and the number of 0s
+    // follows. CountIf16 accumulates in the int16_t space in chunks, which the
+    // compiler vectorizes efficiently and which cannot overflow.
+    const int64_t hist1 = internal::CountIf16(levels.data(), num_levels,
+                                              [](int16_t level) { return level == 1; });
     histogram[0] += num_levels - hist1;
     histogram[1] += hist1;
     return;