diff --git a/cpp/src/arrow/dataset/file_parquet.cc b/cpp/src/arrow/dataset/file_parquet.cc index ba0e93f09d40..7ecf9e6a563f 100644 --- a/cpp/src/arrow/dataset/file_parquet.cc +++ b/cpp/src/arrow/dataset/file_parquet.cc @@ -130,6 +130,8 @@ parquet::ArrowReaderProperties MakeArrowReaderProperties( // Must be set here since the sync ScanTask handles pre-buffering itself arrow_properties.set_pre_buffer( parquet_scan_options.arrow_reader_properties->pre_buffer()); + // Dataset scans consume each row group once, so cached ranges can be released. + arrow_properties.set_auto_evict_read_cache(true); arrow_properties.set_cache_options( parquet_scan_options.arrow_reader_properties->cache_options()); arrow_properties.set_io_context( diff --git a/cpp/src/arrow/io/caching.cc b/cpp/src/arrow/io/caching.cc index 41fdd9f78108..1a7b383dc94a 100644 --- a/cpp/src/arrow/io/caching.cc +++ b/cpp/src/arrow/io/caching.cc @@ -18,7 +18,10 @@ #include #include #include +#include +#include #include +#include #include #include @@ -151,12 +154,12 @@ struct ReadRangeCache::Impl { IOContext ctx; CacheOptions options; - // Ordered by offset (so as to find a matching region by binary search) - std::vector entries; + std::deque entries; // GUARDED_BY(entry_mutex) + std::mutex entry_mutex; virtual ~Impl() = default; - // Get the future corresponding to a range + // Get the future corresponding to a range. Called with entry_mutex held. virtual Future> MaybeRead(RangeCacheEntry* entry) { return entry->future; } @@ -173,23 +176,26 @@ struct ReadRangeCache::Impl { return new_entries; } - // Add the given ranges to the cache, coalescing them where possible - virtual Status Cache(std::vector ranges) { + // Add the given ranges to the cache, coalescing them where possible. + Status Cache(std::vector ranges) { ARROW_ASSIGN_OR_RAISE( ranges, internal::CoalesceReadRanges(std::move(ranges), options.hole_size_limit, options.range_size_limit)); std::vector new_entries = MakeCacheEntries(ranges); - // Add new entries, themselves ordered by offset - if (entries.size() > 0) { - std::vector 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 guard(entry_mutex); + std::deque 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); } - // Prefetch immediately, regardless of executor availability, if possible - auto st = file->WillNeed(ranges); + // Prefetch immediately, regardless of executor availability, if possible. + // Do this outside the lock: WillNeed() may block on an mmap advise / I/O + // hint and we don't want to serialize concurrent Reads on it. + st = file->WillNeed(ranges); // As this is optimisation only, I/O failures should not be treated as fatal if (st.IsIOError()) { return Status::OK(); @@ -197,22 +203,28 @@ struct ReadRangeCache::Impl { return st; } - // Read the given range from the cache, blocking if needed. Cannot read a range - // that spans cache entries. - virtual Result> Read(ReadRange range) { + // Read the given range from the cache, blocking if needed. Cannot read a + // range that spans cache entries. + Result>> ReadIfCached(ReadRange range) { if (range.length == 0) { static const uint8_t byte = 0; - return std::make_shared(&byte, 0); + return std::make_optional(std::make_shared(&byte, 0)); } - const auto it = std::lower_bound( - entries.begin(), entries.end(), range, - [](const RangeCacheEntry& entry, const ReadRange& range) { - return entry.range.offset + entry.range.length < range.offset + range.length; - }); - if (it != entries.end() && it->range.Contains(range)) { - auto fut = MaybeRead(&*it); - ARROW_ASSIGN_OR_RAISE(auto buf, fut.result()); + Future> fut; + int64_t slice_offset = 0; + { + std::unique_lock guard(entry_mutex); + const auto it = std::lower_bound( + entries.begin(), entries.end(), range, + [](const RangeCacheEntry& entry, const ReadRange& range) { + return entry.range.offset + entry.range.length < range.offset + range.length; + }); + if (it == entries.end() || !it->range.Contains(range)) { + return std::nullopt; + } + fut = MaybeRead(&*it); + slice_offset = range.offset - it->range.offset; if (options.lazy && options.prefetch_limit > 0) { int64_t num_prefetched = 0; for (auto next_it = it + 1; @@ -226,53 +238,78 @@ 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 std::make_optional(SliceBuffer(std::move(buf), slice_offset, range.length)); + } + + Result> Read(ReadRange range) { + ARROW_ASSIGN_OR_RAISE(auto buffer, ReadIfCached(range)); + if (!buffer) { + return Status::Invalid("ReadRangeCache did not find matching cache entry"); + } + return std::move(*buffer); } - virtual Future<> Wait() { + Future<> Wait() { std::vector> futures; - for (auto& entry : entries) { - futures.emplace_back(MaybeRead(&entry)); + { + std::unique_lock 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. + void EvictEntriesBefore(int64_t end_offset) { + std::unique_lock guard(entry_mutex); + const auto first_kept = std::find_if( + entries.cbegin(), entries.cend(), [end_offset](const RangeCacheEntry& entry) { + return entry.range.offset + entry.range.length > end_offset; + }); + entries.erase(entries.cbegin(), first_kept); + } + // Return a Future that completes when the given ranges have been read. - virtual Future<> WaitFor(std::vector ranges) { + Future<> WaitFor(std::vector ranges) { auto end = std::remove_if(ranges.begin(), ranges.end(), [](const ReadRange& range) { return range.length == 0; }); ranges.resize(end - ranges.begin()); std::vector> futures; futures.reserve(ranges.size()); - for (auto& range : ranges) { - const auto it = std::lower_bound( - entries.begin(), entries.end(), range, - [](const RangeCacheEntry& entry, const ReadRange& range) { - return entry.range.offset + entry.range.length < range.offset + range.length; - }); - if (it != entries.end() && it->range.Contains(range)) { - futures.push_back(Future<>(MaybeRead(&*it))); - } else { - return Status::Invalid("Range was not requested for caching: offset=", - range.offset, " length=", range.length); + { + std::unique_lock guard(entry_mutex); + for (auto& range : ranges) { + const auto it = + std::lower_bound(entries.begin(), entries.end(), range, + [](const RangeCacheEntry& entry, const ReadRange& range) { + return entry.range.offset + entry.range.length < + range.offset + range.length; + }); + if (it != entries.end() && it->range.Contains(range)) { + futures.push_back(Future<>(MaybeRead(&*it))); + } else { + return Status::Invalid("Range was not requested for caching: offset=", + range.offset, " length=", range.length); + } } } return AllComplete(futures); } }; -// Don't read ranges when they're first added. Instead, wait until they're requested -// (either through Read or WaitFor). +// Don't read ranges when they're first added. Instead, wait until they're +// requested (either through Read or WaitFor). struct ReadRangeCache::LazyImpl : public ReadRangeCache::Impl { - // Protect against concurrent modification of entries[i]->future - std::mutex entry_mutex; - virtual ~LazyImpl() = default; Future> MaybeRead(RangeCacheEntry* entry) override { - // Called by superclass Read()/WaitFor() so we have the lock if (!entry->future.is_valid()) { entry->future = file->ReadAsync(ctx, entry->range.offset, entry->range.length, /*allow_short_read=*/false); @@ -291,26 +328,6 @@ struct ReadRangeCache::LazyImpl : public ReadRangeCache::Impl { } return new_entries; } - - Status Cache(std::vector ranges) override { - std::unique_lock guard(entry_mutex); - return ReadRangeCache::Impl::Cache(std::move(ranges)); - } - - Result> Read(ReadRange range) override { - std::unique_lock guard(entry_mutex); - return ReadRangeCache::Impl::Read(range); - } - - Future<> Wait() override { - std::unique_lock guard(entry_mutex); - return ReadRangeCache::Impl::Wait(); - } - - Future<> WaitFor(std::vector ranges) override { - std::unique_lock guard(entry_mutex); - return ReadRangeCache::Impl::WaitFor(std::move(ranges)); - } }; ReadRangeCache::ReadRangeCache(std::shared_ptr owned_file, @@ -333,12 +350,21 @@ Result> ReadRangeCache::Read(ReadRange range) { return impl_->Read(range); } +Result>> ReadRangeCache::ReadIfCached( + ReadRange range) { + return impl_->ReadIfCached(range); +} + Future<> ReadRangeCache::Wait() { return impl_->Wait(); } Future<> ReadRangeCache::WaitFor(std::vector ranges) { return impl_->WaitFor(std::move(ranges)); } +void ReadRangeCache::EvictEntriesBefore(int64_t end_offset) { + impl_->EvictEntriesBefore(end_offset); +} + } // namespace internal } // namespace io } // namespace arrow diff --git a/cpp/src/arrow/io/caching.h b/cpp/src/arrow/io/caching.h index e2b911fafdbb..0e99c8cf3151 100644 --- a/cpp/src/arrow/io/caching.h +++ b/cpp/src/arrow/io/caching.h @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -136,12 +137,23 @@ class ARROW_EXPORT ReadRangeCache { /// \brief Read a range previously given to Cache(). Result> Read(ReadRange range); + /// \brief Read a range if it is still cached. + /// + /// A cache miss returns an empty optional. I/O errors are propagated. + Result>> ReadIfCached(ReadRange range); + /// \brief Wait until all ranges added so far have been cached. Future<> Wait(); /// \brief Wait until all given ranges have been cached. Future<> WaitFor(std::vector ranges); + /// \brief Evict cache entries ending at or before `end_offset`. + /// + /// An entry straddling `end_offset` is retained. Buffers already returned by + /// Read() stay valid through shared ownership. + void EvictEntriesBefore(int64_t end_offset); + protected: struct Impl; struct LazyImpl; diff --git a/cpp/src/arrow/io/memory_test.cc b/cpp/src/arrow/io/memory_test.cc index 3e95e5257a97..4b1dc651e53f 100644 --- a/cpp/src/arrow/io/memory_test.cc +++ b/cpp/src/arrow/io/memory_test.cc @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +#include #include #include #include @@ -25,6 +26,7 @@ #include #include #include +#include #include #include @@ -744,6 +746,16 @@ class CountingBufferReader : public BufferReader { int64_t read_count_ = 0; }; +class InvalidReadBufferReader : public BufferReader { + public: + using BufferReader::BufferReader; + + Future> ReadAsync(const IOContext&, int64_t, int64_t, + bool) override { + return Future>::MakeFinished(Status::Invalid("read failed")); + } +}; + TEST(RangeReadCache, Basics) { std::string data = "abcdefghijklmnopqrstuvwxyz"; @@ -784,6 +796,8 @@ TEST(RangeReadCache, Basics) { ASSERT_RAISES(Invalid, cache.Read({19, 3})); ASSERT_RAISES(Invalid, cache.Read({0, 3})); ASSERT_RAISES(Invalid, cache.Read({25, 2})); + ASSERT_OK_AND_ASSIGN(auto missing, cache.ReadIfCached({25, 2})); + ASSERT_FALSE(missing); ASSERT_FINISHES_AND_RAISES(Invalid, cache.WaitFor({{25, 2}})); ASSERT_FINISHES_AND_RAISES(Invalid, cache.WaitFor({{1, 2}, {25, 2}})); @@ -793,6 +807,14 @@ TEST(RangeReadCache, Basics) { } } +TEST(RangeReadCache, ReadIfCachedPropagatesErrors) { + auto file = std::make_shared(Buffer::FromString("data")); + internal::ReadRangeCache cache(file, {}, CacheOptions::Defaults()); + + ASSERT_OK(cache.Cache({{0, 4}})); + ASSERT_RAISES(Invalid, cache.ReadIfCached({0, 4})); +} + TEST(RangeReadCache, Concurrency) { std::string data = "abcdefghijklmnopqrstuvwxyz"; @@ -918,6 +940,155 @@ TEST(RangeReadCache, LazyWithPrefetching) { ASSERT_RAISES(Invalid, cache.Read({25, 2})); } +TEST(RangeReadCache, EvictEntriesBefore) { + // GH-39808: entries cached by PreBuffer()-style code need to be evictable so + // memory stays bounded while iterating a large Parquet file. + std::string data = "abcdefghijklmnopqrstuvwxyz"; + + for (auto lazy : std::vector{false, true}) { + SCOPED_TRACE(lazy); + CacheOptions options = CacheOptions::Defaults(); + options.hole_size_limit = 0; // disable coalescing: one entry per range + options.range_size_limit = 10; + options.lazy = lazy; + + auto file = std::make_shared(std::make_shared(data)); + internal::ReadRangeCache cache(file, {}, options); + + // Entries: [1,3), [10,14), [20,22). + ASSERT_OK(cache.Cache({{1, 2}, {10, 4}, {20, 2}})); + ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({10, 4})); + AssertBufferEqual(*buf, "klmn"); + + // An offset that splits no entry frees nothing. + cache.EvictEntriesBefore(0); + cache.EvictEntriesBefore(2); // [1,3) extends past 2 -> retained + ASSERT_OK(cache.Read({1, 2})); + + // end_offset == an entry's end frees exactly that entry ([1,3) ends at 3). + cache.EvictEntriesBefore(3); + ASSERT_RAISES(Invalid, cache.Read({1, 2})); + ASSERT_OK_AND_ASSIGN(buf, cache.Read({10, 4})); // others intact + AssertBufferEqual(*buf, "klmn"); + + // An offset inside an entry leaves it in place. + cache.EvictEntriesBefore(12); // [10,14) straddles 12 + ASSERT_OK(cache.Read({10, 4})); + + // A wide offset frees every remaining entry in one call. + cache.EvictEntriesBefore(100); + ASSERT_RAISES(Invalid, cache.Read({10, 4})); + ASSERT_RAISES(Invalid, cache.Read({20, 2})); + + // Empty cache is a safe no-op. + cache.EvictEntriesBefore(100); + } +} + +TEST(RangeReadCache, ConcurrentReadAndEvict) { + // GH-39808: the Parquet dataset scanner calls EvictEntriesBefore from the + // thread-pool continuation that runs after a row group is decoded, while + // other threads may still be calling Read() for column chunks of other + // in-flight row groups. Exercise that pattern explicitly by slamming the + // cache with parallel Read()s interleaved with Evict()s and make sure we + // don't hit UB (iterator invalidation, torn reads, etc.). + constexpr int kNumRanges = 64; + constexpr int kRangeSize = 64; + constexpr int kIterations = 50; + std::string data(kNumRanges * kRangeSize, 'x'); + + for (auto lazy : std::vector{false, true}) { + SCOPED_TRACE(lazy); + CacheOptions options = CacheOptions::Defaults(); + // No coalescing: each range is its own entry so we can evict them + // individually without fighting the coalescing heuristic. + options.hole_size_limit = 0; + options.range_size_limit = kRangeSize; + options.lazy = lazy; + + auto file = std::make_shared(std::make_shared(data)); + internal::ReadRangeCache cache(file, {}, options); + + std::vector ranges; + ranges.reserve(kNumRanges); + for (int i = 0; i < kNumRanges; ++i) { + ranges.push_back({static_cast(i * kRangeSize), kRangeSize}); + } + ASSERT_OK(cache.Cache(ranges)); + + // Half of the threads repeatedly read the upper half of the ranges. + // The other half repeatedly evict and re-cache the lower half. Under + // the old code this would race on the shared `entries` vector. + std::atomic stop{false}; + std::atomic failures{0}; + + auto reader_fn = [&]() { + while (!stop.load(std::memory_order_relaxed)) { + for (int i = kNumRanges / 2; i < kNumRanges; ++i) { + auto result = cache.Read(ranges[i]); + if (!result.ok() || (*result)->size() != kRangeSize) { + failures.fetch_add(1, std::memory_order_relaxed); + } + } + } + }; + auto evictor_fn = [&]() { + for (int iter = 0; iter < kIterations; ++iter) { + // Evict the entire lower half in one call (entries ending before the + // midpoint offset). + const int64_t mid_offset = + ranges[kNumRanges / 2 - 1].offset + ranges[kNumRanges / 2 - 1].length; + cache.EvictEntriesBefore(mid_offset); + // Re-cache them so the next iteration has something to evict. + std::vector lower(ranges.begin(), ranges.begin() + kNumRanges / 2); + if (!cache.Cache(lower).ok()) { + failures.fetch_add(1, std::memory_order_relaxed); + } + } + stop.store(true, std::memory_order_relaxed); + }; + + std::vector threads; + for (int i = 0; i < 4; ++i) threads.emplace_back(reader_fn); + threads.emplace_back(evictor_fn); + for (auto& t : threads) t.join(); + ASSERT_EQ(0, failures.load()); + + // Every upper-half range is still readable after the torture loop. + for (int i = kNumRanges / 2; i < kNumRanges; ++i) { + ASSERT_OK_AND_ASSIGN(auto buf, cache.Read(ranges[i])); + ASSERT_EQ(kRangeSize, buf->size()); + } + } +} + +TEST(RangeReadCache, EvictEntriesBeforeSpanningEntry) { + // A coalesced entry must not be dropped until end_offset passes its end, + // otherwise we drop bytes a later consumer still needs. + std::string data(40, 'x'); + + CacheOptions options = CacheOptions::Defaults(); + options.hole_size_limit = 100; // force coalescing into one entry + options.range_size_limit = 200; + + auto file = std::make_shared(std::make_shared(data)); + internal::ReadRangeCache cache(file, {}, options); + + // {1,3} and {10,4} coalesce into a single entry [1, 14). + ASSERT_OK(cache.Cache({{1, 3}, {10, 4}})); + + // An offset inside the entry (e.g. just past the first logical range) keeps it. + cache.EvictEntriesBefore(4); + cache.EvictEntriesBefore(13); + ASSERT_OK_AND_ASSIGN(auto buf, cache.Read({10, 4})); + ASSERT_EQ(4, buf->size()); + + // end_offset at/after the entry's end (14) frees it. + cache.EvictEntriesBefore(14); + ASSERT_RAISES(Invalid, cache.Read({1, 3})); + ASSERT_RAISES(Invalid, cache.Read({10, 4})); +} + TEST(CacheOptions, Basics) { auto check = [](const CacheOptions actual, const double expected_hole_size_limit_MiB, const double expected_range_size_limit_MiB) -> void { diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 2bdbc38b3647..abe3e85068d2 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -24,8 +24,10 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" +#include #include #include +#include #include #include #include @@ -2726,6 +2728,201 @@ TEST(TestArrowReadWrite, GetRecordBatchReaderNoColumns) { ASSERT_EQ(actual_batch->num_rows(), num_rows); } +// GH-39808: bytes cached by PreBuffer() for a decoded row group must be +// releasable, else Dataset.to_batches accumulates memory over the reader's life. +TEST(TestArrowReadWrite, EvictPreBufferedDataBefore) { + ArrowReaderProperties properties = default_arrow_reader_properties(); + properties.set_pre_buffer(true); + const int num_rows = 1024; + const int row_group_size = 256; + const int num_columns = 3; + const std::vector row_groups = {0, 1, 2, 3}; + const std::vector column_indices = {0, 1, 2}; + + std::shared_ptr table; + ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); + + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, + default_arrow_writer_properties(), &buffer)); + + std::unique_ptr reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(properties)->Build(&reader)); + ASSERT_EQ(reader->num_row_groups(), static_cast(row_groups.size())); + + reader->parquet_reader()->PreBuffer(row_groups, column_indices, + ::arrow::io::IOContext(), + ::arrow::io::CacheOptions::LazyDefaults()); + ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); + + std::shared_ptr
rg_table; + ASSERT_OK(reader->ReadRowGroup(/*i=*/0, column_indices, &rg_table)); + ASSERT_EQ(rg_table->num_rows(), row_group_size); + + // The lowest byte offset row group 1 needs: evicting before it frees row + // group 0's (separately coalesced) entries while leaving row group 1 intact. + ASSERT_OK_AND_ASSIGN(auto rg1_ranges, + reader->parquet_reader()->GetReadRanges({1}, column_indices)); + int64_t rg1_min = std::numeric_limits::max(); + for (const auto& r : rg1_ranges) rg1_min = std::min(rg1_min, r.offset); + + reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min); + ASSERT_RAISES(Invalid, + reader->parquet_reader()->WhenBuffered({0}, column_indices).status()); + ASSERT_OK(reader->ReadRowGroup(/*i=*/1, column_indices, &rg_table)); + ASSERT_EQ(rg_table->num_rows(), row_group_size); + + reader->parquet_reader()->EvictPreBufferedDataBefore( + std::numeric_limits::max()); + ASSERT_RAISES(Invalid, + reader->parquet_reader()->WhenBuffered({1}, column_indices).status()); + std::shared_ptr reread_column; + ASSERT_OK(reader->RowGroup(1)->Column(0)->Read(&reread_column)); + ASSERT_EQ(reread_column->length(), row_group_size); + + // Re-evicting the same window frees nothing more. + reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min); + + // A reader that never called PreBuffer is a no-op. + std::unique_ptr no_prebuffer_reader; + FileReaderBuilder no_prebuffer_builder; + ASSERT_OK(no_prebuffer_builder.Open(std::make_shared(buffer))); + ASSERT_OK(no_prebuffer_builder.Build(&no_prebuffer_reader)); + no_prebuffer_reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min); +} + +// GH-39808: with coalescing a single cache entry can span adjacent row groups. +// Offset-based eviction frees such an entry once end_offset passes its end. +TEST(TestArrowReadWrite, EvictPreBufferedDataBeforeReleasesCrossRowGroupEntry) { + ArrowReaderProperties properties = default_arrow_reader_properties(); + properties.set_pre_buffer(true); + const int num_rows = 1024; + const int row_group_size = 256; // 4 row groups + const int num_columns = 2; + const std::vector row_groups = {0, 1, 2, 3}; + const std::vector column_indices = {0, 1}; + + std::shared_ptr
table; + ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); + + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, + default_arrow_writer_properties(), &buffer)); + + std::unique_ptr reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(properties)->Build(&reader)); + + // Huge limits coalesce every column chunk of every row group into ONE entry + // spanning all row-group boundaries. + ::arrow::io::CacheOptions options = ::arrow::io::CacheOptions::LazyDefaults(); + options.hole_size_limit = static_cast(buffer->size()); + options.range_size_limit = static_cast(buffer->size()) + 1; + reader->parquet_reader()->PreBuffer(row_groups, column_indices, + ::arrow::io::IOContext(), options); + ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); + + // Any offset short of the spanning entry's end frees nothing. + ASSERT_OK_AND_ASSIGN(auto rg3_ranges, + reader->parquet_reader()->GetReadRanges({3}, column_indices)); + int64_t rg3_min = std::numeric_limits::max(); + for (const auto& r : rg3_ranges) rg3_min = std::min(rg3_min, r.offset); + reader->parquet_reader()->EvictPreBufferedDataBefore(rg3_min); + ASSERT_OK(reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); + + // An offset past the whole file frees the spanning entry. + reader->parquet_reader()->EvictPreBufferedDataBefore( + static_cast(buffer->size())); + ASSERT_RAISES( + Invalid, + reader->parquet_reader()->WhenBuffered(row_groups, column_indices).status()); +} + +// GH-39808: when Dataset.to_batches-style iteration drives the async +// RecordBatchGenerator, each row group's pre-buffered bytes should be +// released as soon as the row group has been converted into record batches, +// so the overall memory footprint is independent of how many row groups the +// file contains. +TEST(TestArrowReadWrite, GetRecordBatchGeneratorReleasesPreBufferedRowGroups) { + ArrowReaderProperties properties = default_arrow_reader_properties(); + properties.set_pre_buffer(true); + // Read one row group at a time so the test deterministically exercises the + // per-row-group eviction path. + properties.set_batch_size(256); + + const int num_rows = 1024; + const int row_group_size = 256; + const int num_columns = 2; + + std::shared_ptr
table; + ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); + + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, row_group_size, + default_arrow_writer_properties(), &buffer)); + + // General FileReaders retain their cache because callers may create another + // generator for the same row group. + ASSERT_FALSE(properties.auto_evict_read_cache()); + std::shared_ptr retaining_reader; + { + std::unique_ptr unique_reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(properties)->Build(&unique_reader)); + retaining_reader = std::move(unique_reader); + } + ASSERT_OK_AND_ASSIGN( + auto retaining_generator, + retaining_reader->GetRecordBatchGenerator(retaining_reader, {0}, {0})); + auto retained_fut = retaining_generator(); + ASSERT_OK_AND_ASSIGN(auto retained_batch, retained_fut.result()); + ASSERT_NE(retained_batch, nullptr); + auto retained_end_fut = retaining_generator(); + ASSERT_OK_AND_ASSIGN(auto retained_end, retained_end_fut.result()); + ASSERT_EQ(retained_end, nullptr); + ASSERT_OK(retaining_reader->parquet_reader()->WhenBuffered({0}, {0}).status()); + + properties.set_auto_evict_read_cache(true); + + std::shared_ptr reader; + { + std::unique_ptr unique_reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(properties)->Build(&unique_reader)); + reader = std::move(unique_reader); + } + ASSERT_EQ(reader->num_row_groups(), num_rows / row_group_size); + + // Drive the generator exactly as ScanBatchesAsync does. + ASSERT_OK_AND_ASSIGN( + auto batch_generator, + reader->GetRecordBatchGenerator(reader, {0, 1, 2, 3}, {0, 1}, + /*cpu_executor=*/nullptr, + /*rows_to_readahead=*/2 * row_group_size)); + std::vector> batches; + for (int i = 0; i < reader->num_row_groups(); ++i) { + auto fut = batch_generator(); + ASSERT_OK_AND_ASSIGN(auto batch, fut.result()); + ASSERT_NE(batch, nullptr); + batches.push_back(std::move(batch)); + } + // Generator is drained. + auto fut_end = batch_generator(); + ASSERT_OK_AND_ASSIGN(auto end_batch, fut_end.result()); + ASSERT_EQ(end_batch, nullptr); + + ASSERT_OK_AND_ASSIGN(auto actual, + ::arrow::Table::FromRecordBatches(batches[0]->schema(), batches)); + AssertTablesEqual(*table, *actual, /*same_chunk_layout=*/false); + ASSERT_RAISES(Invalid, + reader->parquet_reader()->WhenBuffered({0, 1, 2, 3}, {0, 1}).status()); +} + TEST(TestArrowReadWrite, GetRecordBatchGenerator) { ArrowReaderProperties properties = default_arrow_reader_properties(); const int num_rows = 1024; diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index d6fe369301b7..a9f1a7de4e6e 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -1207,12 +1208,14 @@ class RowGroupGenerator { explicit RowGroupGenerator(std::shared_ptr arrow_reader, ::arrow::internal::Executor* cpu_executor, std::vector row_groups, std::vector column_indices, - int64_t min_rows_in_flight) + int64_t min_rows_in_flight, + std::vector evict_before_offsets) : arrow_reader_(std::move(arrow_reader)), cpu_executor_(cpu_executor), row_groups_(std::move(row_groups)), column_indices_(std::move(column_indices)), min_rows_in_flight_(min_rows_in_flight), + evict_before_offsets_(std::move(evict_before_offsets)), rows_in_flight_(0), index_(0), readahead_index_(0) {} @@ -1221,13 +1224,21 @@ class RowGroupGenerator { if (index_ >= row_groups_.size()) { return ::arrow::AsyncGeneratorEnd(); } - index_++; + const size_t request_index = index_++; FillReadahead(); - ReadRequest next = std::move(in_flight_reads_.front()); DCHECK(!in_flight_reads_.empty()); + ReadRequest next = std::move(in_flight_reads_.front()); in_flight_reads_.pop(); rows_in_flight_ -= next.num_rows; - return next.read; + if (evict_before_offsets_.empty()) { + return next.read; + } + auto reader = arrow_reader_; + return next.read.Then([reader, offset = evict_before_offsets_[request_index + 1]]( + RecordBatchGenerator generator) { + reader->parquet_reader()->EvictPreBufferedDataBefore(offset); + return generator; + }); } private: @@ -1244,8 +1255,7 @@ class RowGroupGenerator { } void FetchNext() { - size_t row_group_index = readahead_index_++; - int row_group = row_groups_[row_group_index]; + int row_group = row_groups_[readahead_index_++]; std::vector column_indices = column_indices_; auto reader = arrow_reader_; int64_t num_rows = @@ -1303,6 +1313,7 @@ class RowGroupGenerator { std::vector row_groups_; std::vector column_indices_; int64_t min_rows_in_flight_; + std::vector evict_before_offsets_; std::queue in_flight_reads_; int64_t rows_in_flight_; size_t index_; @@ -1325,10 +1336,27 @@ FileReaderImpl::GetRecordBatchGenerator(std::shared_ptr reader, reader_properties_.cache_options()); END_PARQUET_CATCH_EXCEPTIONS } + std::vector evict_before_offsets; + if (reader_properties_.pre_buffer() && reader_properties_.auto_evict_read_cache() && + !column_indices.empty() && !row_group_indices.empty()) { + const int64_t kNoMoreRanges = std::numeric_limits::max(); + evict_before_offsets.resize(row_group_indices.size() + 1, kNoMoreRanges); + for (int64_t i = static_cast(row_group_indices.size()) - 1; i >= 0; --i) { + ARROW_ASSIGN_OR_RAISE( + auto ranges, reader_->GetReadRanges({row_group_indices[static_cast(i)]}, + column_indices)); + int64_t rg_min = kNoMoreRanges; + for (const auto& range : ranges) { + rg_min = std::min(rg_min, range.offset); + } + evict_before_offsets[static_cast(i)] = + std::min(evict_before_offsets[static_cast(i + 1)], rg_min); + } + } ::arrow::AsyncGenerator row_group_generator = RowGroupGenerator(::arrow::internal::checked_pointer_cast(reader), cpu_executor, row_group_indices, column_indices, - rows_to_readahead); + rows_to_readahead, std::move(evict_before_offsets)); ::arrow::AsyncGenerator> concatenated = ::arrow::MakeConcatenatedGenerator(std::move(row_group_generator)); WRAP_ASYNC_GENERATOR(std::move(concatenated)); diff --git a/cpp/src/parquet/arrow/reader.h b/cpp/src/parquet/arrow/reader.h index 642546335f16..4fc97568a377 100644 --- a/cpp/src/parquet/arrow/reader.h +++ b/cpp/src/parquet/arrow/reader.h @@ -224,6 +224,10 @@ class PARQUET_EXPORT FileReader { /// The FileReader must outlive the generator, so this requires that you pass in a /// shared_ptr. /// + /// If automatic read-cache eviction is enabled, this is intended for one-pass + /// consumption. Row groups may be supplied in any order, but file order allows + /// pre-buffered memory to be released progressively. + /// /// \returns error Result if either row_group_indices or column_indices contains an /// invalid index virtual ::arrow::Result< diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index 2f46a5e296f8..955ebcc9d07d 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/cpp/src/parquet/file_reader.cc @@ -248,9 +248,12 @@ 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 { + PARQUET_ASSIGN_OR_THROW(auto buffer, cached_source_->ReadIfCached(col_range)); + if (buffer) { + stream = std::make_shared<::arrow::io::BufferReader>(*std::move(buffer)); + } + } + if (stream == nullptr) { stream = properties_.GetStream(source_, col_range.offset, col_range.length); } @@ -433,6 +436,12 @@ class SerializedFile : public ParquetFileReader::Contents { return cached_source_->WaitFor(ranges); } + void EvictPreBufferedDataBefore(int64_t end_offset) { + if (cached_source_) { + cached_source_->EvictEntriesBefore(end_offset); + } + } + // Metadata/footer parsing. Divided up to separate sync/async paths, and to use // exceptions for error handling (with the async path converting to Future/Status). @@ -617,8 +626,7 @@ class SerializedFile : public ParquetFileReader::Contents { ReaderProperties properties_; std::shared_ptr page_index_reader_; std::unique_ptr bloom_filter_reader_; - // Maps row group ordinal and prebuffer status of its column chunks in the form of a - // bitmap buffer. + std::unordered_map> prebuffered_column_chunks_; // \return The true length of the metadata in bytes @@ -909,6 +917,13 @@ void ParquetFileReader::PreBuffer(const std::vector& row_groups, file->PreBuffer(row_groups, column_indices, ctx, options); } +void ParquetFileReader::EvictPreBufferedDataBefore(int64_t end_offset) { + // Access private methods here + SerializedFile* file = + ::arrow::internal::checked_cast(contents_.get()); + file->EvictPreBufferedDataBefore(end_offset); +} + Result> ParquetFileReader::GetReadRanges( const std::vector& row_groups, const std::vector& column_indices, int64_t hole_size_limit, int64_t range_size_limit) { diff --git a/cpp/src/parquet/file_reader.h b/cpp/src/parquet/file_reader.h index c42163276cda..8a2ea80487ce 100644 --- a/cpp/src/parquet/file_reader.h +++ b/cpp/src/parquet/file_reader.h @@ -201,6 +201,11 @@ 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 if PreBuffer() was not called. + void EvictPreBufferedDataBefore(int64_t end_offset); + /// Retrieve the list of byte ranges that would need to be read to retrieve /// the data for the specified row groups and column indices. /// diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index e2244a1176e3..7086c2dcda88 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -1151,6 +1151,7 @@ class PARQUET_EXPORT ArrowReaderProperties { read_dict_indices_(), batch_size_(kArrowDefaultBatchSize), pre_buffer_(true), + auto_evict_read_cache_(false), cache_options_(::arrow::io::CacheOptions::LazyDefaults()), coerce_int96_timestamp_unit_(::arrow::TimeUnit::NANO), binary_type_(kArrowDefaultBinaryType), @@ -1235,6 +1236,14 @@ class PARQUET_EXPORT ArrowReaderProperties { /// Return whether read coalescing is enabled. bool pre_buffer() const { return pre_buffer_; } + /// Enable eviction of pre-buffered ranges as record batch generators consume them. + /// + /// This is intended for one-pass consumers. It is safe with any row group order, but + /// file order allows the cache to release memory progressively. Default is false. + void set_auto_evict_read_cache(bool auto_evict) { auto_evict_read_cache_ = auto_evict; } + /// Return whether consumed pre-buffered ranges are evicted. + bool auto_evict_read_cache() const { return auto_evict_read_cache_; } + /// Set options for read coalescing. This can be used to tune the /// implementation for characteristics of different filesystems. void set_cache_options(::arrow::io::CacheOptions options) { cache_options_ = options; } @@ -1300,6 +1309,7 @@ class PARQUET_EXPORT ArrowReaderProperties { std::unordered_set read_dict_indices_; int64_t batch_size_; bool pre_buffer_; + bool auto_evict_read_cache_; ::arrow::io::IOContext io_context_; ::arrow::io::CacheOptions cache_options_; ::arrow::TimeUnit::type coerce_int96_timestamp_unit_;