diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index 4bd08fd6530..21341694662 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 16b702c23d4..91b7bcc764d 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); diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 56c65fd233b..ff80457707f 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 07064a9cf37..34df9837d35 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. diff --git a/cpp/src/parquet/column_reader_benchmark.cc b/cpp/src/parquet/column_reader_benchmark.cc index 83f661bf9ad..5de96888430 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