pitrou commented on code in PR #50422:
URL: https://github.com/apache/arrow/pull/50422#discussion_r3551799643
##########
cpp/src/arrow/util/rle_encoding_internal.h:
##########
@@ -1269,68 +1444,27 @@ template <typename V>
auto RleBitPackedDecoder<T>::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;
}
internal::DictionaryConverter<V, value_type> 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<value_type,
decltype(run)>::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(
Review Comment:
Neat refactor!
##########
cpp/src/arrow/util/rle_encoding_test.cc:
##########
@@ -335,6 +335,31 @@ TEST(Rle, RleDecoder) {
/* expected_value= */ 16777749);
}
+TEST(Rle, RleDecoderCountUpTo) {
Review Comment:
`RleBitPackedDecoder::CountUpTo` doesn't seem tested otherwise, or did I
miss something?
##########
cpp/src/parquet/column_reader.h:
##########
@@ -78,26 +65,44 @@ 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);
+
+ /// 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.
+ std::pair<int, int> CountUpTo(int16_t value, int batch_size);
Review Comment:
Can we return a `struct` with field names to make meaning of the results
more explicit?
##########
cpp/src/parquet/column_reader.cc:
##########
@@ -629,6 +657,70 @@ std::unique_ptr<PageReader>
PageReader::Open(std::shared_ptr<ArrowInputStream> 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.
+template <typename DType, int64_t kScratchValueCount_>
Review Comment:
Nit: remove trailing _
```suggestion
template <typename DType, int64_t kScratchValueCount>
```
##########
cpp/src/arrow/util/rle_encoding_internal.h:
##########
@@ -515,6 +573,85 @@ 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 <typename T>
+class BitPackedDecoder : private BitPackedRunDecoder<T> {
+ private:
+ using Base = BitPackedRunDecoder<T>;
+
+ 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.
Review Comment:
Doesn't the caller always know the value count? It's the number of levels
declared in the page header.
##########
cpp/src/arrow/util/rle_encoding_internal.h:
##########
@@ -814,23 +947,65 @@ auto RleBitPackedDecoder<T>::GetBatch(value_type* out,
using RunDecoder =
typename RleBitPackedDecoderGetRunDecoder<value_type,
decltype(run)>::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
Review Comment:
Keep this comment perhaps?
##########
cpp/src/parquet/column_reader.cc:
##########
@@ -95,52 +96,73 @@ constexpr std::string_view
kErrorRepDefLevelNotMatchesNumValues =
} // namespace
-LevelDecoder::LevelDecoder() : num_values_remaining_(0) {}
+/******************
+ * LevelDecoder *
+ ******************/
+
+struct LevelDecoder::Impl {
+ using RleBitPackedDecoder = ::arrow::util::RleBitPackedDecoder<int16_t>;
+ using BitPackedDecoder = ::arrow::util::BitPackedDecoder<int16_t>;
+
+ std::variant<RleBitPackedDecoder, BitPackedDecoder> decoder = {};
Review Comment:
I'm curious why you're using a variant rather than a polymorphic base class
with derived implementations for RLE and BIT_PACKED.
(not saying you should necessarily switch to the latter, but it seems more
idiomatic and more flexible - you may want to add implementation-specific
fields, for example)
##########
cpp/src/arrow/util/rle_encoding_internal.h:
##########
@@ -258,6 +259,18 @@ class RleBitPackedParser {
std::pair<rle_size_t, ControlFlow> PeekImpl(Handler&&) const;
};
+template <typename T>
+struct RleCountUpToParams {
+ T value;
+ rle_size_t batch_size;
+ rle_size_t value_bit_width;
+};
+
+struct RleCountUpToResult {
+ rle_size_t count;
+ rle_size_t advanced_count;
Review Comment:
Nit, but I don't find the names very understandable, how about
`matching_count` and `processed_count` (or `num_matching_values` and
`num_processed_values`)?
(feel free to disregard if you don't like this suggestion, btw)
##########
cpp/src/parquet/column_reader.cc:
##########
@@ -629,6 +657,70 @@ std::unique_ptr<PageReader>
PageReader::Open(std::shared_ptr<ArrowInputStream> 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.
Review Comment:
Can you open a separate issue for this improvement?
##########
cpp/src/parquet/column_reader.h:
##########
@@ -78,26 +65,44 @@ 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);
+
+ /// 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.
+ std::pair<int, int> 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<int16_t>> rle_decoder_;
- std::unique_ptr<::arrow::bit_util::BitReader> bit_packed_decoder_;
+ struct Impl;
+
+ std::unique_ptr<Impl> decoder_;
Review Comment:
Typical name for this would be `impl_` :)
##########
cpp/src/parquet/column_reader.cc:
##########
@@ -629,6 +657,70 @@ std::unique_ptr<PageReader>
PageReader::Open(std::shared_ptr<ArrowInputStream> 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.
+template <typename DType, int64_t kScratchValueCount_>
+class SkippableTypedDecoder {
+ public:
+ using Decoder = TypedDecoder<DType>;
+ using T = typename Decoder::T;
+
+ static constexpr int64_t kValueByteSize =
type_traits<DType::type_num>::value_byte_size;
+ static constexpr int64_t kScratchValueCount = kScratchValueCount_;
+ static constexpr int64_t kScratchByteSize = kScratchValueCount *
kValueByteSize;
+
+ SkippableTypedDecoder() = default;
+
+ 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, ::arrow::MemoryPool* pool = nullptr) {
Review Comment:
I'd rather pass the `MemoryPool` to the `SkippableTypedDecoder` constructor,
as is already the case in `MakeDecoder`.
##########
cpp/src/parquet/column_reader.cc:
##########
@@ -95,52 +96,73 @@ constexpr std::string_view
kErrorRepDefLevelNotMatchesNumValues =
} // namespace
-LevelDecoder::LevelDecoder() : num_values_remaining_(0) {}
+/******************
+ * LevelDecoder *
+ ******************/
+
+struct LevelDecoder::Impl {
+ using RleBitPackedDecoder = ::arrow::util::RleBitPackedDecoder<int16_t>;
+ using BitPackedDecoder = ::arrow::util::BitPackedDecoder<int16_t>;
+
+ std::variant<RleBitPackedDecoder, BitPackedDecoder> decoder = {};
Review Comment:
Note that at some point we may have more specialized implementations, for
example an `AlwaysMaxLevelDecoderImpl` and a `AlwaysZeroLevelDecoderImpl`
depending on the `num_nulls` in the header.
##########
cpp/src/arrow/util/rle_encoding_test.cc:
##########
@@ -1102,6 +1142,60 @@ 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 <typename Type>
+void CheckAdvance(const Array& data, int bit_width) {
+ 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 = static_cast<const
ArrayType&>(data).raw_values();
+
+ ARROW_SCOPED_TRACE("bit_width = ", bit_width, ", data_size = ", data_size);
+
+ // Encode all values into `buffer`.
+ const auto buffer = EncodeTestArray<Type>(data, bit_width);
+
+ // Interleave Advance and GetBatch, verifying decoded values against the
original.
+ RleBitPackedDecoder<value_type> decoder(buffer.data(),
static_cast<int>(buffer.size()),
+ bit_width);
+ rle_size_t pos = 0;
+ auto advance = [&](rle_size_t* pos, rle_size_t n) {
Review Comment:
It's weird to pass `rle_size_t* pos` explicitly even though you're capturing
the entire environment by reference.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
