AntoinePrv commented on code in PR #50629:
URL: https://github.com/apache/arrow/pull/50629#discussion_r3656400839
##########
cpp/src/parquet/column_reader.cc:
##########
@@ -726,393 +917,702 @@ 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);
+}
+
+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;
+ 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);
+ int64_t reset_capacity() { return std::exchange(capacity_, 0); }
+
+ private:
+ int64_t values_count_ = 0;
+ int64_t capacity_ = 0;
+};
+
+/*********************
+ * ValueSinkBuffer *
+ *********************/
+
+template <typename T>
+class ValueSinkBuffer : private ValueSinkCursor {
+ public:
+ using value_type = T;
+
+ using ValueSinkCursor::capacity;
+ using ValueSinkCursor::values_count;
+
+ explicit ValueSinkBuffer(MemoryPool* pool) : values_(AllocateBuffer(pool)) {}
+
+ value_type* data() const { return values_->mutable_data_as<value_type>(); }
+
+ void OnNewDictionary(auto& /* decoder */) {}
+
+ [[nodiscard]] auto ReadValuesDense(auto& decoder, int32_t batch_size) {
+ ReserveValues(batch_size);
+ const auto decoded = decoder.Decode(write_start(), batch_size);
+ set_values_count(values_count() + batch_size);
+ return decoded;
}
- // 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;
- }
- return definition_level_decoder_.Decode(static_cast<int>(batch_size),
levels);
+ [[nodiscard]] auto ReadValuesSpaced(auto& decoder, int32_t batch_size,
+ int32_t null_count, const uint8_t*
valid_bits,
+ int64_t valid_bits_offset) {
+ ReserveValues(batch_size);
+ const auto decoded = decoder.DecodeSpaced(write_start(), batch_size,
null_count,
+ valid_bits, valid_bits_offset);
+ set_values_count(values_count() + batch_size);
+ return decoded;
}
- 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;
- }
- }
- return true;
+ std::shared_ptr<ResizableBuffer> ReleaseValues(MemoryPool* pool) {
+ // TODO should we set values_written to zero?
+ auto result = values_;
+ const auto byte_count = bytes_for_values(values_count());
+ PARQUET_THROW_NOT_OK(result->Resize(byte_count, /*shrink_to_fit=*/true));
+ values_ = AllocateBuffer(pool);
+ reset_capacity();
+ return result;
}
- // 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;
+ void ReserveValues(int64_t extra_values) {
+ const auto old_capacity = fit_capacity_for_extra(extra_values);
+ if (capacity() > old_capacity) {
+ const auto byte_count = bytes_for_values(capacity());
+ PARQUET_THROW_NOT_OK(values_->Resize(byte_count,
/*shrink_to_fit=*/false));
}
- return repetition_level_decoder_.Decode(static_cast<int>(batch_size),
levels);
}
- // 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;
- }
+ void ResetValues() {
+ if (values_count() > 0) {
+ PARQUET_THROW_NOT_OK(values_->Resize(0, /*shrink_to_fit=*/false));
+ ValueSinkCursor::operator=({});
+ }
+ }
- 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 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);
- }
+ private:
+ std::shared_ptr<::arrow::ResizableBuffer> values_;
- auto it = decoders_.find(encoding);
- if (it != decoders_.end()) {
- throw ParquetException("Column cannot have more than one dictionary.");
+ static int64_t bytes_for_values(int64_t nitems) {
+ constexpr auto kValueByteSize = static_cast<int64_t>(sizeof(value_type));
+ int64_t bytes = -1;
+ if (MultiplyWithOverflow(nitems, kValueByteSize, &bytes)) {
+ throw ParquetException("Total size of items too large");
}
+ return bytes;
+ }
- if (page->encoding() == Encoding::PLAIN_DICTIONARY ||
- page->encoding() == Encoding::PLAIN) {
- auto dictionary = MakeTypedDecoder<DType>(Encoding::PLAIN, descr_,
pool_);
- dictionary->SetData(page->num_values(), page->data(), page->size());
+ value_type* write_start() { return data() + values_count(); }
+};
- // The dictionary is fully decoded during DictionaryDecoder::Init, so the
- // DictionaryPage buffer is no longer required after this step
- //
- // TODO(wesm): investigate whether this all-or-nothing decoding of the
- // dictionary makes sense and whether performance can be improved
+/************************
+ * ValiditySinkBuffer *
+ ************************/
- std::unique_ptr<DictDecoder<DType>> decoder =
MakeDictDecoder<DType>(descr_, pool_);
- decoder->SetDict(dictionary.get());
- decoders_[encoding] =
-
std::unique_ptr<DecoderType>(dynamic_cast<DecoderType*>(decoder.release()));
- } else {
- ParquetException::NYI("only plain dictionary encoding has been
implemented");
- }
+class ValiditySinkBuffer : private ValueSinkCursor {
+ public:
+ using ValueSinkCursor::capacity;
+ using ValueSinkCursor::values_count;
+
+ ValiditySinkBuffer() = default;
+
+ explicit ValiditySinkBuffer(MemoryPool* pool) :
values_(AllocateBuffer(pool)) {}
+
+ uint8_t* data() const { return values_->mutable_data_as<uint8_t>(); }
- new_dictionary_ = true;
- current_decoder_.SetDecoder(decoders_[encoding].get());
- ARROW_DCHECK(current_decoder_);
+ template <typename Int>
+ struct ReadResult {
+ Int values_read = 0;
+ Int null_count = 0;
+ };
+
+ ReadResult<int64_t> ReadFromLevels(const int16_t* def_levels, int64_t
num_def_levels,
+ const internal::LevelInfo& level_info) {
+ // At most one validity bit is written per definition level.
+ ReserveValues(num_def_levels);
+ internal::ValidityBitmapInputOutput validity_io{};
+ validity_io.values_read_upper_bound = num_def_levels;
+ validity_io.valid_bits = data();
+ validity_io.valid_bits_offset = values_count();
+ DefLevelsToBitmap(def_levels, num_def_levels, level_info, &validity_io);
+ ARROW_DCHECK_GE(validity_io.values_read, 0);
+ ARROW_DCHECK_GE(validity_io.null_count, 0);
+
+ // Advance by the number of leaf values (one validity bit each), not by the
+ // number of definition levels: for repeated/nested columns some def levels
+ // describe empty or absent lists that produce no leaf value, so
+ // values_read <= num_def_levels.
+ set_values_count(values_count() + validity_io.values_read);
+ return {
+ .values_read = validity_io.values_read,
+ .null_count = validity_io.null_count,
+ };
}
- // Initialize repetition and definition level decoders on the next data page.
+ ReadResult<int32_t> ReadFromDecoder(LevelToBitmapDecoder& decoder, int32_t
batch_size) {
+ using BitmapSpanMut = LevelToBitmapDecoder::BitmapSpanMut;
- // If the data page includes repetition and definition levels, we
- // initialize the level decoders and return the number of encoded level
bytes.
- // The return value helps determine the number of bytes in the encoded data.
- int64_t InitializeLevelDecoders(const DataPage& page,
- Encoding::type repetition_level_encoding,
- Encoding::type definition_level_encoding) {
- // Read a data page.
- num_buffered_values_ = page.num_values();
+ ReserveValues(batch_size);
+ const auto write_pos = BitmapSpanMut(data() + values_count() / 8,
+ static_cast<int32_t>(values_count() %
8));
+ const auto decoded = decoder.Decode(batch_size, write_pos);
+ const auto non_null =
+ ::arrow::internal::CountSetBits(write_pos.data(),
write_pos.bit_start(), decoded);
- // Have not decoded any values from the data page yet
- num_decoded_values_ = 0;
+ ARROW_DCHECK_LE(non_null, decoded);
+ set_values_count(values_count() + decoded);
+ return {
+ .values_read = decoded,
+ .null_count = static_cast<int32_t>(decoded - non_null),
+ };
+ }
- const uint8_t* buffer = page.data();
- int32_t levels_byte_size = 0;
- int32_t max_size = page.size();
+ std::shared_ptr<ResizableBuffer> ReleaseValues(MemoryPool* pool) {
+ // TODO should we set values_written to zero?
+ auto result = values_;
+ const auto byte_count = bytes_for_values(values_count());
+ PARQUET_THROW_NOT_OK(result->Resize(byte_count, /*shrink_to_fit=*/true));
+ values_ = AllocateBuffer(pool);
+ reset_capacity();
+ return result;
+ }
- // 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) {
- int32_t rep_levels_bytes = repetition_level_decoder_.SetData(
- 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;
- max_size -= rep_levels_bytes;
+ void ReserveValues(int64_t extra_values) {
+ const auto old_capacity = fit_capacity_for_extra(extra_values);
+ if (capacity() > old_capacity) {
+ const auto byte_count = bytes_for_values(capacity());
+ PARQUET_THROW_NOT_OK(values_->Resize(byte_count,
/*shrink_to_fit=*/false));
+ // Zero the newly grown region so that appending bits at a
non-byte-aligned
+ // offset never reads uninitialized memory (avoids Valgrind/MSAN
warnings).
+ const auto old_byte_count = bytes_for_values(old_capacity);
+ std::memset(data() + old_byte_count, 0,
+ static_cast<std::size_t>(byte_count - old_byte_count));
}
- // TODO figure a way to set max_def_level_ to 0
- // if the initial value is invalid
+ }
- // Init definition levels
- if (max_def_level() > 0) {
- int32_t def_levels_bytes = definition_level_decoder_.SetData(
- 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;
+ void ResetValues() {
+ if (values_count() > 0) {
+ PARQUET_THROW_NOT_OK(values_->Resize(0, /*shrink_to_fit=*/false));
+ ValueSinkCursor::operator=({});
}
-
- return levels_byte_size;
}
- int64_t InitializeLevelDecodersV2(const DataPageV2& page) {
- // Read a data page.
- num_buffered_values_ = page.num_values();
+ private:
+ std::shared_ptr<::arrow::ResizableBuffer> values_;
- // Have not decoded any values from the data page yet
- num_decoded_values_ = 0;
- const uint8_t* buffer = page.data();
+ static int64_t bytes_for_values(int64_t num_values) {
+ return bit_util::BytesForBits(num_values);
+ }
+};
- const int64_t total_levels_length =
- static_cast<int64_t>(page.repetition_levels_byte_length()) +
- page.definition_levels_byte_length();
+/***********************
+ * ColumnChunkReader *
+ ***********************/
- if (total_levels_length > page.size()) {
- throw ParquetException("Data page too small for levels (corrupt
header?)");
- }
+/// Initialize repetition and definition level decoders on the given data page.
+///
+/// If the data page includes repetition and definition levels, we initialize
the level
+/// decoders and return the number of encoded level bytes.
+/// The return value helps determine the number of bytes in the encoded data.
+int64_t InitializeV1Levels(const DataPageV1& page, auto& def_dec, auto&
rep_dec) {
+ const auto num_values = static_cast<int>(page.num_values());
- if (max_rep_level() > 0) {
- repetition_level_decoder_.SetDataV2(page.repetition_levels_byte_length(),
- max_rep_level(),
-
static_cast<int>(num_buffered_values_), buffer);
- }
- // ARROW-17453: Even if max_rep_level_ is 0, there may still be
- // repetition level bytes written and/or reported in the header by
- // some writers (e.g. Athena)
- buffer += page.repetition_levels_byte_length();
+ const uint8_t* buffer = page.data();
+ int32_t levels_byte_size = 0;
+ int32_t max_size = page.size();
- if (max_def_level() > 0) {
- definition_level_decoder_.SetDataV2(page.definition_levels_byte_length(),
- max_def_level(),
-
static_cast<int>(num_buffered_values_), buffer);
- }
+ if (const auto max_rep_lvl = rep_dec.max_level(); max_rep_lvl > 0) {
+ const int32_t rep_levels_bytes = rep_dec.SetData(
+ page.repetition_level_encoding(), max_rep_lvl, num_values, buffer,
max_size);
+ buffer += rep_levels_bytes;
+ levels_byte_size += rep_levels_bytes;
+ max_size -= rep_levels_bytes;
+ }
- return total_levels_length;
+ if (const auto max_def_lvl = def_dec.max_level(); max_def_lvl > 0) {
+ const int32_t def_levels_bytes = def_dec.SetData(
+ page.definition_level_encoding(), max_def_lvl, num_values, buffer,
max_size);
+ levels_byte_size += def_levels_bytes;
+ max_size -= def_levels_bytes;
}
- // Get a decoder object for this page or create a new decoder if this is the
- // first page with this encoding.
- void InitializeDataDecoder(const DataPage& page, int64_t levels_byte_size) {
- const uint8_t* buffer = page.data() + levels_byte_size;
- const int64_t data_size = page.size() - levels_byte_size;
+ return levels_byte_size;
+}
- if (data_size < 0) {
- throw ParquetException("Page smaller than size of encoded levels");
- }
+/// Initialize repetition and definition level decoders on the given data page.
+///
+/// If the data page includes repetition and definition levels, we initialize
the level
+/// decoders and return the number of encoded level bytes.
+/// The return value helps determine the number of bytes in the encoded data.
+int64_t InitializeV2Levels(const DataPageV2& page, auto& def_dec, auto&
rep_dec) {
+ const auto num_values = static_cast<int>(page.num_values());
- Encoding::type encoding = page.encoding();
- if (IsDictionaryIndexEncoding(encoding)) {
- // Normalizing the PLAIN_DICTIONARY to RLE_DICTIONARY encoding
- // in decoder.
- encoding = Encoding::RLE_DICTIONARY;
- }
+ const int64_t total_levels_length =
+ static_cast<int64_t>(page.repetition_levels_byte_length()) +
+ page.definition_levels_byte_length();
+ if (total_levels_length > page.size()) {
+ throw ParquetException("Data page too small for levels (corrupt header?)");
+ }
- auto it = decoders_.find(static_cast<int>(encoding));
- if (it != decoders_.end()) {
- ARROW_DCHECK(it->second.get() != nullptr);
- current_decoder_.SetDecoder(it->second.get());
- } else {
- switch (encoding) {
- case Encoding::PLAIN:
- case Encoding::BYTE_STREAM_SPLIT:
- case Encoding::RLE:
- case Encoding::DELTA_BINARY_PACKED:
- case Encoding::DELTA_BYTE_ARRAY:
- case Encoding::DELTA_LENGTH_BYTE_ARRAY: {
- auto decoder = MakeTypedDecoder<DType>(encoding, descr_, pool_);
- current_decoder_.SetDecoder(decoder.get());
- decoders_[static_cast<int>(encoding)] = std::move(decoder);
- break;
- }
+ const uint8_t* buffer = page.data();
- case Encoding::RLE_DICTIONARY:
- throw ParquetException("Dictionary page must be before data page.");
+ if (const auto max_rep_lvl = rep_dec.max_level(); max_rep_lvl > 0) {
+ rep_dec.SetDataV2(page.repetition_levels_byte_length(), max_rep_lvl,
num_values,
+ buffer);
+ }
+ // ARROW-17453: Even if max_rep_level_ is 0, there may still be
+ // repetition level bytes written and/or reported in the header by
+ // some writers (e.g. Athena)
+ buffer += page.repetition_levels_byte_length();
- default:
- throw ParquetException("Unknown encoding type.");
- }
- }
- current_encoding_ = encoding;
- current_decoder_->SetData(static_cast<int>(num_buffered_values_), buffer,
- static_cast<int>(data_size));
+ if (const auto max_def_lvl = def_dec.max_level(); max_def_lvl > 0) {
+ def_dec.SetDataV2(page.definition_levels_byte_length(), max_def_lvl,
num_values,
+ buffer);
}
- // Available values in the current data page, value includes repeated values
- // and nulls.
- int64_t available_values_current_page() const {
- return num_buffered_values_ - num_decoded_values_;
+ return total_levels_length;
+}
+
+/// Read through the multiple pages of a column chunk.
+template <typename Traits>
+class ColumnChunkReader {
+ public:
+ using DType = typename Traits::DType;
+ using DefLevelDecoder = typename Traits::DefLevelDecoder;
+ using RepLevelDecoder = typename Traits::RepLevelDecoder;
Review Comment:
Looking at `column_reader.cc.o` in release mode it goes from 720K in main to
2.1M in this PR (but it's not the worst offender).
We could first remove the template inheritance in the array and binary and
readers that is not as critical if we want to save some space.
--
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]
