GH-45847: [C++] Optimize Parquet column reader by fusing decoding and counting

This PR implements the optimization to fuse definition level decoding with counting in the Parquet column reader, addressing the TODO in column_reader.cc.

Author: tennisleng

cpp/src/arrow/util/rle_encoding_internal.h

@@ -308,6 +308,23 @@ class RleRunDecoder {
     return to_read;
   }
 
+  /// Get a batch of values and count how many equal match_value
+  [[nodiscard]] rle_size_t GetBatchWithCount(value_type* out, rle_size_t batch_size,
+                                             rle_size_t value_bit_width,
+                                             value_type match_value, int64_t* out_count) {
+    if (ARROW_PREDICT_FALSE(remaining_count_ == 0)) {
+      return 0;
+    }
+
+    const auto to_read = std::min(remaining_count_, batch_size);
+    std::fill(out, out + to_read, value_);
+    if (value_ == match_value) {
+      *out_count += to_read;
+    }
+    remaining_count_ -= to_read;
+    return to_read;
+  }
+
  private:
   value_type value_ = {};
   rle_size_t remaining_count_ = 0;
@@ -377,6 +394,15 @@ class BitPackedRunDecoder {
     return steps;
   }
 
+  /// Get a batch of values and count how many equal match_value
+  [[nodiscard]] rle_size_t GetBatchWithCount(value_type* out, rle_size_t batch_size,
+                                             rle_size_t value_bit_width,
+                                             value_type match_value, int64_t* out_count) {
+    auto steps = GetBatch(out, batch_size, value_bit_width);
+    *out_count += std::count(out, out + steps, match_value);
+    return steps;
+  }
+
  private:
   /// The pointer to the beginning of the run
   const uint8_t* data_ = nullptr;
@@ -438,6 +464,10 @@ class RleBitPackedDecoder {
   /// left or if an error occurred.
   [[nodiscard]] rle_size_t GetBatch(value_type* out, rle_size_t batch_size);
 
+  /// Get a batch of values and count how many equal match_value
+  [[nodiscard]] rle_size_t GetBatchWithCount(value_type* out, rle_size_t batch_size,
+                                             value_type match_value, int64_t* out_count);
+
   /// Like GetBatch but add spacing for null entries.
   ///
   /// Null entries will be set to an arbistrary value to avoid leaking private data.
@@ -483,6 +513,18 @@ class RleBitPackedDecoder {
         decoder_);
   }
 
+  /// Get a batch of values from the current run and return the number elements read.
+  [[nodiscard]] rle_size_t RunGetBatchWithCount(value_type* out, rle_size_t batch_size,
+                                                value_type match_value,
+                                                int64_t* out_count) {
+    return std::visit(
+        [&](auto& dec) {
+          return dec.GetBatchWithCount(out, batch_size, value_bit_width_, match_value,
+                                       out_count);
+        },
+        decoder_);
+  }
+
   /// Call the parser with a single callable for all event types.
   template <typename Callable>
   void ParseWithCallable(Callable&& func);
@@ -1474,4 +1516,49 @@ inline void RleBitPackedEncoder::Clear() {
   bit_writer_.Clear();
 }
 
+template <typename T>
+auto RleBitPackedDecoder<T>::GetBatchWithCount(value_type* out, rle_size_t batch_size,
+                                               value_type match_value, int64_t* out_count)
+    -> rle_size_t {
+  using ControlFlow = RleBitPackedParser::ControlFlow;
+
+  rle_size_t values_read = 0;
+
+  // Remaining from a previous call that would have left some unread data from a run.
+  if (ARROW_PREDICT_FALSE(run_remaining() > 0)) {
+    const auto read = RunGetBatchWithCount(out, batch_size, match_value, out_count);
+    values_read += read;
+    out += read;
+
+    // Either we fulfilled all the batch to be read or we finished remaining run.
+    if (ARROW_PREDICT_FALSE(values_read == batch_size)) {
+      return values_read;
+    }
+    ARROW_DCHECK(run_remaining() == 0);
+  }
+
+  ParseWithCallable([&](auto run) {
+    using RunDecoder = typename decltype(run)::template DecoderType<value_type>;
+
+    ARROW_DCHECK_LT(values_read, batch_size);
+    RunDecoder decoder(run, value_bit_width_);
+    const auto read =
+        decoder.GetBatchWithCount(out, batch_size - values_read, value_bit_width_,
+                                  match_value, out_count);
+    ARROW_DCHECK_LE(read, batch_size - values_read);
+    values_read += read;
+    out += read;
+
+    // Stop reading and store remaining decoder
+    if (ARROW_PREDICT_FALSE(values_read == batch_size || read == 0)) {
+      decoder_ = std::move(decoder);
+      return ControlFlow::Break;
+    }
+
+    return ControlFlow::Continue;
+  });
+
+  return values_read;
+}
+
 }  // namespace arrow::util

cpp/src/arrow/util/rle_encoding_test.cc

