pitrou commented on code in PR #50629: URL: https://github.com/apache/arrow/pull/50629#discussion_r4323217280
########## cpp/src/parquet/level_decoder_internal.h: ########## @@ -0,0 +1,314 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include <cstdint> +#include <sstream> +#include <variant> + +#include "arrow/util/bit_util.h" +#include "arrow/util/int_util_overflow.h" +#include "arrow/util/logging.h" +#include "arrow/util/macros.h" +#include "arrow/util/rle_bitmap_internal.h" +#include "arrow/util/rle_encoding_internal.h" +#include "arrow/util/ubsan.h" +#include "parquet/exception.h" +#include "parquet/level_comparison.h" +#include "parquet/types.h" + +namespace parquet { + +/******************************* + * PageLevelDecoder Adapters * + *******************************/ + +/// BitPackedDecoder adapter for PageLevelDecoder. +struct LevelBitPackedDecoder : ::arrow::util::BitPackedDecoder<int16_t> { + using Base = BitPackedDecoder<int16_t>; + + LevelBitPackedDecoder(const uint8_t* data, int32_t data_size, const auto& params) Review Comment: Is `params` truly generic or can we specify its type explicitly for clarity? If it's only `PageLevelDecoder::DataParams` perhaps we can extract it as a standalone `LevelDecoderParams` struct? ########## cpp/src/parquet/column_reader.h: ########## @@ -63,52 +63,63 @@ struct PARQUET_EXPORT DataPageStats { std::optional<int32_t> num_rows; }; +/// Decoder for repetition or definition level. +/// +/// This decoder is used with either a deprecated bit packed (`BIT_PACKED = 4`) +/// encoding or a mixed bit packed and RLE one (`RLE = 3`). +/// Because it takes as input a single buffer, `SetData` and `Decode` are typically +/// used on each of the Parquet `DataPage`. +/// The number of levels is guaranteed to fit into an `int32_t` by the specification. +/// +/// @see https://research.google.com/pubs/archive/36632.pdf class PARQUET_EXPORT LevelDecoder { public: explicit LevelDecoder(int16_t max_level = 0); + LevelDecoder(LevelDecoder&&); + LevelDecoder& operator=(LevelDecoder&&); + ~LevelDecoder(); /// 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); + int32_t SetData(Encoding::type encoding, int16_t max_level, int32_t 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, + void SetDataV2(int32_t num_bytes, int16_t max_level, int32_t num_buffered_values, const uint8_t* data); - /// Decode a batch of levels into an array and returns the number of levels decoded. - int Decode(int batch_size, int16_t* levels); + /// Decode a batch of levels into an array and return the number of levels decoded. + int32_t Decode(int32_t batch_size, int16_t* levels); /// Advance the decoder and throw away decoder levels. - int Skip(int batch_size); + int32_t Skip(int32_t batch_size); struct CountUpToResult { - int matching_count; - int processed_count; + int32_t matching_count; + int32_t 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); + CountUpToResult CountUpTo(int16_t value, int32_t batch_size); /// Return the max level used in this decoder. - int max_level() const { return max_level_; } + int16_t max_level() const; + + /// Return the number of values left to be decoded. Review Comment: ```suggestion /// Return the number of levels left to be decoded. ``` ########## cpp/src/parquet/column_reader.h: ########## @@ -351,105 +383,50 @@ class PARQUET_EXPORT RecordReader { virtual const void* ReadDictionary(int32_t* dictionary_length) = 0; /// \brief Decoded definition levels Review Comment: Can you expand the docstring to explain when this will return a valid pointer? ########## cpp/src/parquet/column_reader.cc: ########## @@ -91,137 +83,100 @@ inline void CheckNumberDecoded(int64_t number_decoded, int64_t expected) { } } -constexpr std::string_view kErrorRepDefLevelNotMatchesNumValues = - "Number of decoded rep / def levels do not match num_values in page header"; +/// True if a T can hold a U. +template <typename T, typename U> +inline constexpr bool can_hold_v = std::in_range<T>(std::numeric_limits<U>::min()) && + std::in_range<T>(std::numeric_limits<U>::max()); + +/// Clamp to a given type. +template <typename T, typename U> +constexpr T clamp_to(U val) { + if constexpr (can_hold_v<T, U>) { + return static_cast<T>(val); + } + constexpr U kMax = std::numeric_limits<T>::max(); + constexpr U kMin = std::numeric_limits<T>::min(); + return static_cast<T>(std::clamp<U>(val, kMin, kMax)); +} + +/// Return the int that can be represented by the other. +template <typename T, typename U> +using smallest_int_t = std::conditional_t<can_hold_v<T, U>, U, T>; + +/// Return the min in the smallest integer size of its inputs. +template <typename T, typename U> +constexpr auto narrow_min(T x, U y) -> smallest_int_t<T, U> { + using V = smallest_int_t<T, U>; + return std::min<V>(clamp_to<V>(x), clamp_to<V>(y)); +} + +/// Whether values of type T can be printed to `std::cout`. +template <typename T> +concept can_cout = requires(std::ostream& os, const T& value) { + os << value; +}; // NOLINT(readability/braces) } // namespace /****************** * 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); - } +// Replacement for legacy LevelDecoder +using NewLevelDecoder = PageLevelDecoder<>; Review Comment: Things named "New" get old quickly, can we choose another name perhaps? ########## cpp/src/parquet/column_reader.cc: ########## @@ -724,358 +690,816 @@ class SkippableTypedDecoder { } }; -// ---------------------------------------------------------------------- -// Impl base class for TypedColumnReader and RecordReader +/********************* + * ValueSinkCursor * + *********************/ -template <typename DType> -class ColumnReaderImplBase { +inline int64_t compute_capacity_pow2(int64_t capacity, int64_t size, int64_t extra_size) { Review Comment: Can we use the project's naming conventions, e.g. `ComputeCapacityPow2` or something? ########## cpp/src/parquet/column_reader_test.cc: ########## @@ -643,19 +650,60 @@ internal::LevelInfo ComputeLevelInfo(const ColumnDescriptor* descr) { } // namespace -using ReadDenseForNullable = bool; +struct RecordReaderTestParam { + bool read_dense_for_nullable; + bool allow_levels_optimization; +}; + +std::vector<RecordReaderTestParam> RecordReaderTestParams( + std::initializer_list<bool> allow_levels_optimizations) { + std::vector<RecordReaderTestParam> params; + for (bool read_dense_for_nullable : {true, false}) { + for (bool allow_levels_optimization : allow_levels_optimizations) { + params.push_back({ + .read_dense_for_nullable = read_dense_for_nullable, + .allow_levels_optimization = allow_levels_optimization, + }); + } + } + return params; +} + +std::string RecordReaderTestParamName( + const testing::TestParamInfo<RecordReaderTestParam>& info) { + return std::string(info.param.read_dense_for_nullable ? "Dense" : "Spaced") + + (info.param.allow_levels_optimization ? "LevelsOptimization" + : "NoLevelsOptimization"); +} + class RecordReaderPrimitiveTypeTest - : public ::testing::TestWithParam<ReadDenseForNullable> { + : public ::testing::TestWithParam<RecordReaderTestParam> { public: const int32_t kNullValue = -1; + bool read_dense_for_nullable() const { return GetParam().read_dense_for_nullable; } + + bool allow_levels_optimization() const { return GetParam().allow_levels_optimization; } + void Init(NodePtr column) { NodePtr root = GroupNode::Make("root", Repetition::REQUIRED, {column}); schema_descriptor_.Init(root); descr_ = schema_descriptor_.Column(0); - record_reader_ = internal::RecordReader::Make(descr_, ComputeLevelInfo(descr_), - ::arrow::default_memory_pool(), - /*read_dictionary=*/false, GetParam()); + record_reader_ = internal::RecordReader::Make({ + .descr = descr_, + .leaf_info = ComputeLevelInfo(descr_), + .read_dense_for_nullable = read_dense_for_nullable(), + .allow_levels_optimization = allow_levels_optimization(), + }); + } + + // Whether the record reader for this column exposes definition and repetition + // levels. The flat optional reader decodes definition levels straight into the + // validity bitmap of the records it produces, so it never materializes levels and + // reports none: `def_levels()` is null and both level counters stay at zero. + bool exposes_levels() const { Review Comment: Interesting, perhaps this can be a method on RecordReader instead? Something like `RecordReader::exposes_levels`. ########## cpp/src/parquet/column_reader.cc: ########## @@ -724,358 +690,816 @@ class SkippableTypedDecoder { } }; -// ---------------------------------------------------------------------- -// Impl base class for TypedColumnReader and RecordReader +/********************* + * ValueSinkCursor * + *********************/ -template <typename DType> -class ColumnReaderImplBase { +inline int64_t compute_capacity_pow2(int64_t capacity, int64_t size, int64_t extra_size) { + if (extra_size < 0) { + throw ParquetException("Negative size (corrupt file?)"); + } + int64_t target_size = -1; + if (AddWithOverflow(size, extra_size, &target_size)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + if (target_size >= (1LL << 62)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + if (capacity >= target_size) { + return capacity; + } + return bit_util::NextPower2(target_size); +} + +/// A simple cursor with a number of values and an available capacity. +/// +/// This is a reused foundation for creating data sinks. +/// Since decoders write to an already available buffer, we need to manually track a +/// capacity where the user is allowed to write (contrary to say `std::vector` where +/// it is UB to write in range `[size(), capacity()[`). +class ValueSinkCursor { public: - using T = typename DType::c_type; + int64_t capacity() const { return capacity_; } - ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) - : descr_(descr), - definition_level_decoder_(descr->max_definition_level()), - repetition_level_decoder_(descr->max_repetition_level()), - pool_(pool), - current_decoder_(pool) {} + int64_t values_count() const { return values_count_; } - virtual ~ColumnReaderImplBase() = default; + void set_values_count(int64_t vals) { values_count_ = vals; } - protected: - // Read up to batch_size values from the current data page into the - // pre-allocated memory T* - // - // @returns: the number of values read into the out buffer - int64_t ReadValues(int64_t batch_size, T* out) { - int64_t num_decoded = current_decoder_->Decode(out, static_cast<int>(batch_size)); - return num_decoded; + void increase_values_count(int64_t extra) { values_count_ += extra; } + + int64_t fit_capacity_for_extra(int64_t extra_values) { + auto new_capacity = compute_capacity_pow2(capacity_, values_count_, extra_values); + ARROW_DCHECK_GE(new_capacity, capacity()); + return std::exchange(capacity_, new_capacity); } - // Read up to batch_size values from the current data page into the - // pre-allocated memory T*, leaving spaces for null entries according - // to the def_levels. - // - // @returns: the number of values read into the out buffer - int64_t ReadValuesSpaced(int64_t batch_size, T* out, int64_t null_count, - uint8_t* valid_bits, int64_t valid_bits_offset) { - return current_decoder_->DecodeSpaced(out, static_cast<int>(batch_size), - static_cast<int>(null_count), valid_bits, - valid_bits_offset); + void reset() { + values_count_ = 0; + capacity_ = 0; } - // Read multiple definition levels into preallocated memory - // - // Returns the number of decoded definition levels - int64_t ReadDefinitionLevels(int64_t batch_size, int16_t* levels) { - if (max_def_level() == 0) { - return 0; + private: + int64_t values_count_ = 0; + int64_t capacity_ = 0; +}; + +/******************** + * DataSinkBuffer * + ********************/ + +template <typename T> +struct BytesCounterForValues { + /// Values are always fully written before being read, no need to zero new capacity. + static constexpr bool kZeroNewCapacity = false; + + static int64_t bytes_for_values(int64_t nitems) { + constexpr auto kValueByteSize = static_cast<int64_t>(sizeof(T)); + int64_t bytes = -1; + if (MultiplyWithOverflow(nitems, kValueByteSize, &bytes)) { + throw ParquetException("Total size of items too large"); } - return definition_level_decoder_.Decode(static_cast<int>(batch_size), levels); + return bytes; } +}; - bool HasNextInternal() { - // Either there is no data page available yet, or the data page has been - // exhausted - if (num_buffered_values_ == 0 || num_decoded_values_ == num_buffered_values_) { - if (!ReadNewPage() || num_buffered_values_ == 0) { - return false; - } +/// A base data sink that writes into an Arrow buffer. +template <typename T, typename BytesCounter = BytesCounterForValues<T>> +class DataSinkBuffer : private ValueSinkCursor { + public: + using value_type = T; + + using ValueSinkCursor::capacity; + using ValueSinkCursor::values_count; + + explicit DataSinkBuffer(std::shared_ptr<::arrow::ResizableBuffer> buffer) + : values_(std::move(buffer)) {} + + bool is_void() const { return values_ == nullptr; } + + value_type* data() const { + return is_void() ? nullptr : values_->mutable_data_as<value_type>(); + } + + /// Erase a range of element from the buffer. + void Erase(int64_t start, int64_t end) { + if (is_void() || start >= values_count() || start >= end) { + return; } - return true; + const auto last = std::min(end, values_count()); + const auto count = last - start; + std::copy(data() + last, write_start(), data() + start); + set_values_count(values_count() - count); } - // 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) { - return 0; + /// Transfer ownership of the values already decoded to the caller. + std::shared_ptr<ResizableBuffer> ReleaseValues(MemoryPool* pool) { + if (is_void()) { + return nullptr; } - return repetition_level_decoder_.Decode(static_cast<int>(batch_size), levels); + + auto result = values_; + const auto byte_count = BytesCounter::bytes_for_values(values_count()); + PARQUET_THROW_NOT_OK(result->Resize(byte_count, /*shrink_to_fit=*/true)); + values_ = AllocateBuffer(pool); + ValueSinkCursor::reset(); + return result; } - // Advance to the next data page - bool ReadNewPage() { - // Loop until we find the next data page. - while (true) { - current_page_ = pager_->NextPage(); - if (!current_page_) { - // EOS - return false; + /// Exponentially reserve more capacity if needed. + void ReserveValues(int64_t extra_values) { + const auto old_capacity = fit_capacity_for_extra(extra_values); + if (capacity() > old_capacity && !is_void()) { + const auto byte_count = BytesCounter::bytes_for_values(capacity()); + PARQUET_THROW_NOT_OK(values_->Resize(byte_count, /*shrink_to_fit=*/false)); + if constexpr (BytesCounter::kZeroNewCapacity) { + const auto old_byte_count = BytesCounter::bytes_for_values(old_capacity); + std::memset(data() + old_byte_count, 0, + static_cast<std::size_t>(byte_count - old_byte_count)); } + } + } - if (current_page_->type() == PageType::DICTIONARY_PAGE) { - ConfigureDictionary(static_cast<const DictionaryPage*>(current_page_.get())); - continue; - } else if (current_page_->type() == PageType::DATA_PAGE) { - const auto* page = static_cast<const DataPageV1*>(current_page_.get()); - const int64_t levels_byte_size = InitializeLevelDecoders( - *page, page->repetition_level_encoding(), page->definition_level_encoding()); - InitializeDataDecoder(*page, levels_byte_size); - return true; - } else if (current_page_->type() == PageType::DATA_PAGE_V2) { - const auto* page = static_cast<const DataPageV2*>(current_page_.get()); - int64_t levels_byte_size = InitializeLevelDecodersV2(*page); - InitializeDataDecoder(*page, levels_byte_size); - return true; + void ResetValues() { + // Resizing a never allocated buffer would make a zero-size allocation + if (!is_void() && values_count() > 0) { + PARQUET_THROW_NOT_OK(values_->Resize(0, /*shrink_to_fit=*/false)); + ValueSinkCursor::reset(); + } + } + + void DebugPrintState() { + const T* vals = data(); + for (int64_t i = 0; i < values_count(); ++i) { + if constexpr (can_cout<T>) { + std::cout << vals[i] << ' '; } else { - // We don't know what this page type is. We're allowed to skip non-data - // pages. - continue; + std::cout << "? "; } } - return true; } - void ConfigureDictionary(const DictionaryPage* page) { - int encoding = static_cast<int>(page->encoding()); - if (page->encoding() == Encoding::PLAIN_DICTIONARY || - page->encoding() == Encoding::PLAIN) { - encoding = static_cast<int>(Encoding::RLE_DICTIONARY); + protected: + /// Decode values contiguously into the sink buffer. + template <typename Int> + Int ReadFromCallback(auto&& decode, Int batch_size) { Review Comment: Is the `Int` templating actually required or is it only obscuring things? ########## cpp/src/parquet/column_reader.cc: ########## @@ -724,358 +690,816 @@ class SkippableTypedDecoder { } }; -// ---------------------------------------------------------------------- -// Impl base class for TypedColumnReader and RecordReader +/********************* + * ValueSinkCursor * + *********************/ -template <typename DType> -class ColumnReaderImplBase { +inline int64_t compute_capacity_pow2(int64_t capacity, int64_t size, int64_t extra_size) { + if (extra_size < 0) { + throw ParquetException("Negative size (corrupt file?)"); + } + int64_t target_size = -1; + if (AddWithOverflow(size, extra_size, &target_size)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + if (target_size >= (1LL << 62)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + if (capacity >= target_size) { + return capacity; + } + return bit_util::NextPower2(target_size); +} + +/// A simple cursor with a number of values and an available capacity. +/// +/// This is a reused foundation for creating data sinks. +/// Since decoders write to an already available buffer, we need to manually track a +/// capacity where the user is allowed to write (contrary to say `std::vector` where +/// it is UB to write in range `[size(), capacity()[`). +class ValueSinkCursor { public: - using T = typename DType::c_type; + int64_t capacity() const { return capacity_; } - ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) - : descr_(descr), - definition_level_decoder_(descr->max_definition_level()), - repetition_level_decoder_(descr->max_repetition_level()), - pool_(pool), - current_decoder_(pool) {} + int64_t values_count() const { return values_count_; } - virtual ~ColumnReaderImplBase() = default; + void set_values_count(int64_t vals) { values_count_ = vals; } - protected: - // Read up to batch_size values from the current data page into the - // pre-allocated memory T* - // - // @returns: the number of values read into the out buffer - int64_t ReadValues(int64_t batch_size, T* out) { - int64_t num_decoded = current_decoder_->Decode(out, static_cast<int>(batch_size)); - return num_decoded; + void increase_values_count(int64_t extra) { values_count_ += extra; } + + int64_t fit_capacity_for_extra(int64_t extra_values) { Review Comment: `FitCapacityForExtra`? ########## cpp/src/parquet/column_reader.cc: ########## @@ -724,358 +690,816 @@ class SkippableTypedDecoder { } }; -// ---------------------------------------------------------------------- -// Impl base class for TypedColumnReader and RecordReader +/********************* + * ValueSinkCursor * + *********************/ -template <typename DType> -class ColumnReaderImplBase { +inline int64_t compute_capacity_pow2(int64_t capacity, int64_t size, int64_t extra_size) { + if (extra_size < 0) { + throw ParquetException("Negative size (corrupt file?)"); + } + int64_t target_size = -1; + if (AddWithOverflow(size, extra_size, &target_size)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + if (target_size >= (1LL << 62)) { + throw ParquetException("Allocation size too large (corrupt file?)"); + } + if (capacity >= target_size) { + return capacity; + } + return bit_util::NextPower2(target_size); +} + +/// A simple cursor with a number of values and an available capacity. +/// +/// This is a reused foundation for creating data sinks. +/// Since decoders write to an already available buffer, we need to manually track a +/// capacity where the user is allowed to write (contrary to say `std::vector` where +/// it is UB to write in range `[size(), capacity()[`). +class ValueSinkCursor { public: - using T = typename DType::c_type; + int64_t capacity() const { return capacity_; } - ColumnReaderImplBase(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) - : descr_(descr), - definition_level_decoder_(descr->max_definition_level()), - repetition_level_decoder_(descr->max_repetition_level()), - pool_(pool), - current_decoder_(pool) {} + int64_t values_count() const { return values_count_; } - virtual ~ColumnReaderImplBase() = default; + void set_values_count(int64_t vals) { values_count_ = vals; } - protected: - // Read up to batch_size values from the current data page into the - // pre-allocated memory T* - // - // @returns: the number of values read into the out buffer - int64_t ReadValues(int64_t batch_size, T* out) { - int64_t num_decoded = current_decoder_->Decode(out, static_cast<int>(batch_size)); - return num_decoded; + void increase_values_count(int64_t extra) { values_count_ += extra; } + + int64_t fit_capacity_for_extra(int64_t extra_values) { + auto new_capacity = compute_capacity_pow2(capacity_, values_count_, extra_values); + ARROW_DCHECK_GE(new_capacity, capacity()); + return std::exchange(capacity_, new_capacity); } - // Read up to batch_size values from the current data page into the - // pre-allocated memory T*, leaving spaces for null entries according - // to the def_levels. - // - // @returns: the number of values read into the out buffer - int64_t ReadValuesSpaced(int64_t batch_size, T* out, int64_t null_count, - uint8_t* valid_bits, int64_t valid_bits_offset) { - return current_decoder_->DecodeSpaced(out, static_cast<int>(batch_size), - static_cast<int>(null_count), valid_bits, - valid_bits_offset); + void reset() { Review Comment: `Reset` -- 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]
