From adf33446f8569cf3a2cd7f76abaabf5389a20470 Mon Sep 17 00:00:00 2001 From: shockp Date: Mon, 31 Aug 2026 17:48:03 +0000 Subject: [PATCH 1/3] GH-48561: [C++] Add decode-and-count support to RLE decoders Add GetBatchAndCount support to RleRunDecoder, BitPackedRunDecoder, RleBitPackedDecoder, and BitPackedDecoder. The new operation decodes values while also counting occurrences of a target value. RLE runs compute the count directly from the run value and length, avoiding an additional scan of the decoded output. Bit-packed runs count the values after unpacking them into the output buffer. Add tests covering decoded output, matching counts, partial batches, and mixed RLE/bit-packed runs. --- cpp/src/arrow/util/rle_encoding_internal.h | 81 +++++++++++++++++++ cpp/src/arrow/util/rle_encoding_test.cc | 92 ++++++++++++++++++++++ 2 files changed, 173 insertions(+) diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index 4bd08fd65306..21341694662b 100644 --- a/cpp/src/arrow/util/rle_encoding_internal.h +++ b/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& 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& p) { + const auto read = GetBatch(out, p.batch_size, p.value_bit_width); + const auto matching_count = + static_cast(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 { 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::GetBatch(value_type* out, batch_size); } +template +RleCountUpToResult RleBitPackedDecoder::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 diff --git a/cpp/src/arrow/util/rle_encoding_test.cc b/cpp/src/arrow/util/rle_encoding_test.cc index 16b702c23d49..91b7bcc764d9 100644 --- a/cpp/src/arrow/util/rle_encoding_test.cc +++ b/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 bytes = {0x88, 0xc6, 0xfa}; + const std::vector expected = {0, 1, 2, 3, 4, 5, 6, 7}; + + BitPackedDecoder decoder( + bytes.data(), + static_cast(bytes.size()), + /* value_bit_width= */ 3, + /* value_count= */ 8); + + std::vector decoded(expected.size()); + + const auto result = + decoder.GetBatchAndCount(decoded.data(), /* value= */ 3, + /* batch_size= */ static_cast(decoded.size())); + + EXPECT_EQ(result.processed_count, 8); + EXPECT_EQ(result.matching_count, 1); + EXPECT_EQ(decoded, expected); +} + template void TestRleBitPackedParser(std::vector bytes, rle_size_t bit_width, std::vector 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 +void CheckGetBatchAndCount(const Array& data, int bit_width, + typename Type::c_type value) { + using ArrayType = typename TypeTraits::ArrayType; + using value_type = typename Type::c_type; + + const auto data_size = static_cast(data.length()); + const value_type* data_values = + arrow::internal::checked_cast(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(data, bit_width); + + RleBitPackedDecoder decoder(buffer.data(), + static_cast(buffer.size()), + bit_width); + + // Decode in chunks small enough to repeatedly cross run boundaries. + const rle_size_t step = std::max(data_size / 16, 1); + std::vector 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(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(total_expected)); +} + template struct DataTestRleBitPackedRandomPart { using value_type = T; @@ -1419,6 +1505,12 @@ void DoTestGetBatchSpacedRoundtrip() { CheckCountUpTo(*array, case_.bit_width, max_value); CheckCountUpTo(*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(*array, case_.bit_width, first); + CheckGetBatchAndCount(*array, case_.bit_width, max_value); + CheckGetBatchAndCount(*array->Slice(1), case_.bit_width, first); + // Tests for GetBatchSpaced CheckRoundTrip(*array, case_.bit_width, /* spaced= */ true, /* parts= */ 1); From 21dd44b03e0ccf37a71f697745411aefbe847fd8 Mon Sep 17 00:00:00 2001 From: shockp Date: Mon, 31 Aug 2026 18:36:43 +0000 Subject: [PATCH 2/3] GH-48561: [C++][Parquet] Fuse definition level decoding and counting Use the new decode-and-count operation when reading definition levels. Add LevelDecoder::DecodeAndCount to decode levels while counting occurrences of the maximum definition level, preserving the existing level validation. Update the column reader to use the matching count directly when determining the number of physical values to decode, removing the separate std::count pass over the definition level buffer. --- cpp/src/parquet/column_reader.cc | 53 +++++++++++++++++++++++++++++--- cpp/src/parquet/column_reader.h | 7 +++++ 2 files changed, 55 insertions(+), 5 deletions(-) diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 56c65fd233b1..ff80457707f3 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/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(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(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, // 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) { diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index 07064a9cf37d..34df9837d356 100644 --- a/cpp/src/parquet/column_reader.h +++ b/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. From 9e4a87139427b7f9171621db7addec013bc4fbc3 Mon Sep 17 00:00:00 2001 From: shockp Date: Mon, 31 Aug 2026 19:49:27 +0000 Subject: [PATCH 3/3] GH-48561: [C++][Parquet] Add decode-and-count benchmarks Add benchmarks comparing separate definition level decoding and counting with the fused DecodeAndCount path. Cover RLE and bit-packed encodings across different maximum levels, batch sizes, and level repeat counts. --- cpp/src/parquet/column_reader_benchmark.cc | 110 +++++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/cpp/src/parquet/column_reader_benchmark.cc b/cpp/src/parquet/column_reader_benchmark.cc index 83f661bf9ad6..5de96888430c 100644 --- a/cpp/src/parquet/column_reader_benchmark.cc +++ b/cpp/src/parquet/column_reader_benchmark.cc @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +#include #include #include "benchmark/benchmark.h" #include "parquet/column_page.h" @@ -336,6 +337,57 @@ static void DecodeLevels(Encoding::type level_encoding, int16_t max_level, int n state.SetItemsProcessed(state.iterations() * num_levels); } +template +static void DecodeAndCountLevels(Encoding::type level_encoding, int16_t max_level, + int num_levels, int batch_size, int level_repeat_count, + ::benchmark::State& state) { + std::vector bytes; + { + std::vector input_levels; + GenerateLevels(/*level_repeats=*/level_repeat_count, + /*max_repeat_factor=*/max_level, num_levels, &input_levels); + EncodeLevels(level_encoding, max_level, num_levels, input_levels.data(), &bytes); + } + + LevelDecoder decoder; + std::vector output_levels(batch_size); + + for (auto _ : state) { + state.PauseTiming(); + decoder.SetData(level_encoding, max_level, num_levels, bytes.data(), + static_cast(bytes.size())); + int64_t matching_count = 0; + state.ResumeTiming(); + + while (true) { + if constexpr (Fused) { + const auto result = + decoder.DecodeAndCount(batch_size, output_levels.data(), max_level); + + if (result.processed_count == 0) { + break; + } + + matching_count += result.matching_count; + } else { + const int levels_decoded = decoder.Decode(batch_size, output_levels.data()); + + if (levels_decoded == 0) { + break; + } + + matching_count += std::count(output_levels.data(), + output_levels.data() + levels_decoded, max_level); + } + } + + DoNotOptimize(matching_count); + } + + state.SetBytesProcessed(state.iterations() * num_levels * sizeof(int16_t)); + state.SetItemsProcessed(state.iterations() * num_levels); +} + static void ReadLevels_Rle(::benchmark::State& state) { int16_t max_level = static_cast(state.range(0)); int num_levels = static_cast(state.range(1)); @@ -354,6 +406,46 @@ static void ReadLevels_BitPack(::benchmark::State& state) { level_repeat_count, state); } +static void ReadLevelsAndCount_Rle_Separate(::benchmark::State& state) { + const int16_t max_level = static_cast(state.range(0)); + const int num_levels = static_cast(state.range(1)); + const int batch_size = static_cast(state.range(2)); + const int level_repeat_count = static_cast(state.range(3)); + + DecodeAndCountLevels(Encoding::RLE, max_level, num_levels, batch_size, + level_repeat_count, state); +} + +static void ReadLevelsAndCount_Rle_Fused(::benchmark::State& state) { + const int16_t max_level = static_cast(state.range(0)); + const int num_levels = static_cast(state.range(1)); + const int batch_size = static_cast(state.range(2)); + const int level_repeat_count = static_cast(state.range(3)); + + DecodeAndCountLevels(Encoding::RLE, max_level, num_levels, batch_size, + level_repeat_count, state); +} + +static void ReadLevelsAndCount_BitPack_Separate(::benchmark::State& state) { + const int16_t max_level = static_cast(state.range(0)); + const int num_levels = static_cast(state.range(1)); + const int batch_size = static_cast(state.range(2)); + const int level_repeat_count = static_cast(state.range(3)); + + DecodeAndCountLevels(Encoding::BIT_PACKED, max_level, num_levels, batch_size, + level_repeat_count, state); +} + +static void ReadLevelsAndCount_BitPack_Fused(::benchmark::State& state) { + const int16_t max_level = static_cast(state.range(0)); + const int num_levels = static_cast(state.range(1)); + const int batch_size = static_cast(state.range(2)); + const int level_repeat_count = static_cast(state.range(3)); + + DecodeAndCountLevels(Encoding::BIT_PACKED, max_level, num_levels, batch_size, + level_repeat_count, state); +} + static void ReadLevelsArguments(::benchmark::internal::Benchmark* b) { b->ArgNames({"MaxLevel", "NumLevels", "BatchSize", "LevelRepeatCount"}) ->Args({1, 8096, 1024, 1}) @@ -365,8 +457,26 @@ static void ReadLevelsArguments(::benchmark::internal::Benchmark* b) { ->Args({3, 8096, 1024, 7}); } +static void ReadLevelsAndCountArguments(::benchmark::internal::Benchmark* b) { + b->ArgNames({"MaxLevel", "NumLevels", "BatchSize", "LevelRepeatCount"}) + ->Args({1, 8096, 1024, 1}) + ->Args({1, 8096, 1024, 7}) + ->Args({1, 8096, 1024, 1024}) + ->Args({3, 8096, 1024, 1}) + ->Args({3, 8096, 2048, 1}) + ->Args({3, 8096, 1024, 7}); +} + BENCHMARK(ReadLevels_Rle)->Apply(ReadLevelsArguments); BENCHMARK(ReadLevels_BitPack)->Apply(ReadLevelsArguments); +BENCHMARK(ReadLevelsAndCount_Rle_Separate)->Apply(ReadLevelsAndCountArguments); + +BENCHMARK(ReadLevelsAndCount_Rle_Fused)->Apply(ReadLevelsAndCountArguments); + +BENCHMARK(ReadLevelsAndCount_BitPack_Separate)->Apply(ReadLevelsAndCountArguments); + +BENCHMARK(ReadLevelsAndCount_BitPack_Fused)->Apply(ReadLevelsAndCountArguments); + } // namespace benchmarks } // namespace parquet