Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions include/paimon/read_context.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -116,8 +116,8 @@ class PAIMON_EXPORT ReadContext {
return specific_file_system_;
}

Comment thread
lucasfang marked this conversation as resolved.
bool EnablePrefetchCache() const {
return enable_prefetch_cache_;
PrefetchCacheMode GetPrefetchCacheMode() const {
return prefetch_cache_mode_;
}

const CacheConfig& GetCacheConfig() const {
Expand All @@ -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_;
};

Expand Down Expand Up @@ -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);
Comment thread
lucasfang marked this conversation as resolved.

/// Set the cache configuration for prefetch read operations.
///
Expand Down
18 changes: 18 additions & 0 deletions include/paimon/utils/read_ahead_cache.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment thread
lucasfang marked this conversation as resolved.
};
Comment thread
lucasfang marked this conversation as resolved.

/// Configuration parameters for the read-ahead cache behavior.
///
/// This struct controls various limits and prefetching strategies used by
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
37 changes: 28 additions & 9 deletions src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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.");
Expand All @@ -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);
}
Comment thread
lucasfang marked this conversation as resolved.
Expand Down Expand Up @@ -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.
Expand All @@ -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++) {
Expand All @@ -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();
}

Expand Down Expand Up @@ -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;
Comment thread
lucasfang marked this conversation as resolved.
}
Comment thread
lucasfang marked this conversation as resolved.
}

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;
Expand All @@ -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;
}
}

Expand Down
8 changes: 6 additions & 2 deletions src/paimon/common/reader/prefetch_file_batch_reader_impl.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand All @@ -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_;
Expand All @@ -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_;
Expand All @@ -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_;
Expand Down
Loading
Loading