This is an automated email from the ASF dual-hosted git repository.
pitrou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 5815dc0b4e GH-50421: [C++][Parquet] Add LevelDecoder Skip and Count
(#50422)
5815dc0b4e is described below
commit 5815dc0b4e3d9239828b47d427c9759fbc506e09
Author: Antoine Prouvost <[email protected]>
AuthorDate: Fri Jul 10 16:22:42 2026 +0200
GH-50421: [C++][Parquet] Add LevelDecoder Skip and Count (#50422)
### Rationale for this change
Adding new methods to the `LevelDecoder` to start reducing the complexity
in `ColumnReaderImplBase`, `TypedColumnReaderImpl`, and `TypedRecordReader`.
### What changes are included in this PR?
- Simplify `LevelDecoder` members based on a single variant
- Add `LevelDecoder::Skip`
- Add `LevelDecoder::CountUpTo`
- Simplify `Skip` implementation in `TypedColumnReaderImpl`
- Add `SkippableTypedDecoder` to wrap `TypedDecoder` with the remaining
usage of the `scratch_for_skip_` buffer
### Are these changes tested?
Yes.
### Are there any user-facing changes?
No.
* GitHub Issue: #50421
Lead-authored-by: AntoinePrv <[email protected]>
Co-authored-by: Antoine Prouvost <[email protected]>
Co-authored-by: Copilot Autofix powered by AI
<[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
cpp/src/arrow/util/rle_encoding_internal.h | 291 +++++++++++++++++-------
cpp/src/arrow/util/rle_encoding_test.cc | 206 +++++++++++++++--
cpp/src/parquet/column_reader.cc | 352 ++++++++++++++++++-----------
cpp/src/parquet/column_reader.h | 54 +++--
4 files changed, 644 insertions(+), 259 deletions(-)
diff --git a/cpp/src/arrow/util/rle_encoding_internal.h
b/cpp/src/arrow/util/rle_encoding_internal.h
index 5dc94f368d..4bd08fd653 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<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 matching_count;
+ rle_size_t processed_count;
+};
+
/// Decoder class for a single run of RLE encoded data.
template <typename T>
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<value_type>& 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<value_type>& 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<rle_size_t>(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 <typename Callable>
+ rle_size_t ProcessValues(Callable&& func, rle_size_t batch_size);
/// Utility methods for retrieving spaced values.
template <typename Converter>
@@ -515,6 +572,82 @@ 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.
+ ///
+ /// @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) 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) noexcept {
+ ARROW_DCHECK_GE(value_count, 0);
+ ARROW_DCHECK_GE(value_bit_width, 0);
+ ARROW_DCHECK_LE(value_bit_width, 64);
+
+ value_bit_width_ = value_bit_width;
+
+ ARROW_DCHECK_LE(value_count, std::numeric_limits<rle_size_t>::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 +915,26 @@ struct RleBitPackedDecoderGetRunDecoder<T, BitPackedRun> {
};
template <typename T>
-bool RleBitPackedDecoder<T>::Get(value_type* val) {
- return GetBatch(val, 1) == 1;
-}
-
-template <typename T>
-auto RleBitPackedDecoder<T>::GetBatch(value_type* out,
- rle_size_t batch_size) -> rle_size_t {
+template <typename Callable>
+auto RleBitPackedDecoder<T>::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 +943,66 @@ 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
- 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 <typename T>
+auto RleBitPackedDecoder<T>::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 <typename T>
+RleCountUpToResult RleBitPackedDecoder<T>::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 <typename T>
+bool RleBitPackedDecoder<T>::Get(value_type* val) {
+ return GetBatch(val, 1) == 1;
+}
+
+template <typename T>
+auto RleBitPackedDecoder<T>::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 +1441,6 @@ 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;
@@ -1278,59 +1448,20 @@ auto RleBitPackedDecoder<T>::GetBatchWithDict(const V*
dictionary,
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(
+ [&](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 <typename T>
diff --git a/cpp/src/arrow/util/rle_encoding_test.cc
b/cpp/src/arrow/util/rle_encoding_test.cc
index 788450c67c..16b702c23d 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 <algorithm>
#include <bit>
#include <cstdint>
#include <cstring>
@@ -335,6 +336,31 @@ TEST(Rle, RleDecoder) {
/* expected_value= */ 16777749);
}
+TEST(Rle, RleDecoderCountUpTo) {
+ // A run of value 21 repeated 23 times.
+ const std::array<uint8_t, 3> bytes = {21, 0, 0};
+ const auto run = RleRun(bytes.data(), /* value_count= */ 23, /* bit_width=
*/ 5);
+ auto decoder = RleRunDecoder<uint8_t>(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 <typename T>
void TestBitPackedDecoder(std::vector<uint8_t> bytes, rle_size_t value_count,
rle_size_t bit_width, std::vector<T> 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 <typename Type>
+std::vector<uint8_t> EncodeTestArray(const Array& data, int bit_width,
+ bool spaced = false) {
+ using ArrayType = typename TypeTraits<Type>::ArrayType;
+
+ const auto data_size = static_cast<rle_size_t>(data.length());
+ const auto values_count =
+ static_cast<rle_size_t>(data.length() - (spaced ? data.null_count() :
0));
+ const int buffer_size =
+ static_cast<int>(RleBitPackedEncoder::MaxBufferSize(bit_width,
values_count) +
+ RleBitPackedEncoder::MinBufferSize(bit_width));
+ const auto* data_values = static_cast<const ArrayType&>(data).raw_values();
+
+ std::vector<uint8_t> 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<uint64_t>(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<int>(data.length());
- const int data_values_count =
- static_cast<int>(data.length() - spaced * data.null_count());
- const int buffer_size = static_cast<int>(
- ::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<const
ArrayType&>(data).raw_values();
// Encode the data into `buffer` using the encoder.
- std::vector<uint8_t> 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<uint64_t>(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<Type>(data, bit_width, spaced);
// Now we verify batch read
- RleBitPackedDecoder<value_type> decoder(buffer.data(), encoded_byte_size,
bit_width);
+ RleBitPackedDecoder<value_type> decoder(buffer.data(),
static_cast<int>(buffer.size()),
+ bit_width);
// We will only use one of them depending on whether this is a dictionary
tests
std::vector<float> dict_read;
std::vector<value_type> 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 <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) {
+ 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<value_type> 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<rle_size_t>(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 <typename Type>
+void CheckCountUpTo(const Array& data, int bit_width, typename Type::c_type
value) {
+ using ArrayType = typename TypeTraits<Type>::ArrayType;
+ using value_type = typename Type::c_type;
+
+ const auto data_size = static_cast<rle_size_t>(data.length());
+ const value_type* data_values =
+ arrow::internal::checked_cast<const ArrayType&>(data).raw_values();
+
+ ARROW_SCOPED_TRACE("bit_width = ", bit_width, ", data_size = ", data_size,
+ ", value = ", value);
+
+ // Encode all values into `buffer`.
+ const auto buffer = EncodeTestArray<Type>(data, bit_width);
+
+ RleBitPackedDecoder<value_type> decoder(buffer.data(),
static_cast<int>(buffer.size()),
+ bit_width);
+
+ // 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<rle_size_t>(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<rle_size_t>(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<rle_size_t>(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 <typename T>
struct DataTestRleBitPackedRandomPart {
using value_type = T;
@@ -1173,6 +1321,7 @@ template <typename T>
void DoTestGetBatchSpacedRoundtrip() {
using Data = DataTestRleBitPacked<T>;
using ArrowType = typename Data::ArrowType;
+ using ArrayType = typename TypeTraits<ArrowType>::ArrayType;
using RandomPart = typename Data::RandomPart;
using NullPart = typename Data::NullPart;
using RepeatPart = typename Data::RepeatPart;
@@ -1257,6 +1406,19 @@ void DoTestGetBatchSpacedRoundtrip() {
CheckRoundTrip<ArrowType>(*array, case_.bit_width, /* spaced= */ false,
/* parts= */ 3);
+ // Tests for Advance
+ CheckAdvance<ArrowType>(*array, case_.bit_width);
+ CheckAdvance<ArrowType>(*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<T>(arrow::internal::checked_cast<const
ArrayType&>(*array).Value(0));
+ const auto max_value = bit_util::LeastSignificantBitMask<T,
true>(case_.bit_width);
+ CheckCountUpTo<ArrowType>(*array, case_.bit_width, first);
+ CheckCountUpTo<ArrowType>(*array, case_.bit_width, max_value);
+ CheckCountUpTo<ArrowType>(*array->Slice(1), case_.bit_width, first);
+
// Tests for GetBatchSpaced
CheckRoundTrip<ArrowType>(*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 7241632e4b..162a72bd15 100644
--- a/cpp/src/parquet/column_reader.cc
+++ b/cpp/src/parquet/column_reader.cc
@@ -27,6 +27,7 @@
#include <type_traits>
#include <unordered_map>
#include <utility>
+#include <variant>
#include <vector>
#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<int16_t>;
+ using BitPackedDecoder = ::arrow::util::BitPackedDecoder<int16_t>;
+
+ std::variant<RleBitPackedDecoder, BitPackedDecoder> 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<Impl>()), 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<int32_t>(data);
+ const auto num_bytes = ::arrow::util::SafeLoadAs<int32_t>(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<int16_t>>(
- 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<int32_t>(bit_util::BytesForBits(num_bits));
+ const auto num_bytes =
static_cast<int32_t>(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<int16_t>>(
- 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>
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.
+///
+/// @todo GH-50453
+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 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<int>::max());
+ const int batch_size =
+ static_cast<int>(std::min(kScratchValueCount, num_values -
total_read));
+
+ iter_read = get()->Decode(scratch_->mutable_data_as<T>(), 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<ResizableBuffer> 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<int>(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<int>(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<int>(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<int>(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<int>(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<int>(num_buffered_values_), buffer);
}
@@ -869,7 +963,7 @@ class ColumnReaderImplBase {
auto it = decoders_.find(static_cast<int>(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<DType>(encoding, descr_, pool_);
- current_decoder_ = decoder.get();
+ current_decoder_.SetDecoder(decoder.get());
decoders_[static_cast<int>(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<PageReader> pager_;
std::shared_ptr<Page> 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<DType>;
- DecoderType* current_decoder_;
- Encoding::type current_encoding_;
+ SkippableTypedDecoder<DType, kSkipScratchBatchSize> 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<DType>,
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<ResizableBuffer> 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<DictDecoder<DType>*>(this->current_decoder_);
+ auto decoder =
dynamic_cast<DictDecoder<DType>*>(this->current_decoder_.get());
return decoder->DecodeIndices(static_cast<int>(indices_to_read), indices);
}
@@ -1006,7 +1100,7 @@ class TypedColumnReaderImpl : public
TypedColumnReader<DType>,
// 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<DictDecoder<DType>*>(this->current_decoder_);
+ auto decoder =
dynamic_cast<DictDecoder<DType>*>(this->current_decoder_.get());
decoder->GetDictionary(dictionary, dictionary_length);
}
@@ -1020,7 +1114,7 @@ class TypedColumnReaderImpl : public
TypedColumnReader<DType>,
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<DType>,
// 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<DType>,
}
// 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<DType>::ReadBatch(int64_t
batch_size, int16_t* def
return total_values;
}
-template <typename DType>
-void TypedColumnReaderImpl<DType>::InitScratchForSkip() {
- if (this->scratch_for_skip_ == nullptr) {
- int value_size = type_traits<DType::type_num>::value_byte_size;
- this->scratch_for_skip_ = AllocateBuffer(
- this->pool_, kSkipScratchBatchSize * std::max<int>(sizeof(int16_t),
value_size));
- }
-}
-
template <typename DType>
int64_t TypedColumnReaderImpl<DType>::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<int>(batch_size),
- scratch_for_skip_->mutable_data_as<int16_t>(),
- scratch_for_skip_->mutable_data_as<int16_t>(),
- scratch_for_skip_->mutable_data_as<T>(),
&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<int>(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<DType>,
}
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<DType>,
<< EncodingToString(this->current_encoding_);
throw ParquetException(ss.str());
}
- auto decoder = dynamic_cast<DictDecoder<DType>*>(this->current_decoder_);
+ auto decoder =
dynamic_cast<DictDecoder<DType>*>(this->current_decoder_.get());
const T* dictionary = nullptr;
decoder->GetDictionary(&dictionary, dictionary_length);
return reinterpret_cast<const void*>(dictionary);
@@ -1323,7 +1417,7 @@ class TypedRecordReader : public
TypedColumnReaderImpl<DType>,
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<DType>,
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<DType>,
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<DType>,
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<DType>,
// 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<DType>,
// 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<DType>,
// 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<DType>,
// 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<int64_t>(kSkipScratchBatchSize,
values_left);
- values_read = this->ReadValues(
- batch_size, this->scratch_for_skip_->template mutable_data_as<T>());
- 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<DType>,
// 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<DType>,
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<DType>,
}
// 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<DType>,
}
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<DType>,
}
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<DType>,
// 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<DType>,
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<DType>,
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<ByteArray
/// insert the new dictionary values
FlushBuilder();
builder_.ResetFull();
- auto decoder = dynamic_cast<BinaryDictDecoder*>(this->current_decoder_);
+ auto decoder =
dynamic_cast<BinaryDictDecoder*>(this->current_decoder_.get());
decoder->InsertDictionary(&builder_);
this->new_dictionary_ = false;
}
@@ -2138,7 +2220,7 @@ class ByteArrayDictionaryRecordReader final : public
TypedRecordReader<ByteArray
int64_t num_decoded = 0;
if (current_encoding_ == Encoding::RLE_DICTIONARY) {
MaybeWriteNewDictionary();
- auto decoder = dynamic_cast<BinaryDictDecoder*>(this->current_decoder_);
+ auto decoder =
dynamic_cast<BinaryDictDecoder*>(this->current_decoder_.get());
num_decoded = decoder->DecodeIndices(static_cast<int>(values_to_read),
&builder_);
} else {
num_decoded = this->current_decoder_->DecodeArrowNonNull(
@@ -2153,7 +2235,7 @@ class ByteArrayDictionaryRecordReader final : public
TypedRecordReader<ByteArray
int64_t num_decoded = 0;
if (current_encoding_ == Encoding::RLE_DICTIONARY) {
MaybeWriteNewDictionary();
- auto decoder = dynamic_cast<BinaryDictDecoder*>(this->current_decoder_);
+ auto decoder =
dynamic_cast<BinaryDictDecoder*>(this->current_decoder_.get());
num_decoded = decoder->DecodeIndicesSpaced(
static_cast<int>(values_to_read), static_cast<int>(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 ad20d4a9e6..07064a9cf3 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 <typename T>
-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<int16_t>> rle_decoder_;
- std::unique_ptr<::arrow::bit_util::BitReader> bit_packed_decoder_;
+ struct Impl;
+
+ std::unique_ptr<Impl> 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_;
};