@@ -1280,4 +1280,48 @@ TEST(RleBitPacked, GetBatchSpacedRoundtripUint64) {
   DoTestGetBatchSpacedRoundtrip<uint64_t>();
 }
 
+TEST(Rle, GetBatchWithCount) {
+  const int bit_width = 1;
+  uint8_t buffer[100];
+  RleBitPackedEncoder encoder(buffer, sizeof(buffer), bit_width);
+
+  // 30 1s
+  for (int i = 0; i < 30; ++i) {
+    bool result = encoder.Put(1);
+    EXPECT_TRUE(result);
+  }
+  // 20 0s
+  for (int i = 0; i < 20; ++i) {
+    bool result = encoder.Put(0);
+    EXPECT_TRUE(result);
+  }
+  // 10 1s
+  for (int i = 0; i < 10; ++i) {
+    bool result = encoder.Put(1);
+    EXPECT_TRUE(result);
+  }
+  encoder.Flush();
+
+  RleBitPackedDecoder<int> decoder(buffer, sizeof(buffer), bit_width);
+  std::vector<int> values(60);
+  int64_t count = 0;
+
+  // Read first 40 values. Should be 30 1s and 10 0s.
+  // Count 1s.
+  int read = decoder.GetBatchWithCount(values.data(), 40, 1, &count);
+  EXPECT_EQ(read, 40);
+  EXPECT_EQ(count, 30);
+  for (int i = 0; i < 30; ++i) EXPECT_EQ(values[i], 1);
+  for (int i = 30; i < 40; ++i) EXPECT_EQ(values[i], 0);
+
+  // Read next 20 values. Should be 10 0s and 10 1s.
+  // Count 1s.
+  count = 0;
+  read = decoder.GetBatchWithCount(values.data(), 20, 1, &count);
+  EXPECT_EQ(read, 20);
+  EXPECT_EQ(count, 10);
+  for (int i = 0; i < 10; ++i) EXPECT_EQ(values[i], 0);
+  for (int i = 10; i < 20; ++i) EXPECT_EQ(values[i], 1);
+}
+
 }  // namespace arrow::util

cpp/src/parquet/column_reader.cc

@@ -191,6 +191,30 @@ int LevelDecoder::Decode(int batch_size, int16_t* levels) {
   return num_decoded;
 }
 
+int LevelDecoder::DecodeAndCount(int batch_size, int16_t* levels, int64_t* count) {
+  int num_decoded = 0;
+
+  int num_values = std::min(num_values_remaining_, batch_size);
+  if (encoding_ == Encoding::RLE) {
+    num_decoded =
+        rle_decoder_->GetBatchWithCount(levels, num_values, max_level_, count);
+  } else {
+    num_decoded = bit_packed_decoder_->GetBatch(bit_width_, levels, num_values);
+    *count += std::count(levels, levels + num_decoded, max_level_);
+  }
+  if (num_decoded > 0) {
+    internal::MinMax min_max = internal::FindMinMax(levels, num_decoded);
+    if (ARROW_PREDICT_FALSE(min_max.min < 0 || min_max.max > max_level_)) {
+      std::stringstream ss;
+      ss << "Malformed levels. min: " << min_max.min << " max: " << min_max.max
+         << " out of range.  Max Level: " << max_level_;
+      throw ParquetException(ss.str());
+    }
+  }
+  num_values_remaining_ -= num_decoded;
+  return num_decoded;
+}
+
 ReaderProperties default_reader_properties() {
   static ReaderProperties default_reader_properties;
   return default_reader_properties;
@@ -1006,14 +1030,11 @@ class TypedColumnReaderImpl : public TypedColumnReader<DType>,
 
     // If the field is required and non-repeated, there are no definition levels
     if (this->max_def_level_ > 0 && def_levels != nullptr) {
-      *num_def_levels = this->ReadDefinitionLevels(batch_size, def_levels);
-      if (ARROW_PREDICT_FALSE(*num_def_levels != batch_size)) {
+      if (ARROW_PREDICT_FALSE(this->definition_level_decoder_.DecodeAndCount(
+              batch_size, def_levels, non_null_values_to_read) != batch_size)) {
         throw ParquetException(kErrorRepDefLevelNotMatchesNumValues);
       }
-      // TODO(wesm): this tallying of values-to-decode can be performed with better
-      // cache-efficiency if fused with the level decoding.
-      *non_null_values_to_read +=
-          std::count(def_levels, def_levels + *num_def_levels, this->max_def_level_);
+      *num_def_levels = batch_size;
     } else {
       // Required field, read all values
       if (num_def_levels != nullptr) {

cpp/src/parquet/column_reader.h

@@ -92,6 +92,10 @@ class PARQUET_EXPORT LevelDecoder {
   // Decodes a batch of levels into an array and returns the number of levels decoded
   int Decode(int batch_size, int16_t* levels);
 
+  /// Decodes a batch of levels into an array and counts the number of levels equal to
+  /// max_level_
+  int DecodeAndCount(int batch_size, int16_t* levels, int64_t* count);
+
  private:
   int bit_width_;
   int num_values_remaining_;