diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index 5dc94f368d3c..7c22b13490fd 100644 --- a/cpp/src/arrow/util/rle_encoding_internal.h +++ b/cpp/src/arrow/util/rle_encoding_internal.h @@ -258,6 +258,18 @@ class RleBitPackedParser { std::pair PeekImpl(Handler&&) const; }; +template +struct RleCountUpToParams { + T value; + rle_size_t batch_size; + rle_size_t value_bit_width; +}; + +struct RleCountUpToResult { + rle_size_t matching_count; + rle_size_t processed_count; +}; + /// Decoder class for a single run of RLE encoded data. template class RleRunDecoder { @@ -300,6 +312,15 @@ class RleRunDecoder { return steps; } + /// Advance 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 processed. + RleCountUpToResult CountUpTo(const RleCountUpToParams& p) { + const auto steps = Advance(p.batch_size); + return {.matching_count = steps * (p.value == value_), .processed_count = steps}; + } + /// Get the next value and return false if there are no more. [[nodiscard]] constexpr bool Get(value_type* out_value, rle_size_t value_bit_width) { return GetBatch(out_value, 1, value_bit_width) == 1; @@ -363,6 +384,29 @@ class BitPackedRunDecoder { return steps; } + /// Advance and count the number of occurrence of a values. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of of element that were processed. + RleCountUpToResult CountUpTo(const RleCountUpToParams& p) { + // Decoding in a stack buffer of 512 values: 1KB for int16_t used in levels + // and up to 4KB for uint64_t. + constexpr rle_size_t kBufferValueCount = 512; + alignas(16) value_type buffer[kBufferValueCount]; // uninitialized + + rle_size_t remaining = p.batch_size; + rle_size_t count = 0; + rle_size_t read_iter = 0; + do { + const auto batch_iter = std::min(remaining, kBufferValueCount); + read_iter = GetBatch(buffer, batch_iter, p.value_bit_width); + count += static_cast(std::count(buffer, buffer + read_iter, p.value)); + remaining -= read_iter; + } while (remaining > 0 && read_iter > 0); + + return {.matching_count = count, .processed_count = p.batch_size - remaining}; + } + /// Get the next value and return false if there are no more. [[nodiscard]] constexpr bool Get(value_type* out_value, rle_size_t value_bit_width) { return GetBatch(out_value, 1, value_bit_width) == 1; @@ -448,6 +492,16 @@ class RleBitPackedDecoder { /// This is how one can check for errors. bool exhausted() const { return (run_remaining() == 0) && parser_.exhausted(); } + /// Advance by as many values as provided or until exhaustion of the decoder. + /// Return the number of values skipped. + [[nodiscard]] rle_size_t Advance(rle_size_t batch_size); + + /// Advance and count the number of occurrence of a values. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of of element that were processed. + RleCountUpToResult CountUpTo(value_type value, rle_size_t batch_size); + /// Gets the next value or returns false if there are no more or an error occurred. /// /// NB: Because the encoding only supports literal runs with lengths @@ -500,12 +554,15 @@ class RleBitPackedDecoder { return std::visit([](const auto& dec) { return dec.remaining(); }, decoder_); } - /// Get a batch of values from the current run and return the number elements read. - [[nodiscard]] rle_size_t RunGetBatch(value_type* out, rle_size_t batch_size) { - return std::visit( - [&](auto& dec) { return dec.GetBatch(out, batch_size, value_bit_width_); }, - decoder_); - } + /// Process data in the current run and subsequent ones. + /// + /// `func` is called as `func(decoder, run_batch_size)` where `decoder` are + /// statically-typed run decoder (not the variant). + /// Must return the number of values it processed. + /// + /// Return the number of values processed. + template + rle_size_t ProcessValues(Callable&& func, rle_size_t batch_size); /// Utility methods for retrieving spaced values. template @@ -515,6 +572,86 @@ class RleBitPackedDecoder { int64_t valid_bits_offset, rle_size_t null_count); }; +/// Minimal decoder for legacy bit packed encoding (BIT_PACKED = 4). +/// +/// The number of values that the decoder can represent is up to 2^31 - 1. +template +class BitPackedDecoder : private BitPackedRunDecoder { + private: + using Base = BitPackedRunDecoder; + + public: + /// The type in which the data should be decoded. + using value_type = T; + + using Base::Advance; + + BitPackedDecoder() noexcept = default; + + /// Create a decoder object. + /// + /// Unless explicitly given, then number of values is deduced from the data size, + /// in which case it may read more values from the input stream than the user + /// intended (to reach final byte alignment). This option is mandatory if the + /// `value_bit_width` is zero. + /// + /// @param data and data_size are the raw bytes to decode. + /// @param value_bit_width is the size in bits of each encoded value. + /// @param value_count is an optional number of values in the run. + BitPackedDecoder(const uint8_t* data, rle_size_t data_size, rle_size_t value_bit_width, + rle_size_t value_count = -1) noexcept { + Reset(data, data_size, value_bit_width, value_count); + } + + void Reset(const uint8_t* data, rle_size_t data_size, rle_size_t value_bit_width, + rle_size_t value_count = -1) noexcept { + ARROW_DCHECK_GE(value_bit_width, 0); + ARROW_DCHECK_LE(value_bit_width, 64); + + value_bit_width_ = value_bit_width; + if (value_count < 0) { + ARROW_DCHECK_GT(value_bit_width, 0); + const int64_t data_bit_size = static_cast(data_size) * 8; + value_count = static_cast(data_bit_size / value_bit_width); + } + + ARROW_DCHECK_LE(value_count, std::numeric_limits::max()); + const auto run = BitPackedRun{ + /* data= */ data, + /* value_count= */ value_count, + /* value_bit_width= */ value_bit_width, + /* max_read_bytes= */ data_size, + }; + return Base::Reset(run, value_bit_width); + } + + /// Whether there is still values to iterate over. + bool exhausted() const { return Base::remaining() == 0; } + + /// Advance and count the number of occurrence of a values. + /// + /// The count is limited to at most the next `batch_size` items. + /// @return The matching value count and number of of element that were processed. + RleCountUpToResult CountUpTo(value_type value, rle_size_t batch_size) { + return Base::CountUpTo({ + .value = value, + .batch_size = batch_size, + .value_bit_width = value_bit_width_, + }); + } + + /// Gets the next value or returns false if there are no more or an error occurred. + [[nodiscard]] bool Get(value_type* val) { return Base::Get(val, value_bit_width_); } + + /// Get a batch of values return the number of decoded elements. + [[nodiscard]] rle_size_t GetBatch(value_type* out, rle_size_t batch_size) { + return Base::GetBatch(out, batch_size, value_bit_width_); + } + + private: + rle_size_t value_bit_width_ = {}; +}; + /// Class to incrementally build the rle data. This class does not allocate any memory. /// The encoding has two modes: encoding repeated runs and literal runs. /// If the run is sufficiently short, it is more efficient to encode as a literal run. @@ -782,30 +919,26 @@ struct RleBitPackedDecoderGetRunDecoder { }; template -bool RleBitPackedDecoder::Get(value_type* val) { - return GetBatch(val, 1) == 1; -} - -template -auto RleBitPackedDecoder::GetBatch(value_type* out, - rle_size_t batch_size) -> rle_size_t { +template +auto RleBitPackedDecoder::ProcessValues(Callable&& func, + rle_size_t batch_size) -> rle_size_t { using ControlFlow = RleBitPackedParser::ControlFlow; if (ARROW_PREDICT_FALSE(batch_size == 0 || exhausted())) { return 0; } - rle_size_t values_read = 0; + rle_size_t values_processed = 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 = RunGetBatch(out, batch_size); - values_read += read; - out += read; + const auto processed = + std::visit([&](auto& decoder) { return func(decoder, batch_size); }, decoder_); + values_processed += processed; // 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; + if (ARROW_PREDICT_FALSE(values_processed == batch_size)) { + return values_processed; } ARROW_DCHECK(run_remaining() == 0); } @@ -814,23 +947,66 @@ auto RleBitPackedDecoder::GetBatch(value_type* out, using RunDecoder = typename RleBitPackedDecoderGetRunDecoder::type; - ARROW_DCHECK_LT(values_read, batch_size); + ARROW_DCHECK_LT(values_processed, batch_size); + // Local decoder to keep its type transparent to the compiler. RunDecoder decoder(run, value_bit_width_); - const auto read = decoder.GetBatch(out, batch_size - values_read, value_bit_width_); - ARROW_DCHECK_LE(read, batch_size - values_read); - values_read += read; - out += read; + const auto read = func(decoder, batch_size - values_processed); + ARROW_DCHECK_LE(read, batch_size - values_processed); + values_processed += read; - // Stop reading and store remaining decoder - if (ARROW_PREDICT_FALSE(values_read == batch_size || read == 0)) { + if (ARROW_PREDICT_FALSE(values_processed == batch_size || read == 0)) { + // Stop reading and store remaining decoder decoder_ = std::move(decoder); return ControlFlow::Break; } - return ControlFlow::Continue; }); - return values_read; + return values_processed; +} + +template +auto RleBitPackedDecoder::Advance(rle_size_t batch_size) -> rle_size_t { + return ProcessValues( + [](auto& decoder, rle_size_t run_batch_size) { + return decoder.Advance(run_batch_size); + }, + batch_size); +} + +template +RleCountUpToResult RleBitPackedDecoder::CountUpTo(value_type value, + rle_size_t batch_size) { + rle_size_t count = 0; + const rle_size_t processed_count = ProcessValues( + [value, this, &count](auto& decoder, rle_size_t run_batch_size) { + const auto result = decoder.CountUpTo({ + .value = value, + .batch_size = run_batch_size, + .value_bit_width = value_bit_width_, + }); + count += result.matching_count; + return result.processed_count; + }, + batch_size); + return {.matching_count = count, .processed_count = processed_count}; +} + +template +bool RleBitPackedDecoder::Get(value_type* val) { + return GetBatch(val, 1) == 1; +} + +template +auto RleBitPackedDecoder::GetBatch(value_type* out, + rle_size_t batch_size) -> rle_size_t { + return ProcessValues( + [&out, this](auto& decoder, rle_size_t run_batch_size) { + const auto read = decoder.GetBatch(out, run_batch_size, value_bit_width_); + out += read; + return read; + }, + batch_size); } namespace internal { @@ -1269,8 +1445,6 @@ template auto RleBitPackedDecoder::GetBatchWithDict(const V* dictionary, int32_t dictionary_length, V* out, rle_size_t batch_size) -> rle_size_t { - using ControlFlow = RleBitPackedParser::ControlFlow; - if (ARROW_PREDICT_FALSE(batch_size <= 0 || dictionary_length == 0)) { // Either empty batch or invalid dictionary return 0; @@ -1278,59 +1452,20 @@ auto RleBitPackedDecoder::GetBatchWithDict(const V* dictionary, internal::DictionaryConverter converter{dictionary, dictionary_length}; - // Make lightweight BitRun class to reuse previous methods. + // Make lightweight BitRun class to reuse the spaced code path with no nulls. constexpr internal::UnreachableBitRunReader validity_reader{}; internal::AllSetBitRun validity_run = {batch_size}; - rle_size_t values_read = 0; - auto batch_values_remaining = [&]() { - ARROW_DCHECK_LE(values_read, batch_size); - return batch_size - values_read; - }; - - if (ARROW_PREDICT_FALSE(run_remaining() > 0)) { - const auto read = internal::RunGetSpaced(&converter, out, batch_size, - /* null_count= */ 0, value_bit_width_, - &validity_reader, &validity_run, &decoder_); - - ARROW_DCHECK_EQ(read.null_read, 0); - values_read += read.values_read; - out += read.values_read; - - // Either we fulfilled all the batch values to be read - if (ARROW_PREDICT_FALSE(values_read >= batch_size)) { - // There may be remaining null if they are not greedily filled - return values_read; - } - - // We finished the remaining run - ARROW_DCHECK(run_remaining() == 0); - } - - parser_.ParseWithCallable([&](auto run) { - using RunDecoder = - typename RleBitPackedDecoderGetRunDecoder::type; - - RunDecoder decoder(run, value_bit_width_); - - const auto read = internal::RunGetSpaced(&converter, out, batch_values_remaining(), - /* null_count= */ 0, value_bit_width_, - &validity_reader, &validity_run, &decoder); - - ARROW_DCHECK_EQ(read.null_read, 0); - values_read += read.values_read; - out += read.values_read; - - // Stop reading and store remaining decoder - if (ARROW_PREDICT_FALSE(read.values_read == 0 || values_read == batch_size)) { - decoder_ = std::move(decoder); - return ControlFlow::Break; - } - - return ControlFlow::Continue; - }); - - return values_read; + return ProcessValues( + [&](auto& decoder, rle_size_t run_batch_size) { + const auto read = internal::RunGetSpaced( + &converter, out, run_batch_size, /* null_count= */ 0, value_bit_width_, + &validity_reader, &validity_run, &decoder); + ARROW_DCHECK_EQ(read.null_read, 0); + out += read.values_read; + return read.values_read; + }, + batch_size); } template diff --git a/cpp/src/arrow/util/rle_encoding_test.cc b/cpp/src/arrow/util/rle_encoding_test.cc index 788450c67c90..16b702c23d49 100644 --- a/cpp/src/arrow/util/rle_encoding_test.cc +++ b/cpp/src/arrow/util/rle_encoding_test.cc @@ -17,6 +17,7 @@ // From Apache Impala (incubating) as of 2016-01-29 +#include #include #include #include @@ -335,6 +336,31 @@ TEST(Rle, RleDecoder) { /* expected_value= */ 16777749); } +TEST(Rle, RleDecoderCountUpTo) { + // A run of value 21 repeated 23 times. + const std::array bytes = {21, 0, 0}; + const auto run = RleRun(bytes.data(), /* value_count= */ 23, /* bit_width= */ 5); + auto decoder = RleRunDecoder(run, /* value_bit_width= */ 5); + + // Counting the repeated value counts every advanced element. + auto res = decoder.CountUpTo({.value = 21, .batch_size = 10, .value_bit_width = 5}); + EXPECT_EQ(res.processed_count, 10); + EXPECT_EQ(res.matching_count, 10); + EXPECT_EQ(decoder.remaining(), 23 - 10); + + // Counting another value matches nothing but still advances. + res = decoder.CountUpTo({.value = 99, .batch_size = 5, .value_bit_width = 5}); + EXPECT_EQ(res.processed_count, 5); + EXPECT_EQ(res.matching_count, 0); + EXPECT_EQ(decoder.remaining(), 23 - 15); + + // Requesting more than remaining is capped to the remaining count. + res = decoder.CountUpTo({.value = 21, .batch_size = 100, .value_bit_width = 5}); + EXPECT_EQ(res.processed_count, 8); + EXPECT_EQ(res.matching_count, 8); + EXPECT_EQ(decoder.remaining(), 0); +} + template void TestBitPackedDecoder(std::vector bytes, rle_size_t value_count, rle_size_t bit_width, std::vector expected) { @@ -989,6 +1015,39 @@ TEST(BitRle, Overflow) { } } +/// Encode the values of an Array with RleBitPackedEncoder. +/// +/// @param spaced If false, treat Nulls as regular data (i.e. their raw value is +/// encoded). If true, Nulls are skipped and only valid values are encoded. +template +std::vector EncodeTestArray(const Array& data, int bit_width, + bool spaced = false) { + using ArrayType = typename TypeTraits::ArrayType; + + const auto data_size = static_cast(data.length()); + const auto values_count = + static_cast(data.length() - (spaced ? data.null_count() : 0)); + const int buffer_size = + static_cast(RleBitPackedEncoder::MaxBufferSize(bit_width, values_count) + + RleBitPackedEncoder::MinBufferSize(bit_width)); + const auto* data_values = static_cast(data).raw_values(); + + std::vector buffer(buffer_size); + RleBitPackedEncoder encoder(buffer.data(), buffer_size, bit_width); + rle_size_t encoded_values_count = 0; + for (rle_size_t i = 0; i < data_size; ++i) { + if (data.IsValid(i) || !spaced) { + EXPECT_TRUE(encoder.Put(static_cast(data_values[i]))) + << "Encoding failed in pos " << i << ", current encoder len: " << encoder.len(); + ++encoded_values_count; + } + } + EXPECT_EQ(encoded_values_count, values_count) + << "All values input were not encoded successfully by the encoder"; + buffer.resize(encoder.Flush()); + return buffer; +} + /// Check RleBitPacked encoding/decoding round trip. /// /// \param spaced If set to false, treat Nulls in the input array as regular data. @@ -1002,38 +1061,20 @@ void CheckRoundTrip(const Array& data, int bit_width, bool spaced, int32_t parts using value_type = typename Type::c_type; const int data_size = static_cast(data.length()); - const int data_values_count = - static_cast(data.length() - spaced * data.null_count()); - const int buffer_size = static_cast( - ::arrow::util::RleBitPackedEncoder::MaxBufferSize(bit_width, data_values_count) + - ::arrow::util::RleBitPackedEncoder::MinBufferSize(bit_width)); ASSERT_GE(parts, 1); ASSERT_LE(parts, data_size); ARROW_SCOPED_TRACE("bit_width = ", bit_width, ", spaced = ", spaced, - ", data_size = ", data_size, ", buffer_size = ", buffer_size); + ", data_size = ", data_size); const value_type* data_values = static_cast(data).raw_values(); // Encode the data into `buffer` using the encoder. - std::vector buffer(buffer_size); - RleBitPackedEncoder encoder(buffer.data(), buffer_size, bit_width); - int32_t encoded_values_size = 0; - for (int i = 0; i < data_size; ++i) { - // Depending on `spaced` we treat nulls as regular values. - if (data.IsValid(i) || !spaced) { - bool success = encoder.Put(static_cast(data_values[i])); - ASSERT_TRUE(success) << "Encoding failed in pos " << i - << ", current encoder len: " << encoder.len(); - ++encoded_values_size; - } - } - int encoded_byte_size = encoder.Flush(); - ASSERT_EQ(encoded_values_size, data_values_count) - << "All values input were not encoded successfully by the encoder"; + const auto buffer = EncodeTestArray(data, bit_width, spaced); // Now we verify batch read - RleBitPackedDecoder decoder(buffer.data(), encoded_byte_size, bit_width); + RleBitPackedDecoder decoder(buffer.data(), static_cast(buffer.size()), + bit_width); // We will only use one of them depending on whether this is a dictionary tests std::vector dict_read; std::vector values_read; @@ -1102,6 +1143,113 @@ void CheckRoundTrip(const Array& data, int bit_width, bool spaced, int32_t parts } } +/// Check RleBitPackedDecoder::Advance, which spans multiple runs through the +/// parser (unlike the per-run decoder Advance). +/// +/// Nulls are treated as regular data (i.e. their raw value is encoded and +/// decoded). Advance and GetBatch are interleaved so that run boundaries are +/// crossed and partial runs consumed, verifying decoded values as we go. +template +void CheckAdvance(const Array& data, int bit_width) { + 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 = static_cast(data).raw_values(); + + ARROW_SCOPED_TRACE("bit_width = ", bit_width, ", data_size = ", data_size); + + // Encode all values into `buffer`. + const auto buffer = EncodeTestArray(data, bit_width); + + // Interleave Advance and GetBatch, verifying decoded values against the original. + RleBitPackedDecoder decoder(buffer.data(), static_cast(buffer.size()), + bit_width); + rle_size_t pos = 0; + auto advance = [&](rle_size_t* pos, rle_size_t n) { + n = std::min(n, data_size - *pos); + EXPECT_EQ(decoder.Advance(n), n); + *pos += n; + }; + auto read = [&](rle_size_t* pos, rle_size_t n) { + n = std::min(n, data_size - *pos); + std::vector got(n); + EXPECT_EQ(decoder.GetBatch(got.data(), n), n); + for (rle_size_t i = 0; i < n; ++i) { + EXPECT_EQ(got[i], data_values[*pos + i]) << "at position " << (*pos + i); + } + *pos += n; + }; + + // Advance in a `step` small relative to the data size, so we span many + // iterations and repeatedly cross run boundaries. + const rle_size_t step = std::max(data_size / 16, 1); + int iter = 0; + while (pos < data_size) { + // Some way to make various successions of `read` of `advance` + if (bit_width % 2 == iter % 3) { + read(&pos, step); + } + advance(&pos, step); + ++iter; + } + // Note: we do not assert exhaustion here because the encoder pads the last + // literal group to a multiple of 8 with zeros, leaving up to 7 extra values. +} + +/// Check RleBitPackedDecoder::CountUpTo, which spans multiple runs through the +/// parser (unlike the per-run decoder CountUpTo). +/// +/// Nulls are treated as regular data (i.e. their raw value is encoded and +/// decoded). The counts returned over successive batches are compared against a +/// naive count over the original data. +template +void CheckCountUpTo(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); + + // Count in a `step` small relative to the data size, so we span many + // iterations and repeatedly cross run boundaries. + const rle_size_t step = std::max(data_size / 16, 1); + 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 res = decoder.CountUpTo(value, to_process); + ASSERT_EQ(res.processed_count, to_process); + + // The matching count of this batch must equal a naive count over the data. + const auto expected = + std::count(data_values + pos, data_values + pos + to_process, value); + EXPECT_EQ(res.matching_count, static_cast(expected)) + << "at position " << pos; + + pos += res.processed_count; + total_matching += res.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)); + + // The decoder is exhausted of the values we requested: counting further only + // yields the padding values the encoder appended to the last literal group. + const auto res = decoder.CountUpTo(value, data_size); + EXPECT_LT(res.processed_count, 8); +} + template struct DataTestRleBitPackedRandomPart { using value_type = T; @@ -1173,6 +1321,7 @@ template void DoTestGetBatchSpacedRoundtrip() { using Data = DataTestRleBitPacked; using ArrowType = typename Data::ArrowType; + using ArrayType = typename TypeTraits::ArrayType; using RandomPart = typename Data::RandomPart; using NullPart = typename Data::NullPart; using RepeatPart = typename Data::RepeatPart; @@ -1257,6 +1406,19 @@ void DoTestGetBatchSpacedRoundtrip() { CheckRoundTrip(*array, case_.bit_width, /* spaced= */ false, /* parts= */ 3); + // Tests for Advance + CheckAdvance(*array, case_.bit_width); + CheckAdvance(*array->Slice(1), case_.bit_width); + + // Tests for CountUpTo, counting a value present in the data (the first one) + // and a value that may not be (the max encodable one). + const auto first = + static_cast(arrow::internal::checked_cast(*array).Value(0)); + const auto max_value = bit_util::LeastSignificantBitMask(case_.bit_width); + CheckCountUpTo(*array, case_.bit_width, first); + CheckCountUpTo(*array, case_.bit_width, max_value); + CheckCountUpTo(*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 7241632e4b62..162a72bd1576 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -27,6 +27,7 @@ #include #include #include +#include #include #include "arrow/array.h" @@ -95,7 +96,32 @@ constexpr std::string_view kErrorRepDefLevelNotMatchesNumValues = } // namespace -LevelDecoder::LevelDecoder() : num_values_remaining_(0) {} +/****************** + * LevelDecoder * + ******************/ + +struct LevelDecoder::Impl { + using RleBitPackedDecoder = ::arrow::util::RleBitPackedDecoder; + using BitPackedDecoder = ::arrow::util::BitPackedDecoder; + + std::variant decoder = {}; + + [[nodiscard]] int GetBatch(int16_t* out, int batch_size) { + return std::visit([&](auto& dec) { return dec.GetBatch(out, batch_size); }, decoder); + } + + [[nodiscard]] int Advance(int batch_size) { + return std::visit([&](auto& dec) { return dec.Advance(batch_size); }, decoder); + } + + auto CountUpTo(int16_t value, int batch_size) { + return std::visit([&](auto& dec) { return dec.CountUpTo(value, batch_size); }, + decoder); + } +}; + +LevelDecoder::LevelDecoder(int16_t max_level) + : impl_(std::make_unique()), max_level_(max_level) {} LevelDecoder::~LevelDecoder() = default; @@ -103,44 +129,40 @@ int LevelDecoder::SetData(Encoding::type encoding, int16_t max_level, int num_buffered_values, const uint8_t* data, int32_t data_size) { max_level_ = max_level; - int32_t num_bytes = 0; - encoding_ = encoding; num_values_remaining_ = num_buffered_values; - bit_width_ = bit_util::Log2(max_level + 1); + const int value_bit_width = bit_util::Log2(max_level + 1); + switch (encoding) { case Encoding::RLE: { if (data_size < 4) { throw ParquetException("Received invalid levels (corrupt data page?)"); } - num_bytes = ::arrow::util::SafeLoadAs(data); + const auto num_bytes = ::arrow::util::SafeLoadAs(data); if (num_bytes < 0 || num_bytes > data_size - 4) { throw ParquetException("Received invalid number of bytes (corrupt data page?)"); } - const uint8_t* decoder_data = data + 4; - if (!rle_decoder_) { - rle_decoder_ = std::make_unique<::arrow::util::RleBitPackedDecoder>( - decoder_data, num_bytes, bit_width_); - } else { - rle_decoder_->Reset(decoder_data, num_bytes, bit_width_); - } + this->impl_->decoder = Impl::RleBitPackedDecoder( // + /* data= */ data + 4, + /* data_size =*/num_bytes, + /* value_bit_width= */ value_bit_width); return 4 + num_bytes; } case Encoding::BIT_PACKED: { int num_bits = 0; - if (MultiplyWithOverflow(num_buffered_values, bit_width_, &num_bits)) { + if (MultiplyWithOverflow(num_buffered_values, value_bit_width, &num_bits)) { throw ParquetException( "Number of buffered values too large (corrupt data page?)"); } - num_bytes = static_cast(bit_util::BytesForBits(num_bits)); + const auto num_bytes = static_cast(bit_util::BytesForBits(num_bits)); if (num_bytes < 0 || num_bytes > data_size) { throw ParquetException("Received invalid number of bytes (corrupt data page?)"); } - if (!bit_packed_decoder_) { - bit_packed_decoder_ = - std::make_unique<::arrow::bit_util::BitReader>(data, num_bytes); - } else { - bit_packed_decoder_->Reset(data, num_bytes); - } + // Also adding `value_count` so that the decoder also works with zero-width runs. + this->impl_->decoder = Impl::BitPackedDecoder( // + /* data= */ data, + /* data_size =*/num_bytes, + /* value_bit_width= */ value_bit_width, + /* value_count= */ num_buffered_values); return num_bytes; } default: @@ -157,27 +179,17 @@ void LevelDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, if (num_bytes < 0) { throw ParquetException("Invalid page header (corrupt data page?)"); } - encoding_ = Encoding::RLE; num_values_remaining_ = num_buffered_values; - bit_width_ = bit_util::Log2(max_level + 1); - if (!rle_decoder_) { - rle_decoder_ = std::make_unique<::arrow::util::RleBitPackedDecoder>( - data, num_bytes, bit_width_); - } else { - rle_decoder_->Reset(data, num_bytes, bit_width_); - } + this->impl_->decoder = Impl::RleBitPackedDecoder( // + /* data= */ data, + /* data_size =*/num_bytes, + /* value_bit_width= */ bit_util::Log2(max_level + 1)); } int LevelDecoder::Decode(int batch_size, int16_t* levels) { - int num_decoded = 0; - - int num_values = std::min(num_values_remaining_, batch_size); - if (encoding_ == Encoding::RLE) { - num_decoded = rle_decoder_->GetBatch(levels, num_values); - } else { - num_decoded = bit_packed_decoder_->GetBatch(bit_width_, levels, num_values); - } + const int num_values = std::min(num_values_remaining_, batch_size); + const int num_decoded = impl_->GetBatch(levels, num_values); 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_)) { @@ -191,6 +203,25 @@ int LevelDecoder::Decode(int batch_size, int16_t* levels) { return 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); + ARROW_DCHECK_EQ(num_values, num_advanced); + num_values_remaining_ -= num_advanced; + return num_advanced; +} + +auto LevelDecoder::CountUpTo(int16_t value, int batch_size) -> CountUpToResult { + const int num_values = std::min(num_values_remaining_, batch_size); + const auto result = impl_->CountUpTo(value, num_values); + ARROW_DCHECK_EQ(num_values, result.processed_count); + num_values_remaining_ -= result.processed_count; + return { + .matching_count = result.matching_count, + .processed_count = result.processed_count, + }; +} + ReaderProperties default_reader_properties() { static ReaderProperties default_reader_properties; return default_reader_properties; @@ -629,6 +660,72 @@ std::unique_ptr PageReader::Open(std::shared_ptr s namespace { +/// Wrapper around a `TypedDecoder` pointer to skip values. +/// +/// Use a scratch buffer to decode values into that buffer. +/// This was migrated here from a historical implementation. +/// Ideally all decoders would implement a `Skip` functionality that would at best +/// avoid decoding, and at worst, decode without intermediary allocation. +/// +/// @todo GH-50453 +template +class SkippableTypedDecoder { + public: + using Decoder = TypedDecoder; + using T = typename Decoder::T; + + static constexpr int64_t kValueByteSize = type_traits::value_byte_size; + static constexpr int64_t kScratchByteSize = kScratchValueCount * kValueByteSize; + + explicit SkippableTypedDecoder(::arrow::MemoryPool* pool = nullptr) : pool_(pool) {} + + explicit SkippableTypedDecoder(Decoder* decoder) : decoder_(decoder) {} + + void SetDecoder(Decoder* decoder) { decoder_ = decoder; } + + const Decoder* get() const { return decoder_; } + + Decoder* get() { return decoder_; } + + const Decoder* operator->() const { return decoder_; } + + Decoder* operator->() { return decoder_; } + + explicit operator bool() const { return decoder_ != nullptr; } + + int64_t Skip(int64_t num_values) { + EnsureScratch(); + + int64_t total_read = 0; + int iter_read = 0; + do { + static_assert(kScratchValueCount <= std::numeric_limits::max()); + const int batch_size = + static_cast(std::min(kScratchValueCount, num_values - total_read)); + + iter_read = get()->Decode(scratch_->mutable_data_as(), batch_size); + total_read += iter_read; + } while (iter_read > 0 && total_read < num_values); + + return total_read; + } + + private: + /// Scratch space to skip decode values that need skipping. + /// We actually do not need the whole shared_ptr machinery but it was historically + /// chosen for ease of use with ``AllocateBuffer`` and migrated here. + std::shared_ptr scratch_ = nullptr; + Decoder* decoder_ = nullptr; + ::arrow::MemoryPool* pool_ = nullptr; + + void EnsureScratch() { + if (this->scratch_ == nullptr) { + this->scratch_ = AllocateBuffer(pool_, kScratchByteSize); + } + ARROW_DCHECK_NE(this->scratch_, nullptr); + } +}; + // ---------------------------------------------------------------------- // Impl base class for TypedColumnReader and RecordReader @@ -639,13 +736,10 @@ class ColumnReaderImplBase { ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) : descr_(descr), - max_def_level_(descr->max_definition_level()), - max_rep_level_(descr->max_repetition_level()), - num_buffered_values_(0), - num_decoded_values_(0), + definition_level_decoder_(descr->max_definition_level()), + repetition_level_decoder_(descr->max_repetition_level()), pool_(pool), - current_decoder_(nullptr), - current_encoding_(Encoding::UNKNOWN) {} + current_decoder_(pool) {} virtual ~ColumnReaderImplBase() = default; @@ -675,7 +769,7 @@ class ColumnReaderImplBase { // // Returns the number of decoded definition levels int64_t ReadDefinitionLevels(int64_t batch_size, int16_t* levels) { - if (max_def_level_ == 0) { + if (max_def_level() == 0) { return 0; } return definition_level_decoder_.Decode(static_cast(batch_size), levels); @@ -695,7 +789,7 @@ class ColumnReaderImplBase { // Read multiple repetition levels into preallocated memory // Returns the number of decoded repetition levels int64_t ReadRepetitionLevels(int64_t batch_size, int16_t* levels) { - if (max_rep_level_ == 0) { + if (max_rep_level() == 0) { return 0; } return repetition_level_decoder_.Decode(static_cast(batch_size), levels); @@ -766,7 +860,7 @@ class ColumnReaderImplBase { } new_dictionary_ = true; - current_decoder_ = decoders_[encoding].get(); + current_decoder_.SetDecoder(decoders_[encoding].get()); ARROW_DCHECK(current_decoder_); } @@ -791,9 +885,9 @@ class ColumnReaderImplBase { // Data page Layout: Repetition Levels - Definition Levels - encoded values. // Levels are encoded as rle or bit-packed. // Init repetition levels - if (max_rep_level_ > 0) { + if (max_rep_level() > 0) { int32_t rep_levels_bytes = repetition_level_decoder_.SetData( - repetition_level_encoding, max_rep_level_, + repetition_level_encoding, max_rep_level(), static_cast(num_buffered_values_), buffer, max_size); buffer += rep_levels_bytes; levels_byte_size += rep_levels_bytes; @@ -803,9 +897,9 @@ class ColumnReaderImplBase { // if the initial value is invalid // Init definition levels - if (max_def_level_ > 0) { + if (max_def_level() > 0) { int32_t def_levels_bytes = definition_level_decoder_.SetData( - definition_level_encoding, max_def_level_, + definition_level_encoding, max_def_level(), static_cast(num_buffered_values_), buffer, max_size); levels_byte_size += def_levels_bytes; max_size -= def_levels_bytes; @@ -830,9 +924,9 @@ class ColumnReaderImplBase { throw ParquetException("Data page too small for levels (corrupt header?)"); } - if (max_rep_level_ > 0) { + if (max_rep_level() > 0) { repetition_level_decoder_.SetDataV2(page.repetition_levels_byte_length(), - max_rep_level_, + max_rep_level(), static_cast(num_buffered_values_), buffer); } // ARROW-17453: Even if max_rep_level_ is 0, there may still be @@ -840,9 +934,9 @@ class ColumnReaderImplBase { // some writers (e.g. Athena) buffer += page.repetition_levels_byte_length(); - if (max_def_level_ > 0) { + if (max_def_level() > 0) { definition_level_decoder_.SetDataV2(page.definition_levels_byte_length(), - max_def_level_, + max_def_level(), static_cast(num_buffered_values_), buffer); } @@ -869,7 +963,7 @@ class ColumnReaderImplBase { auto it = decoders_.find(static_cast(encoding)); if (it != decoders_.end()) { ARROW_DCHECK(it->second.get() != nullptr); - current_decoder_ = it->second.get(); + current_decoder_.SetDecoder(it->second.get()); } else { switch (encoding) { case Encoding::PLAIN: @@ -879,7 +973,7 @@ class ColumnReaderImplBase { case Encoding::DELTA_BYTE_ARRAY: case Encoding::DELTA_LENGTH_BYTE_ARRAY: { auto decoder = MakeTypedDecoder(encoding, descr_, pool_); - current_decoder_ = decoder.get(); + current_decoder_.SetDecoder(decoder.get()); decoders_[static_cast(encoding)] = std::move(decoder); break; } @@ -902,17 +996,25 @@ class ColumnReaderImplBase { return num_buffered_values_ - num_decoded_values_; } + int16_t max_def_level() const { + // max level indirectly part of this object storage + return definition_level_decoder_.max_level(); + } + + int16_t max_rep_level() const { + // max level indirectly part of this object storage + return repetition_level_decoder_.max_level(); + } + const ColumnDescriptor* descr_; - const int16_t max_def_level_; - const int16_t max_rep_level_; std::unique_ptr pager_; std::shared_ptr current_page_; - // Not set if full schema for this field has no optional or repeated elements + // No data set if full schema for this field has no optional or repeated elements LevelDecoder definition_level_decoder_; - // Not set for flat schemas. + // No data set for flat schemas. LevelDecoder repetition_level_decoder_; // The total number of values stored in the data page. This is the maximum of @@ -921,17 +1023,17 @@ class ColumnReaderImplBase { // values. For repeated or optional values, there may be fewer data values // than levels, and this tells you how many encoded levels there are in that // case. - int64_t num_buffered_values_; + int64_t num_buffered_values_ = 0; // The number of values from the current data page that have been decoded // into memory or skipped over. - int64_t num_decoded_values_; + int64_t num_decoded_values_ = 0; ::arrow::MemoryPool* pool_; using DecoderType = TypedDecoder; - DecoderType* current_decoder_; - Encoding::type current_encoding_; + SkippableTypedDecoder current_decoder_; + Encoding::type current_encoding_ = Encoding::UNKNOWN; /// Flag to signal when a new dictionary has been set, for the benefit of /// DictionaryRecordReader @@ -986,19 +1088,11 @@ class TypedColumnReaderImpl : public TypedColumnReader, this->exposed_encoding_ = encoding; } - // Allocate enough scratch space to accommodate skipping 16-bit levels or any - // value type. - void InitScratchForSkip(); - - // Scratch space for reading and throwing away rep/def levels and values when - // skipping. - std::shared_ptr scratch_for_skip_; - private: // Read dictionary indices. Similar to ReadValues but decode data to dictionary indices. // This function is called only by ReadBatchWithDictionary(). int64_t ReadDictionaryIndices(int64_t indices_to_read, int32_t* indices) { - auto decoder = dynamic_cast*>(this->current_decoder_); + auto decoder = dynamic_cast*>(this->current_decoder_.get()); return decoder->DecodeIndices(static_cast(indices_to_read), indices); } @@ -1006,7 +1100,7 @@ class TypedColumnReaderImpl : public TypedColumnReader, // owned by the internal decoder and is destroyed when the reader is destroyed. This // function is called only by ReadBatchWithDictionary() after dictionary is configured. void GetDictionary(const T** dictionary, int32_t* dictionary_length) { - auto decoder = dynamic_cast*>(this->current_decoder_); + auto decoder = dynamic_cast*>(this->current_decoder_.get()); decoder->GetDictionary(dictionary, dictionary_length); } @@ -1020,7 +1114,7 @@ class TypedColumnReaderImpl : public TypedColumnReader, batch_size = std::min(batch_size, this->available_values_current_page()); // If the field is required and non-repeated, there are no definition levels - if (this->max_def_level_ > 0 && def_levels != nullptr) { + 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)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); @@ -1028,7 +1122,7 @@ class TypedColumnReaderImpl : public TypedColumnReader, // 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_); + std::count(def_levels, def_levels + *num_def_levels, this->max_def_level()); } else { // Required field, read all values if (num_def_levels != nullptr) { @@ -1038,7 +1132,7 @@ class TypedColumnReaderImpl : public TypedColumnReader, } // Not present for non-repeated fields - if (this->max_rep_level_ > 0 && rep_levels != nullptr) { + if (this->max_rep_level() > 0 && rep_levels != nullptr) { int64_t num_rep_levels = this->ReadRepetitionLevels(batch_size, rep_levels); if (batch_size != num_rep_levels) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); @@ -1137,40 +1231,40 @@ int64_t TypedColumnReaderImpl::ReadBatch(int64_t batch_size, int16_t* def return total_values; } -template -void TypedColumnReaderImpl::InitScratchForSkip() { - if (this->scratch_for_skip_ == nullptr) { - int value_size = type_traits::value_byte_size; - this->scratch_for_skip_ = AllocateBuffer( - this->pool_, kSkipScratchBatchSize * std::max(sizeof(int16_t), value_size)); - } -} - template int64_t TypedColumnReaderImpl::Skip(int64_t num_values_to_skip) { int64_t values_to_skip = num_values_to_skip; // Optimization: Do not call HasNext() when values_to_skip == 0. while (values_to_skip > 0 && HasNext()) { // If the number of values to skip is more than the number of undecoded values, skip - // the Page. + // the whole Page without decoding levels or values. const int64_t available_values = this->available_values_current_page(); if (values_to_skip >= available_values) { values_to_skip -= available_values; this->ConsumeBufferedValues(available_values); } else { - // We need to read this Page - // Jump to the right offset in the Page - int64_t values_read = 0; - InitScratchForSkip(); - ARROW_DCHECK_NE(this->scratch_for_skip_, nullptr); - do { - int64_t batch_size = std::min(kSkipScratchBatchSize, values_to_skip); - values_read = ReadBatch(static_cast(batch_size), - scratch_for_skip_->mutable_data_as(), - scratch_for_skip_->mutable_data_as(), - scratch_for_skip_->mutable_data_as(), &values_read); - values_to_skip -= values_read; - } while (values_read > 0 && values_to_skip > 0); + // Skip within the current Page. Since `values_to_skip < available_values`, the + // whole batch fits inside this Page and no page boundary is crossed. + const int batch_size = static_cast(values_to_skip); + + // Advance the definition levels, counting how many correspond to present + // (non-null) values that must be skipped in the data decoder. + int64_t non_null_values_to_skip = batch_size; + if (this->max_def_level() > 0) { + const auto count = + this->definition_level_decoder_.CountUpTo(this->max_def_level(), batch_size); + non_null_values_to_skip = count.matching_count; + ARROW_DCHECK_EQ(count.processed_count, batch_size); + } + // Advance the repetition levels; their values are not needed. + if (this->max_rep_level() > 0) { + this->repetition_level_decoder_.Skip(batch_size); + } + // Skip the corresponding data values. + this->current_decoder_.Skip(non_null_values_to_skip); + + this->ConsumeBufferedValues(batch_size); + values_to_skip -= batch_size; } } return num_values_to_skip - values_to_skip; @@ -1267,7 +1361,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, } const void* ReadDictionary(int32_t* dictionary_length) override { - if (this->current_decoder_ == nullptr && !this->HasNextInternal()) { + if (!this->current_decoder_ && !this->HasNextInternal()) { *dictionary_length = 0; return nullptr; } @@ -1280,7 +1374,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, << EncodingToString(this->current_encoding_); throw ParquetException(ss.str()); } - auto decoder = dynamic_cast*>(this->current_decoder_); + auto decoder = dynamic_cast*>(this->current_decoder_.get()); const T* dictionary = nullptr; decoder->GetDictionary(&dictionary, dictionary_length); return reinterpret_cast(dictionary); @@ -1323,7 +1417,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, break; } - if (this->max_def_level_ > 0) { + if (this->max_def_level() > 0) { ReserveLevels(batch_size); int16_t* def_levels = this->def_levels() + levels_written_; @@ -1333,7 +1427,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); } - if (this->max_rep_level_ > 0) { + if (this->max_rep_level() > 0) { int64_t rep_levels_read = this->ReadRepetitionLevels(batch_size, rep_levels); if (ARROW_PREDICT_FALSE(rep_levels_read != batch_size)) { throw ParquetException(kErrorRepDefLevelNotMatchesNumValues); @@ -1359,7 +1453,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, void ThrowAwayLevels(int64_t start_levels_position) { ARROW_DCHECK_LE(levels_position_, levels_written_); ARROW_DCHECK_LE(start_levels_position, levels_position_); - ARROW_DCHECK_GT(this->max_def_level_, 0); + ARROW_DCHECK_GT(this->max_def_level(), 0); ARROW_DCHECK_NE(def_levels_, nullptr); int64_t gap = levels_position_ - start_levels_position; @@ -1377,7 +1471,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, left_shift(def_levels_.get()); - if (this->max_rep_level_ > 0) { + if (this->max_rep_level() > 0) { ARROW_DCHECK_NE(rep_levels_, nullptr); left_shift(rep_levels_.get()); } @@ -1390,7 +1484,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, // Skip records that we have in our buffer. This function is only for // non-repeated fields. int64_t SkipRecordsInBufferNonRepeated(int64_t num_records) { - ARROW_DCHECK_EQ(this->max_rep_level_, 0); + ARROW_DCHECK_EQ(this->max_rep_level(), 0); if (!this->has_values_to_process() || num_records == 0) return 0; int64_t remaining_records = levels_written_ - levels_position_; @@ -1404,7 +1498,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, // First we need to figure out how many present/not-null values there are. int64_t values_to_read = std::count(def_levels() + start_levels_position, def_levels() + levels_position_, - this->max_def_level_); + this->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. @@ -1451,7 +1545,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, // desired number of records or we run out of values in the column chunk. // Returns number of skipped records. int64_t SkipRecordsRepeated(int64_t num_records) { - ARROW_DCHECK_GT(this->max_rep_level_, 0); + ARROW_DCHECK_GT(this->max_rep_level(), 0); int64_t skipped_records = 0; // First consume what is in the buffer. @@ -1516,20 +1610,8 @@ class TypedRecordReader : public TypedColumnReaderImpl, // Read 'num_values' values and throw them away. // Throws an error if it could not read 'num_values'. void ReadAndThrowAwayValues(int64_t num_values) { - int64_t values_left = num_values; - int64_t values_read = 0; - - // Allocate enough scratch space to accommodate 16-bit levels or any - // value type - this->InitScratchForSkip(); - ARROW_DCHECK_NE(this->scratch_for_skip_, nullptr); - do { - int64_t batch_size = std::min(kSkipScratchBatchSize, values_left); - values_read = this->ReadValues( - batch_size, this->scratch_for_skip_->template mutable_data_as()); - values_left -= values_read; - } while (values_read > 0 && values_left > 0); - if (values_left > 0) { + const int64_t values_read = this->current_decoder_.Skip(num_values); + if (values_read < num_values) { std::stringstream ss; ss << "Could not read and throw away " << num_values << " values"; throw ParquetException(ss.str()); @@ -1541,11 +1623,11 @@ class TypedRecordReader : public TypedColumnReaderImpl, // Top level required field. Number of records equals to number of levels, // and there is not read-ahead for levels. - if (this->max_rep_level_ == 0 && this->max_def_level_ == 0) { + if (this->max_rep_level() == 0 && this->max_def_level() == 0) { return this->Skip(num_records); } int64_t skipped_records = 0; - if (this->max_rep_level_ == 0) { + if (this->max_rep_level() == 0) { // Non-repeated optional field. // First consume whatever is in the buffer. skipped_records = SkipRecordsInBufferNonRepeated(num_records); @@ -1606,7 +1688,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, int64_t records_read = 0; const int16_t* const rep_levels = this->rep_levels(); const int16_t* const def_levels = this->def_levels(); - ARROW_DCHECK_GT(this->max_rep_level_, 0); + ARROW_DCHECK_GT(this->max_rep_level(), 0); // If at_record_start_ is true, we are seeing the start of a record // for the second time, such as after repeated calls to // DelimitRecords. In this case we must continue until we find @@ -1652,7 +1734,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, } // Scan definition levels to find number of physical values *values_seen = std::count(def_levels + level, def_levels + levels_position_, - this->max_def_level_); + this->max_def_level()); return records_read; } @@ -1679,7 +1761,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, } void ReserveLevels(int64_t extra_levels) { - if (this->max_def_level_ > 0) { + if (this->max_def_level() > 0) { const int64_t new_levels_capacity = UpdateCapacity(levels_capacity_, levels_written_, extra_levels); if (new_levels_capacity > levels_capacity_) { @@ -1690,7 +1772,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, } PARQUET_THROW_NOT_OK( def_levels_->Resize(capacity_in_bytes, /*shrink_to_fit=*/false)); - if (this->max_rep_level_ > 0) { + if (this->max_rep_level() > 0) { PARQUET_THROW_NOT_OK( rep_levels_->Resize(capacity_in_bytes, /*shrink_to_fit=*/false)); } @@ -1830,7 +1912,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, // 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_); + def_levels + levels_position_, this->max_def_level()); ReadValuesDense(*values_to_read); } @@ -1869,11 +1951,11 @@ class TypedRecordReader : public TypedColumnReaderImpl, int64_t records_read = 0; int64_t values_to_read = 0; int64_t null_count = 0; - if (this->max_rep_level_ > 0) { + if (this->max_rep_level() > 0) { // Repeated fields may be nullable or not. // This call updates levels_position_. records_read = ReadRepeatedRecords(num_records, &values_to_read, &null_count); - } else if (this->max_def_level_ > 0) { + } else if (this->max_def_level() > 0) { // Non-repeated optional values are always nullable. // This call updates levels_position_. ARROW_DCHECK(nullable_values()); @@ -1896,7 +1978,7 @@ class TypedRecordReader : public TypedColumnReaderImpl, null_count_ += null_count; } // Total values, including null spaces, if any - if (this->max_def_level_ > 0) { + if (this->max_def_level() > 0) { // Optional, repeated, or some mix thereof this->ConsumeBufferedValues(levels_position_ - start_levels_position); } else { @@ -2128,7 +2210,7 @@ class ByteArrayDictionaryRecordReader final : public TypedRecordReader(this->current_decoder_); + auto decoder = dynamic_cast(this->current_decoder_.get()); decoder->InsertDictionary(&builder_); this->new_dictionary_ = false; } @@ -2138,7 +2220,7 @@ class ByteArrayDictionaryRecordReader final : public TypedRecordReader(this->current_decoder_); + auto decoder = dynamic_cast(this->current_decoder_.get()); num_decoded = decoder->DecodeIndices(static_cast(values_to_read), &builder_); } else { num_decoded = this->current_decoder_->DecodeArrowNonNull( @@ -2153,7 +2235,7 @@ class ByteArrayDictionaryRecordReader final : public TypedRecordReader(this->current_decoder_); + auto decoder = dynamic_cast(this->current_decoder_.get()); num_decoded = decoder->DecodeIndicesSpaced( static_cast(values_to_read), static_cast(null_count), valid_bits_->mutable_data(), values_written_, &builder_); diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index ad20d4a9e6e6..07064a9cf37d 100644 --- a/cpp/src/parquet/column_reader.h +++ b/cpp/src/parquet/column_reader.h @@ -32,19 +32,6 @@ #include "parquet/schema.h" #include "parquet/types.h" -namespace arrow { - -namespace bit_util { -class BitReader; -} // namespace bit_util - -namespace util { -template -class RleBitPackedDecoder; -} // namespace util - -} // namespace arrow - namespace parquet { class Decryptor; @@ -78,26 +65,49 @@ struct PARQUET_EXPORT DataPageStats { class PARQUET_EXPORT LevelDecoder { public: - LevelDecoder(); + explicit LevelDecoder(int16_t max_level = 0); + ~LevelDecoder(); - // Initialize the LevelDecoder state with new data - // and return the number of bytes consumed + /// Initialize the LevelDecoder state with new data from a legacy (V1) page. + /// + /// @return the number of bytes consumed int SetData(Encoding::type encoding, int16_t max_level, int num_buffered_values, const uint8_t* data, int32_t data_size); + /// Initialize the LevelDecoder state with new data from a V2 page. + /// + /// Repetition and definition levels in V2 pages are always RLE encoded. void SetDataV2(int32_t num_bytes, int16_t max_level, int num_buffered_values, const uint8_t* data); - // Decodes a batch of levels into an array and returns the number of levels decoded + /// Decode a batch of levels into an array and returns the number of levels decoded. int Decode(int batch_size, int16_t* levels); + /// Advance the decoder and throw away decoder levels. + int Skip(int batch_size); + + struct CountUpToResult { + int matching_count; + int processed_count; + }; + + /// Advance 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 processed. + CountUpToResult CountUpTo(int16_t value, int batch_size); + + /// Return the max level used in this decoder. + int max_level() const { return max_level_; } + private: - int bit_width_; - int num_values_remaining_; - Encoding::type encoding_; - std::unique_ptr<::arrow::util::RleBitPackedDecoder> rle_decoder_; - std::unique_ptr<::arrow::bit_util::BitReader> bit_packed_decoder_; + struct Impl; + + std::unique_ptr impl_; + /// Number of value remaining. The underlying decoder zero pads bit packed values + /// up to a multiple of 8 so it cannot know the exact number of remaining values. + int num_values_remaining_ = 0; int16_t max_level_; };