GH-48561: [C++][Parquet] Optimize column reader by fusing definition level decoding with counting

cpp/src/arrow/util/rle_encoding_internal.h

@@ -341,6 +341,19 @@ class RleRunDecoder {
     return to_read;
   }
 
+  /// Decode a batch of values and count the number of occurrences of a value.
+  ///
+  /// The count is limited to at most the next `batch_size` items.
+  /// @return The matching value count and number of elements that were decoded.
+  RleCountUpToResult GetBatchAndCount(value_type* out,
+                                      const RleCountUpToParams<value_type>& p) {
+    const auto read = GetBatch(out, p.batch_size, p.value_bit_width);
+    return {
+        .matching_count = read * (p.value == value_),
+        .processed_count = read,
+    };
+  }
+
  private:
   value_type value_ = {};
   rle_size_t remaining_count_ = 0;
@@ -443,6 +456,22 @@ class BitPackedRunDecoder {
     return opts.batch_size;
   }
 
+  /// Decode a batch of values and count the number of occurrences of a value.
+  ///
+  /// The count is limited to at most the next `batch_size` items.
+  /// @return The matching value count and number of elements that were decoded.
+  RleCountUpToResult GetBatchAndCount(value_type* out,
+                                      const RleCountUpToParams<value_type>& p) {
+    const auto read = GetBatch(out, p.batch_size, p.value_bit_width);
+    const auto matching_count =
+        static_cast<rle_size_t>(std::count(out, out + read, p.value));
+
+    return {
+      .matching_count = matching_count,
+      .processed_count = read,
+    };
+  }
+
  private:
   /// The pointer to the beginning of the run
   const uint8_t* data_ = nullptr;
@@ -516,6 +545,14 @@ class RleBitPackedDecoder {
   /// left or if an error occurred.
   [[nodiscard]] rle_size_t GetBatch(value_type* out, rle_size_t batch_size);
 
+  /// Decode a batch of values and count the number of occurrences of a value.
+  ///
+  /// May decode fewer elements than requested if there are not enough values left
+  /// or if an error occurred.
+  /// @return The matching value count and number of elements that were decoded.
+  RleCountUpToResult GetBatchAndCount(value_type* out, value_type value,
+                                      rle_size_t batch_size);
+
   /// Like GetBatch but add spacing for null entries.
   ///
   /// Null entries will be set to an arbistrary value to avoid leaking private data.
@@ -644,6 +681,21 @@ class BitPackedDecoder : private BitPackedRunDecoder<T> {
     return Base::GetBatch(out, batch_size, value_bit_width_);
   }
 
+  /// Decode a batch of values and count the number of occurrences of a value.
+  ///
+  /// The count is limited to at most the next `batch_size` items.
+  /// @return The matching value count and number of elements that were decoded.
+  RleCountUpToResult GetBatchAndCount(value_type* out, value_type value,
+                                      rle_size_t batch_size) {
+    return Base::GetBatchAndCount(
+        out,
+        {
+          .value = value,
+          .batch_size = batch_size,
+          .value_bit_width = value_bit_width_,
+        });
+  }
+
  private:
   rle_size_t value_bit_width_ = {};
 };
@@ -1005,6 +1057,35 @@ auto RleBitPackedDecoder<T>::GetBatch(value_type* out,
       batch_size);
 }
 
+template <typename T>
+RleCountUpToResult RleBitPackedDecoder<T>::GetBatchAndCount(
+    value_type* out, value_type value, rle_size_t batch_size) {
+  rle_size_t matching_count = 0;
+
+  const rle_size_t processed_count = ProcessValues(
+      [&out, value, this, &matching_count](auto& decoder,
+                                           rle_size_t run_batch_size) {
+        const auto result = decoder.GetBatchAndCount(
+            out,
+            {
+                .value = value,
+                .batch_size = run_batch_size,
+                .value_bit_width = value_bit_width_,
+            });
+
+        out += result.processed_count;
+        matching_count += result.matching_count;
+
+        return result.processed_count;
+      },
+      batch_size);
+
+  return {
+      .matching_count = matching_count,
+      .processed_count = processed_count,
+  };
+}
+
 namespace internal {
 
 /// Utility class to safely handle values and null count without too error-prone

cpp/src/arrow/util/rle_encoding_test.cc

@@ -434,6 +434,27 @@ TEST(BitPacked, BitPackedDecoder) {
       /* expected= */ {1000, 1, 2, 3, 4, 5, 6, 7});
 }
 
