AntoinePrv commented on code in PR #50629:
URL: https://github.com/apache/arrow/pull/50629#discussion_r4334900323
##########
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:
Hmm, either way. It is called with both `int32_t` and `int64_t`.
`int64_t` could fit all use case, but callers with an `int32_t` input would
then get an `int64_t` back and need to `static_cast` back to `int32_t`. So that
makes a simpler function but a more messy call.
--
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]
