wgtmac commented on code in PR #49855:
URL: https://github.com/apache/arrow/pull/49855#discussion_r3670746814
##########
cpp/src/arrow/io/caching.cc:
##########
@@ -151,11 +152,19 @@ struct ReadRangeCache::Impl {
IOContext ctx;
CacheOptions options;
- // Ordered by offset (so as to find a matching region by binary search)
- std::vector<RangeCacheEntry> entries;
+ // Ordered by offset (so as to find a matching region by binary search).
+ // Mutation of `entries` and of individual entries' futures must be
+ // serialized via `entry_mutex`. Every public method that touches either
+ // acquires the mutex before delegating to the protected *Locked helpers
+ // below, so both the eager and lazy variants are safe to call concurrently
+ // from multiple threads.
+ std::deque<RangeCacheEntry> entries;
Review Comment:
```suggestion
std::deque<RangeCacheEntry> entries; // GUARDED_BY(entry_mutex)
```
No strong opinion. I would prefer simplifying comments like this.
##########
cpp/src/arrow/io/caching.cc:
##########
@@ -151,11 +152,19 @@ struct ReadRangeCache::Impl {
IOContext ctx;
CacheOptions options;
- // Ordered by offset (so as to find a matching region by binary search)
- std::vector<RangeCacheEntry> entries;
+ // Ordered by offset (so as to find a matching region by binary search).
+ // Mutation of `entries` and of individual entries' futures must be
+ // serialized via `entry_mutex`. Every public method that touches either
+ // acquires the mutex before delegating to the protected *Locked helpers
+ // below, so both the eager and lazy variants are safe to call concurrently
+ // from multiple threads.
+ std::deque<RangeCacheEntry> entries;
+ std::mutex entry_mutex;
virtual ~Impl() = default;
+ // -- Polymorphic hooks. Always called with entry_mutex held. --
Review Comment:
```suggestion
```
These comments look more like notes while doing the implementation.
##########
cpp/src/arrow/io/caching.cc:
##########
@@ -172,46 +181,62 @@ struct ReadRangeCache::Impl {
return new_entries;
}
- // Add the given ranges to the cache, coalescing them where possible
- virtual Status Cache(std::vector<ReadRange> ranges) {
+ // -- Public entry points (acquire entry_mutex, then delegate). --
+
Review Comment:
```suggestion
```
##########
cpp/src/arrow/io/caching.cc:
##########
@@ -224,53 +249,80 @@ struct ReadRangeCache::Impl {
++num_prefetched;
}
}
- return SliceBuffer(std::move(buf), range.offset - it->range.offset,
range.length);
}
- return Status::Invalid("ReadRangeCache did not find matching cache entry");
+ // Drop the lock before blocking on the I/O future so other threads can
+ // still do lookups while a previously queued read is in flight.
+ ARROW_ASSIGN_OR_RAISE(auto buf, fut.result());
+ return SliceBuffer(std::move(buf), slice_offset, range.length);
}
- virtual Future<> Wait() {
+ Future<> Wait() {
std::vector<Future<>> futures;
- for (auto& entry : entries) {
- futures.emplace_back(MaybeRead(&entry));
+ {
+ std::unique_lock<std::mutex> guard(entry_mutex);
+ futures.reserve(entries.size());
+ for (auto& entry : entries) {
+ futures.emplace_back(MaybeRead(&entry));
+ }
}
return AllComplete(futures);
}
+ // Cached ranges are sorted and non-overlapping, so entries ending at or
+ // before `end_offset` form a prefix. Keep a straddling entry.
+ int64_t EvictEntriesBefore(int64_t end_offset) {
+ std::unique_lock<std::mutex> guard(entry_mutex);
+ int64_t n_evicted = 0;
+ while (!entries.empty()) {
+ const auto& range = entries.front().range;
+ if (range.offset + range.length > end_offset) {
+ break;
+ }
+ entries.pop_front();
+ ++n_evicted;
+ }
+ return n_evicted;
Review Comment:
Would it be more useful to return total evicted memory?
##########
cpp/src/parquet/arrow/reader.cc:
##########
@@ -1275,10 +1239,31 @@
FileReaderImpl::GetRecordBatchGenerator(std::shared_ptr<FileReader> reader,
reader_properties_.cache_options());
END_PARQUET_CATCH_EXCEPTIONS
}
+ // GH-39808: evict each row group's bytes as the decoded prefix advances, so
+ // memory stays bounded. Only this read-once path evicts, so PreBuffer()'s
+ // contract is unchanged for other callers.
+ std::shared_ptr<ReadCacheEvictionState> eviction_state;
+ if (reader_properties_.pre_buffer() && !column_indices.empty() &&
+ !row_group_indices.empty()) {
+ const int64_t kNoMoreRanges = std::numeric_limits<int64_t>::max();
+ std::vector<int64_t> evict_before(row_group_indices.size() + 1,
kNoMoreRanges);
+ for (int64_t i = static_cast<int64_t>(row_group_indices.size()) - 1; i >=
0; --i) {
Review Comment:
I think it is worth adding a comment to the API w.r.t. this contract.
##########
cpp/src/parquet/file_reader.h:
##########
@@ -201,6 +201,12 @@ class PARQUET_EXPORT ParquetFileReader {
const ::arrow::io::IOContext& ctx,
const ::arrow::io::CacheOptions& options);
+ /// \brief Release cached bytes (from PreBuffer()) ending at or before
+ /// `end_offset`. Call once those row groups are decoded; later reads of
evicted
+ /// ranges fall back to the source file. No-op (returns 0) if PreBuffer()
was not
+ /// called.
+ int64_t EvictPreBufferedDataBefore(int64_t end_offset);
Review Comment:
Why do we need to return a `int64_t`? It would be good to document its
meaning or just make it `void` if not very useful.
##########
cpp/src/parquet/file_reader.cc:
##########
@@ -613,6 +626,7 @@ class SerializedFile : public ParquetFileReader::Contents {
ReaderProperties properties_;
std::shared_ptr<PageIndexReader> page_index_reader_;
std::unique_ptr<BloomFilterReader> bloom_filter_reader_;
+
Review Comment:
```suggestion
```
##########
cpp/src/parquet/arrow/reader_internal.h:
##########
@@ -133,5 +135,50 @@ Status
TransferColumnData(::parquet::internal::RecordReader* reader,
const ColumnDescriptor* descr, const ReaderContext*
ctx,
std::shared_ptr<::arrow::ChunkedArray>* out);
+// ----------------------------------------------------------------------
+// Pre-buffer eviction
+
+// GH-39808: as row groups finish decoding (possibly out of order under
+// readahead), advance a watermark over the leading run of completed ones and
+// evict cache entries ending before the lowest byte any remaining one needs.
+class ReadCacheEvictionState {
+ public:
+ // evict_before_offsets[i] = lowest byte offset row groups i..n-1 (in
+ // generator order) still need; evict_before_offsets[n] == INT64_MAX.
+ explicit ReadCacheEvictionState(std::vector<int64_t> evict_before_offsets)
+ : evict_before_offsets_(std::move(evict_before_offsets)),
+ completed_(evict_before_offsets_.size() - 1, false) {}
+
+ void RowGroupDecoded(ParquetFileReader* reader, size_t row_group_index) {
+ if (auto evict_before = MarkDecodedAndGetEvictOffset(row_group_index)) {
+ reader->EvictPreBufferedDataBefore(*evict_before);
Review Comment:
As I've said earlier, we cannot evict them by default. There are use cases
where users create a file reader by prebuffering some row groups and repeatedly
call `GetRecordBatchGenerator` to read the same row group for more than 1 time.
In that case, followup reads may have issues. We can only enable it by default
to the Dataset parquet reader where each row group is read only once.
##########
cpp/src/parquet/file_reader.cc:
##########
@@ -247,9 +247,15 @@ class SerializedRowGroup : public RowGroupReader::Contents
{
::arrow::bit_util::GetBit(prebuffered_column_chunks_bitmap_->data(),
i)) {
// PARQUET-1698: if read coalescing is enabled, read from pre-buffered
// segments.
- PARQUET_ASSIGN_OR_THROW(auto buffer, cached_source_->Read(col_range));
- stream = std::make_shared<::arrow::io::BufferReader>(buffer);
- } else {
+ auto buffer = cached_source_->Read(col_range);
Review Comment:
This looks hacky to fallback on a specific error code and may hide its
original error. Is it better to introduce a function like
`Result<std::optional<std::shared_ptr<Buffer>>> ReadIfCached(ReadRange range)`
to explicitly indicate that cache miss is not an error?
##########
cpp/src/arrow/io/caching.cc:
##########
@@ -172,46 +181,62 @@ struct ReadRangeCache::Impl {
return new_entries;
}
- // Add the given ranges to the cache, coalescing them where possible
- virtual Status Cache(std::vector<ReadRange> ranges) {
+ // -- Public entry points (acquire entry_mutex, then delegate). --
+
+ // Add the given ranges to the cache, coalescing them where possible.
+ Status Cache(std::vector<ReadRange> ranges) {
ARROW_ASSIGN_OR_RAISE(
ranges, internal::CoalesceReadRanges(std::move(ranges),
options.hole_size_limit,
options.range_size_limit));
- std::vector<RangeCacheEntry> new_entries = MakeCacheEntries(ranges);
- // Add new entries, themselves ordered by offset
- if (entries.size() > 0) {
- std::vector<RangeCacheEntry> merged(entries.size() + new_entries.size());
- std::merge(entries.begin(), entries.end(), new_entries.begin(),
new_entries.end(),
- merged.begin());
- entries = std::move(merged);
- } else {
- entries = std::move(new_entries);
+ Status st;
+ {
+ std::unique_lock<std::mutex> guard(entry_mutex);
+ std::vector<RangeCacheEntry> new_entries = MakeCacheEntries(ranges);
+ // Add new entries, themselves ordered by offset
+ if (entries.size() > 0) {
+ std::deque<RangeCacheEntry> merged(entries.size() +
new_entries.size());
+ std::merge(entries.begin(), entries.end(), new_entries.begin(),
new_entries.end(),
+ merged.begin());
+ entries = std::move(merged);
+ } else {
+ for (auto& entry : new_entries) {
+ entries.push_back(std::move(entry));
+ }
+ }
Review Comment:
```suggestion
std::deque<RangeCacheEntry> merged;
std::merge(std::make_move_iterator(entries.begin()),
std::make_move_iterator(entries.end()),
std::make_move_iterator(new_entries.begin()),
std::make_move_iterator(new_entries.end()),
std::back_inserter(merged));
entries.swap(merged);
```
This might be more efficient.
--
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]