+TEST(BitPacked, BitPackedDecoderGetBatchAndCount) {
+  const std::array<uint8_t, 3> bytes = {0x88, 0xc6, 0xfa};
+  const std::vector<uint16_t> expected = {0, 1, 2, 3, 4, 5, 6, 7};
+
+  BitPackedDecoder<uint16_t> decoder(
+      bytes.data(),
+      static_cast<rle_size_t>(bytes.size()),
+      /* value_bit_width= */ 3,
+      /* value_count= */ 8);
+
+  std::vector<uint16_t> decoded(expected.size());
+
+  const auto result =
+      decoder.GetBatchAndCount(decoded.data(), /* value= */ 3,
+                               /* batch_size= */ static_cast<rle_size_t>(decoded.size()));
+
+  EXPECT_EQ(result.processed_count, 8);
+  EXPECT_EQ(result.matching_count, 1);
+  EXPECT_EQ(decoded, expected);
+}
+
 template <typename T>
 void TestRleBitPackedParser(std::vector<uint8_t> bytes, rle_size_t bit_width,
                             std::vector<T> expected) {
@@ -1250,6 +1271,71 @@ void CheckCountUpTo(const Array& data, int bit_width, typename Type::c_type valu
   EXPECT_LT(res.processed_count, 8);
 }
 
+/// Check RleBitPackedDecoder::GetBatchAndCount, which spans multiple runs
+/// through the parser.
+///
+/// The decoded values are compared against the original data and the matching
+/// counts are compared against a naive count over the original data.
+template <typename Type>
+void CheckGetBatchAndCount(const Array& data, int bit_width,
+                           typename Type::c_type value) {
+  using ArrayType = typename TypeTraits<Type>::ArrayType;
+  using value_type = typename Type::c_type;
+
+  const auto data_size = static_cast<rle_size_t>(data.length());
+  const value_type* data_values =
+      arrow::internal::checked_cast<const ArrayType&>(data).raw_values();
+
+  ARROW_SCOPED_TRACE("bit_width = ", bit_width, ", data_size = ", data_size,
+                     ", value = ", value);
+
+  // Encode all values into `buffer`.
+  const auto buffer = EncodeTestArray<Type>(data, bit_width);
+
+  RleBitPackedDecoder<value_type> decoder(buffer.data(),
+                                           static_cast<int>(buffer.size()),
+                                           bit_width);
+
+  // Decode in chunks small enough to repeatedly cross run boundaries.
+  const rle_size_t step = std::max<rle_size_t>(data_size / 16, 1);
+  std::vector<value_type> decoded(step);
+
+  rle_size_t pos = 0;
+  rle_size_t total_matching = 0;
+
+  while (pos < data_size) {
+    const auto to_process = std::min(step, data_size - pos);
+
+    const auto result =
+        decoder.GetBatchAndCount(decoded.data(), value, to_process);
+
+    ASSERT_EQ(result.processed_count, to_process);
+
+    // The decoded output must exactly match the original input.
+    for (rle_size_t i = 0; i < result.processed_count; ++i) {
+      EXPECT_EQ(decoded[i], data_values[pos + i])
+          << "at position " << (pos + i);
+    }
+
+    // The matching count must equal a naive count over the same input range.
+    const auto expected =
+        std::count(data_values + pos, data_values + pos + to_process, value);
+
+    EXPECT_EQ(result.matching_count, static_cast<rle_size_t>(expected))
+        << "at position " << pos;
+
+    pos += result.processed_count;
+    total_matching += result.matching_count;
+  }
+
+  EXPECT_EQ(pos, data_size) << "Total number of values processed is off";
+
+  const auto total_expected =
+      std::count(data_values, data_values + data_size, value);
+
+  EXPECT_EQ(total_matching, static_cast<rle_size_t>(total_expected));
+}
+
 template <typename T>
 struct DataTestRleBitPackedRandomPart {
   using value_type = T;
@@ -1419,6 +1505,12 @@ void DoTestGetBatchSpacedRoundtrip() {
     CheckCountUpTo<ArrowType>(*array, case_.bit_width, max_value);
     CheckCountUpTo<ArrowType>(*array->Slice(1), case_.bit_width, first);
 
+    // Tests for GetBatchAndCount with both a value present in the data and a value
+    // that may not be present.
+    CheckGetBatchAndCount<ArrowType>(*array, case_.bit_width, first);
+    CheckGetBatchAndCount<ArrowType>(*array, case_.bit_width, max_value);
+    CheckGetBatchAndCount<ArrowType>(*array->Slice(1), case_.bit_width, first);
+
     // Tests for GetBatchSpaced
     CheckRoundTrip<ArrowType>(*array, case_.bit_width, /* spaced= */ true,
                               /* parts= */ 1);

cpp/src/parquet/column_reader.cc

@@ -110,6 +110,12 @@ struct LevelDecoder::Impl {
     return std::visit([&](auto& dec) { return dec.GetBatch(out, batch_size); }, decoder);
   }
 
+  auto GetBatchAndCount(int16_t* out, int16_t value, int batch_size) {
+    return std::visit(
+        [&](auto& dec) { return dec.GetBatchAndCount(out, value, batch_size); },
+        decoder);
+  }
+
   [[nodiscard]] int Advance(int batch_size) {
     return std::visit([&](auto& dec) { return dec.Advance(batch_size); }, decoder);
   }
@@ -203,6 +209,30 @@ int LevelDecoder::Decode(int batch_size, int16_t* levels) {
   return num_decoded;
 }
 
+auto LevelDecoder::DecodeAndCount(int batch_size, int16_t* levels, int16_t value)
+    -> CountUpToResult {
+  const int num_values = std::min(num_values_remaining_, batch_size);
+  const auto result = impl_->GetBatchAndCount(levels, value, num_values);
+  const int num_decoded = result.processed_count;
+
+  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 {
+      .matching_count = result.matching_count,
+      .processed_count = num_decoded,
+  };
+}
+
 int LevelDecoder::Skip(int batch_size) {
   const int num_values = std::min(num_values_remaining_, batch_size);
   const int num_advanced = impl_->Advance(num_values);
@@ -773,6 +803,18 @@ class ColumnReaderImplBase {
     return definition_level_decoder_.Decode(static_cast<int>(batch_size), levels);
   }
 
+  // Read multiple definition levels into preallocated memory and count the
+  // number of physical values.
+  LevelDecoder::CountUpToResult ReadDefinitionLevelsAndCount(int64_t batch_size,
+                                                             int16_t* levels) {
+    if (max_def_level() == 0) {
+      return {};
+    }
+
+    return definition_level_decoder_.DecodeAndCount(
+        static_cast<int>(batch_size), levels, max_def_level());
+  }
+
   bool HasNextInternal() {
     // Either there is no data page available yet, or the data page has been
     // exhausted
@@ -1113,14 +1155,15 @@ 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);
+      const auto result = this->ReadDefinitionLevelsAndCount(batch_size, def_levels);
+
+      *num_def_levels = result.processed_count;
+
       if (ARROW_PREDICT_FALSE(*num_def_levels != 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());
+
+      *non_null_values_to_read += result.matching_count;
     } else {
       // Required field, read all values
       if (num_def_levels != nullptr) {

cpp/src/parquet/column_reader.h

@@ -92,6 +92,13 @@ class PARQUET_EXPORT LevelDecoder {
     int processed_count;
   };
 
+  /// Decode a batch of levels into an array and count the number of occurrences
+  /// of `value`.
+  ///
+  /// The count is limited to at most the next `batch_size` items.
+  /// @return The matching value count and number of elements that were decoded.
+  CountUpToResult DecodeAndCount(int batch_size, int16_t* levels, int16_t value);
+
   /// Advance and count the number of occurrences of `value`.
   ///
   /// The count is limited to at most the next `batch_size` items.