feat: support cache mode for prefetch (#113)

Author: lucasfang

include/paimon/read_context.h

@@ -53,8 +53,8 @@ class PAIMON_EXPORT ReadContext {
                 const std::shared_ptr<Executor>& executor,
                 const std::shared_ptr<FileSystem>& specific_file_system,
                 const std::map<std::string, std::string>& fs_scheme_to_identifier_map,
-                const std::map<std::string, std::string>& options, bool enable_prefetch_cache,
-                const CacheConfig& cache_config);
+                const std::map<std::string, std::string>& options,
+                PrefetchCacheMode prefetch_cache_mode, const CacheConfig& cache_config);
     ~ReadContext();
 
     const std::string& GetPath() const {
@@ -116,8 +116,8 @@ class PAIMON_EXPORT ReadContext {
         return specific_file_system_;
     }
 
-    bool EnablePrefetchCache() const {
-        return enable_prefetch_cache_;
+    PrefetchCacheMode GetPrefetchCacheMode() const {
+        return prefetch_cache_mode_;
     }
 
     const CacheConfig& GetCacheConfig() const {
@@ -142,7 +142,7 @@ class PAIMON_EXPORT ReadContext {
     std::shared_ptr<FileSystem> specific_file_system_;
     std::map<std::string, std::string> fs_scheme_to_identifier_map_;
     std::map<std::string, std::string> options_;
-    bool enable_prefetch_cache_;
+    PrefetchCacheMode prefetch_cache_mode_;
     CacheConfig cache_config_;
 };
 
@@ -226,13 +226,13 @@ class PAIMON_EXPORT ReadContextBuilder {
     /// @return Reference to this builder for method chaining.
     ReadContextBuilder& EnablePrefetch(bool enabled);
 
-    /// Enable or disable prefetch cache for read operations.
+    /// Set prefetch cache mode for read operations.
     ///
-    /// When enabled, a prefetch cache is used to prebuffer data ranges before they are needed,
+    /// A prefetch cache is used to prebuffer data ranges before they are needed,
     /// which can improve read performance by reducing redundant I/O operations.
-    /// @param enabled Whether to enable prefetch cache (default: true)
+    /// @param mode (default: PrefetchCacheMode::ALWAYS)
     /// @return Reference to this builder for method chaining.
-    ReadContextBuilder& EnablePrefetchCache(bool enabled);
+    ReadContextBuilder& SetPrefetchCacheMode(PrefetchCacheMode mode);
 
     /// Set the cache configuration for prefetch read operations.
     ///

include/paimon/utils/read_ahead_cache.h

@@ -34,6 +34,24 @@
 
 namespace paimon {
 
+/// PrefetchCacheMode
+/// Cache prefetch switch modes.
+/// Controls whether to enable cache prefetching under different circumstances, such as queries with
+/// predicates or bitmap indexes.
+///
+/// - ALWAYS: Enable cache in all scenarios.
+/// - EXCLUDE_PREDICATE: Disable cache when query has predicates.
+/// - EXCLUDE_BITMAP: Disable cache when using bitmap index.
+/// - EXCLUDE_BITMAP_OR_PREDICATE: Disable cache if query has predicates or bitmap index.
+/// - NEVER: Always disable cache.
+enum class PAIMON_EXPORT PrefetchCacheMode {
+    ALWAYS = 1,
+    EXCLUDE_PREDICATE = 2,
+    EXCLUDE_BITMAP = 3,
+    EXCLUDE_BITMAP_OR_PREDICATE = 4,
+    NEVER = 5
+};
+
 /// Configuration parameters for the read-ahead cache behavior.
 ///
 /// This struct controls various limits and prefetching strategies used by

src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp

@@ -88,13 +88,14 @@ class ApplyBitmapIndexBatchReaderTest : public ::testing::Test,
             bool enable_prefetch = GetParam();
             if (enable_prefetch) {
                 MockFormatReaderBuilder reader_builder(data, target_type_, batch_size);
-                ASSERT_OK_AND_ASSIGN(file_batch_reader,
-                                     PrefetchFileBatchReaderImpl::Create(
-                                         /*data_file_path=*/"DUMMY", &reader_builder, fs_,
-                                         prefetch_batch_count, batch_size, prefetch_batch_count * 2,
-                                         /*enable_adaptive_prefetch_strategy=*/false, executor_,
-                                         /*initialize_read_ranges=*/true,
-                                         /*enable_prefetch_cache=*/true, CacheConfig(), pool_));
+                ASSERT_OK_AND_ASSIGN(
+                    file_batch_reader,
+                    PrefetchFileBatchReaderImpl::Create(
+                        /*data_file_path=*/"DUMMY", &reader_builder, fs_, prefetch_batch_count,
+                        batch_size, prefetch_batch_count * 2,
+                        /*enable_adaptive_prefetch_strategy=*/false, executor_,
+                        /*initialize_read_ranges=*/true,
+                        /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS, CacheConfig(), pool_));
             } else {
                 file_batch_reader =
                     std::make_unique<MockFileBatchReader>(data, target_type_, batch_size);

src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp

@@ -45,7 +45,7 @@ Result<std::unique_ptr<PrefetchFileBatchReaderImpl>> PrefetchFileBatchReaderImpl
     const std::shared_ptr<FileSystem>& fs, uint32_t prefetch_max_parallel_num, int32_t batch_size,
     uint32_t prefetch_batch_count, bool enable_adaptive_prefetch_strategy,
     const std::shared_ptr<Executor>& executor, bool initialize_read_ranges,
-    bool enable_prefetch_cache, const CacheConfig& cache_config,
+    PrefetchCacheMode prefetch_cache_mode, const CacheConfig& cache_config,
     const std::shared_ptr<MemoryPool>& pool) {
     if (prefetch_max_parallel_num == 0) {
         return Status::Invalid("prefetch max parallel num should be greater than 0.");
@@ -67,7 +67,7 @@ Result<std::unique_ptr<PrefetchFileBatchReaderImpl>> PrefetchFileBatchReaderImpl
     }
 
     std::shared_ptr<ReadAheadCache> cache;
-    if (enable_prefetch_cache) {
+    if (prefetch_cache_mode != PrefetchCacheMode::NEVER) {
         PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InputStream> input_stream, fs->Open(data_file_path));
         cache = std::make_shared<ReadAheadCache>(input_stream, cache_config, pool);
     }
@@ -102,9 +102,9 @@ Result<std::unique_ptr<PrefetchFileBatchReaderImpl>> PrefetchFileBatchReaderImpl
     }
     uint32_t prefetch_queue_capacity = prefetch_batch_count / readers.size();
 
-    auto reader = std::unique_ptr<PrefetchFileBatchReaderImpl>(
-        new PrefetchFileBatchReaderImpl(readers, batch_size, prefetch_queue_capacity,
-                                        enable_adaptive_prefetch_strategy, executor, cache));
+    auto reader = std::unique_ptr<PrefetchFileBatchReaderImpl>(new PrefetchFileBatchReaderImpl(
+        readers, batch_size, prefetch_queue_capacity, enable_adaptive_prefetch_strategy, executor,
+        cache, prefetch_cache_mode));
     if (initialize_read_ranges) {
         // normally initialize read ranges should be false, as set read schema will refresh read
         // ranges, and set read schema will always be called before read.
@@ -116,11 +116,13 @@ Result<std::unique_ptr<PrefetchFileBatchReaderImpl>> PrefetchFileBatchReaderImpl
 PrefetchFileBatchReaderImpl::PrefetchFileBatchReaderImpl(
     const std::vector<std::shared_ptr<PrefetchFileBatchReader>>& readers, int32_t batch_size,
     uint32_t prefetch_queue_capacity, bool enable_adaptive_prefetch_strategy,
-    const std::shared_ptr<Executor>& executor, const std::shared_ptr<ReadAheadCache>& cache)
+    const std::shared_ptr<Executor>& executor, const std::shared_ptr<ReadAheadCache>& cache,
+    PrefetchCacheMode cache_mode)
     : readers_(std::move(readers)),
       batch_size_(batch_size),
       executor_(executor),
       cache_(cache),
+      cache_mode_(cache_mode),
       prefetch_queue_capacity_(prefetch_queue_capacity),
       enable_adaptive_prefetch_strategy_(enable_adaptive_prefetch_strategy) {
     for (size_t i = 0; i < readers_.size(); i++) {
@@ -146,6 +148,7 @@ Status PrefetchFileBatchReaderImpl::SetReadSchema(
         PAIMON_RETURN_NOT_OK(reader->SetReadSchema(c_schema.get(), predicate, selection_bitmap));
     }
     selection_bitmap_ = selection_bitmap;
+    predicate_ = predicate;
     return RefreshReadRanges();
 }
 
@@ -275,10 +278,28 @@ Status PrefetchFileBatchReaderImpl::CleanUp() {
     return Status::OK();
 }
 
+bool PrefetchFileBatchReaderImpl::NeedInitCache() const {
+    switch (cache_mode_) {
+        case PrefetchCacheMode::NEVER:
+            return false;
+        case PrefetchCacheMode::EXCLUDE_PREDICATE:
+            return predicate_ == nullptr;
+        case PrefetchCacheMode::EXCLUDE_BITMAP:
+            return selection_bitmap_ == std::nullopt;
+        case PrefetchCacheMode::EXCLUDE_BITMAP_OR_PREDICATE:
+            return predicate_ == nullptr && selection_bitmap_ == std::nullopt;
+        case PrefetchCacheMode::ALWAYS:
+            return true;
+        default:
+            assert(false);
+            return true;
+    }
+}
+
 void PrefetchFileBatchReaderImpl::Workloop() {
     std::vector<std::future<void>> futures;
     futures.resize(readers_.size());
-    if (cache_) {
+    if (cache_ && NeedInitCache()) {
         auto read_ranges = readers_[0]->PreBufferRange();
         if (read_ranges.ok()) {
             std::vector<ByteRange> ranges;
@@ -288,11 +309,9 @@ void PrefetchFileBatchReaderImpl::Workloop() {
             auto s = cache_->Init(std::move(ranges));
             if (!s.ok()) {
                 SetReadStatus(s);
-                return;
             }
         } else {
             SetReadStatus(read_ranges.status());
-            return;
         }
     }
 

src/paimon/common/reader/prefetch_file_batch_reader_impl.h

@@ -58,7 +58,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader {
         const std::shared_ptr<FileSystem>& fs, uint32_t prefetch_max_parallel_num,
         int32_t batch_size, uint32_t prefetch_batch_count, bool enable_adaptive_prefetch_strategy,
         const std::shared_ptr<Executor>& executor, bool initialize_read_ranges,
-        bool enable_prefetch_cache, const CacheConfig& cache_config,
+        PrefetchCacheMode prefetch_cache_mode, const CacheConfig& cache_config,
         const std::shared_ptr<MemoryPool>& pool);
 
     ~PrefetchFileBatchReaderImpl() override;
@@ -111,7 +111,8 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader {
     PrefetchFileBatchReaderImpl(
         const std::vector<std::shared_ptr<PrefetchFileBatchReader>>& readers, int32_t batch_size,
         uint32_t prefetch_queue_capacity, bool enable_adaptive_prefetch_strategy,
-        const std::shared_ptr<Executor>& executor, const std::shared_ptr<ReadAheadCache>& cache);
+        const std::shared_ptr<Executor>& executor, const std::shared_ptr<ReadAheadCache>& cache,
+        PrefetchCacheMode cache_mode);
 
     Status CleanUp();
     void Workloop();
@@ -134,6 +135,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader {
                                 const std::pair<uint64_t, uint64_t>& read_range) const;
     Status HandleReadResult(size_t reader_idx, const std::pair<uint64_t, uint64_t>& read_range,
                             FileBatchReader::ReadBatchWithBitmap&& read_batch_with_bitmap);
+    bool NeedInitCache() const;
 
  private:
     std::vector<std::shared_ptr<PrefetchFileBatchReader>> readers_;
@@ -143,6 +145,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader {
     std::vector<std::unique_ptr<std::atomic<uint64_t>>> seek_cnt_;
     const int32_t batch_size_;
     std::optional<RoaringBitmap32> selection_bitmap_;
+    std::shared_ptr<Predicate> predicate_;
     std::deque<std::pair<uint64_t, uint64_t>> read_ranges_;
     std::vector<std::vector<std::pair<uint64_t, uint64_t>>> read_ranges_in_group_;
     std::vector<std::unique_ptr<ThreadsafeQueue<PrefetchBatch>>> prefetch_queues_;
@@ -151,6 +154,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader {
     std::condition_variable cv_;
     std::shared_ptr<Executor> executor_;
     std::shared_ptr<ReadAheadCache> cache_;
+    PrefetchCacheMode cache_mode_;
 
     mutable std::shared_mutex rw_mutex_;
     std::unique_ptr<std::thread> background_thread_;