diff --git a/cpp/src/arrow/io/caching.cc b/cpp/src/arrow/io/caching.cc index 74e98170ad0b..5d7df22b738b 100644 --- a/cpp/src/arrow/io/caching.cc +++ b/cpp/src/arrow/io/caching.cc @@ -151,11 +151,19 @@ struct ReadRangeCache::Impl { IOContext ctx; CacheOptions options; - // Ordered by offset (so as to find a matching region by binary search) + // 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::vector entries; + std::mutex entry_mutex; virtual ~Impl() = default; + // -- Polymorphic hooks. Always called with entry_mutex held. -- + // Get the future corresponding to a range virtual Future> MaybeRead(RangeCacheEntry* entry) { return entry->future; @@ -172,23 +180,31 @@ struct ReadRangeCache::Impl { return new_entries; } - // Add the given ranges to the cache, coalescing them where possible - virtual Status Cache(std::vector ranges) { + // -- Public entry points (acquire entry_mutex, then delegate). -- + + // 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::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); + } } - // 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(); @@ -196,22 +212,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> Read(ReadRange range) { if (range.length == 0) { static const uint8_t byte = 0; return 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 Status::Invalid("ReadRangeCache did not find matching cache entry"); + } + 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; @@ -224,53 +246,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 SliceBuffer(std::move(buf), slice_offset, range.length); } - 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); } + // `entries` is sorted by offset; an entry that extends past `end_offset` + // (coalesced with a range a later consumer still needs) is kept. + int64_t EvictEntriesBefore(int64_t end_offset) { + std::unique_lock guard(entry_mutex); + const auto first_kept = + std::remove_if(entries.begin(), entries.end(), + [end_offset](const RangeCacheEntry& entry) { + return entry.range.offset + entry.range.length <= end_offset; + }); + const auto n_evicted = static_cast(entries.end() - first_kept); + entries.erase(first_kept, entries.end()); + return n_evicted; + } + // 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). Thread safety is inherited from +// the base Impl: both MakeCacheEntries and MaybeRead are only ever invoked +// from the public entry points, and those hold the base class's entry_mutex +// before delegating. 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 + // Called by the superclass under entry_mutex, so it is safe to mutate + // entry->future here. if (!entry->future.is_valid()) { entry->future = file->ReadAsync(ctx, entry->range.offset, entry->range.length); } @@ -288,26 +335,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, @@ -336,6 +363,10 @@ Future<> ReadRangeCache::WaitFor(std::vector ranges) { return impl_->WaitFor(std::move(ranges)); } +int64_t ReadRangeCache::EvictEntriesBefore(int64_t end_offset) { + return 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..700ed919228d 100644 --- a/cpp/src/arrow/io/caching.h +++ b/cpp/src/arrow/io/caching.h @@ -142,6 +142,12 @@ class ARROW_EXPORT ReadRangeCache { /// \brief Wait until all given ranges have been cached. Future<> WaitFor(std::vector ranges); + /// \brief Evict cache entries ending at or before `end_offset`, returning the + /// count. An entry straddling `end_offset` is retained (safe when coalescing + /// merged ranges). Buffers already returned by Read() stay valid through + /// shared ownership. + int64_t 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 1b2c7bdbf393..92a36e6a480a 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 @@ -917,6 +919,153 @@ 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. + ASSERT_EQ(0, cache.EvictEntriesBefore(0)); + ASSERT_EQ(0, cache.EvictEntriesBefore(2)); // [1,3) extends past 2 -> retained + + // end_offset == an entry's end frees exactly that entry ([1,3) ends at 3). + ASSERT_EQ(1, 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. + ASSERT_EQ(0, cache.EvictEntriesBefore(12)); // [10,14) straddles 12 + + // A wide offset frees every remaining entry in one call. + ASSERT_EQ(2, cache.EvictEntriesBefore(100)); + ASSERT_RAISES(Invalid, cache.Read({10, 4})); + ASSERT_RAISES(Invalid, cache.Read({20, 2})); + + // Empty cache is a safe no-op. + ASSERT_EQ(0, 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. + ASSERT_EQ(0, cache.EvictEntriesBefore(4)); + ASSERT_EQ(0, 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. + ASSERT_EQ(1, 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 e2384972cf55..9e3e67edf38b 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -24,8 +24,11 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" +#include #include #include +#include +#include #include #include #include @@ -2670,6 +2673,183 @@ 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); + + EXPECT_GT(reader->parquet_reader()->EvictPreBufferedDataBefore(rg1_min), 0); + ASSERT_OK(reader->ReadRowGroup(/*i=*/1, column_indices, &rg_table)); + ASSERT_EQ(rg_table->num_rows(), row_group_size); + + EXPECT_GT(reader->parquet_reader()->EvictPreBufferedDataBefore( + std::numeric_limits::max()), + 0); + 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. + ASSERT_EQ(0, 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)); + ASSERT_EQ(0, + 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); + ASSERT_EQ(0, reader->parquet_reader()->EvictPreBufferedDataBefore(rg3_min)); + + // An offset past the whole file frees the spanning entry. + ASSERT_EQ(1, reader->parquet_reader()->EvictPreBufferedDataBefore( + static_cast(buffer->size()))); +} + +// 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)); + + 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})); + 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); +} + +// Readahead completes row groups out of order, so the watermark advance must +// only evict once the contiguous prefix of completed row groups grows. This +// drives that logic directly, without a reader, for determinism. +TEST(TestArrowReadWrite, ReadCacheEvictionStateOutOfOrder) { + constexpr int64_t kNoMoreRanges = std::numeric_limits::max(); + ReadCacheEvictionState state({100, 200, 300, kNoMoreRanges}); + + // Index 1 completes first: prefix stays at 0 (index 0 not done), no eviction. + EXPECT_EQ(state.MarkDecodedAndGetEvictOffset(1), std::nullopt); + // Index 0 completes: prefix jumps 0 -> 2, evict before row group 2's offset. + EXPECT_EQ(state.MarkDecodedAndGetEvictOffset(0), std::optional(300)); + // Index 2 completes: prefix jumps 2 -> 3, everything is evictable. + EXPECT_EQ(state.MarkDecodedAndGetEvictOffset(2), std::optional(kNoMoreRanges)); + // Re-decoding an already-completed index is idempotent: no eviction. + EXPECT_EQ(state.MarkDecodedAndGetEvictOffset(0), std::nullopt); +} + 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 a60af69aec9f..1092f27a06c1 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -1108,12 +1109,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::shared_ptr eviction_state) : 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), + eviction_state_(std::move(eviction_state)), rows_in_flight_(0), index_(0), readahead_index_(0) {} @@ -1159,10 +1162,19 @@ class RowGroupGenerator { auto ready = reader->parquet_reader()->WhenBuffered({row_group}, column_indices); if (cpu_executor_) ready = cpu_executor_->TransferAlways(ready); row_group_read = - ready.Then([cpu_executor = cpu_executor_, reader, row_group, + ready.Then([cpu_executor = cpu_executor_, reader, row_group, row_group_index, + eviction_state = eviction_state_, column_indices = std::move( column_indices)]() -> ::arrow::Future { - return ReadOneRowGroup(cpu_executor, reader, row_group, column_indices); + return ReadOneRowGroup(cpu_executor, reader, row_group, column_indices) + .Then([reader, row_group_index, eviction_state]( + RecordBatchGenerator generator) -> RecordBatchGenerator { + if (eviction_state) { + eviction_state->RowGroupDecoded(reader->parquet_reader(), + row_group_index); + } + return generator; + }); }); } in_flight_reads_.push({std::move(row_group_read), num_rows}); @@ -1204,6 +1216,7 @@ class RowGroupGenerator { std::vector row_groups_; std::vector column_indices_; int64_t min_rows_in_flight_; + std::shared_ptr eviction_state_; std::queue in_flight_reads_; int64_t rows_in_flight_; size_t index_; @@ -1226,10 +1239,31 @@ FileReaderImpl::GetRecordBatchGenerator(std::shared_ptr 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 eviction_state; + if (reader_properties_.pre_buffer() && !column_indices.empty() && + !row_group_indices.empty()) { + const int64_t kNoMoreRanges = std::numeric_limits::max(); + std::vector evict_before(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[static_cast(i)] = + std::min(evict_before[static_cast(i + 1)], rg_min); + } + eviction_state = std::make_shared(std::move(evict_before)); + } ::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(eviction_state)); ::arrow::AsyncGenerator> concatenated = ::arrow::MakeConcatenatedGenerator(std::move(row_group_generator)); WRAP_ASYNC_GENERATOR(std::move(concatenated)); diff --git a/cpp/src/parquet/arrow/reader_internal.h b/cpp/src/parquet/arrow/reader_internal.h index 4ee3bf98bc54..d36c376d8eaa 100644 --- a/cpp/src/parquet/arrow/reader_internal.h +++ b/cpp/src/parquet/arrow/reader_internal.h @@ -22,6 +22,8 @@ #include #include #include +#include +#include #include #include #include @@ -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 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); + } + } + + // Reader-free half of RowGroupDecoded (so the out-of-order advance is testable): + // the evict-before offset when the prefix grows, else nullopt. + std::optional MarkDecodedAndGetEvictOffset(size_t row_group_index) { + std::lock_guard lock(mutex_); + if (row_group_index >= completed_.size() || completed_[row_group_index]) { + return std::nullopt; + } + completed_[row_group_index] = true; + const size_t old_prefix = completed_prefix_; + while (completed_prefix_ < completed_.size() && completed_[completed_prefix_]) { + ++completed_prefix_; + } + if (completed_prefix_ == old_prefix) { + return std::nullopt; + } + return evict_before_offsets_[completed_prefix_]; + } + + private: + std::mutex mutex_; + std::vector evict_before_offsets_; + std::vector completed_; + size_t completed_prefix_ = 0; +}; + } // namespace arrow } // namespace parquet diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc index af7ccfd7ad7d..0711b479d794 100644 --- a/cpp/src/parquet/file_reader.cc +++ b/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); + if (!buffer.ok() && !buffer.status().IsInvalid()) { + PARQUET_THROW_NOT_OK(buffer.status()); + } + if (buffer.ok()) { + stream = std::make_shared<::arrow::io::BufferReader>(*std::move(buffer)); + } + } + if (stream == nullptr) { stream = properties_.GetStream(source_, col_range.offset, col_range.length); } @@ -432,6 +438,13 @@ class SerializedFile : public ParquetFileReader::Contents { return cached_source_->WaitFor(ranges); } + int64_t EvictPreBufferedDataBefore(int64_t end_offset) { + if (!cached_source_) { + return 0; + } + return 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). @@ -613,6 +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_; @@ -905,6 +919,13 @@ void ParquetFileReader::PreBuffer(const std::vector& row_groups, file->PreBuffer(row_groups, column_indices, ctx, options); } +int64_t ParquetFileReader::EvictPreBufferedDataBefore(int64_t end_offset) { + // Access private methods here + SerializedFile* file = + ::arrow::internal::checked_cast(contents_.get()); + return 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..8bae39e6127f 100644 --- a/cpp/src/parquet/file_reader.h +++ b/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); + /// Retrieve the list of byte ranges that would need to be read to retrieve /// the data for the specified row groups and column indices. ///