From 7b3505e3c66d954fbaadedb0428c6f1487785bae Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Tue, 18 Aug 2026 20:04:51 +0800 Subject: [PATCH] feat: add lru cache --- LICENSE | 1 + src/iceberg/delete_file_index.cc | 98 +++++------- src/iceberg/manifest/manifest_group.cc | 58 +++---- src/iceberg/test/CMakeLists.txt | 1 + src/iceberg/test/cache_test.cc | 113 ++++++++++++++ src/iceberg/test/lazy_test.cc | 8 + src/iceberg/test/meson.build | 1 + src/iceberg/util/cache_internal.h | 201 +++++++++++++++++++++++++ src/iceberg/util/lazy.h | 6 +- 9 files changed, 385 insertions(+), 102 deletions(-) create mode 100644 src/iceberg/test/cache_test.cc create mode 100644 src/iceberg/util/cache_internal.h diff --git a/LICENSE b/LICENSE index 8d8d5ff4e..d0429e24b 100644 --- a/LICENSE +++ b/LICENSE @@ -215,6 +215,7 @@ License: Public Domain This product includes code from Apache Arrow. * Core utilities: + * LRU cache and memoization utilities in src/iceberg/util/cache_internal.h. * checked_cast utility in src/iceberg/util/checked_cast.h. * FnOnce utility in src/iceberg/util/functional.h. * visit_type utility in src/iceberg/util/visit_type.h. diff --git a/src/iceberg/delete_file_index.cc b/src/iceberg/delete_file_index.cc index 634ab63bc..a8c4ef126 100644 --- a/src/iceberg/delete_file_index.cc +++ b/src/iceberg/delete_file_index.cc @@ -22,9 +22,7 @@ #include #include #include -#include #include -#include #include #include "iceberg/expression/expression.h" @@ -38,6 +36,7 @@ #include "iceberg/metrics/scan_report.h" #include "iceberg/partition_spec.h" #include "iceberg/schema.h" +#include "iceberg/util/cache_internal.h" #include "iceberg/util/checked_cast.h" #include "iceberg/util/content_file_util.h" #include "iceberg/util/executor_util_internal.h" @@ -543,12 +542,6 @@ DeleteFileIndex::Builder& DeleteFileIndex::Builder::WithScanMetrics( } Result> DeleteFileIndex::Builder::LoadDeleteFiles() { - // TODO(zehua): Replace with a thread-safe LRU cache. - std::shared_mutex projected_expr_cache_mutex; - std::unordered_map> projected_expr_cache; - std::shared_mutex eval_cache_mutex; - std::unordered_map> eval_cache; - auto data_filter = ignore_residuals_ ? True::Instance() : data_filter_; auto and_filters = @@ -560,59 +553,47 @@ Result> DeleteFileIndex::Builder::LoadDeleteFiles() { return right ? std::move(right) : std::move(left); }; - auto get_projected_expr = [&](int32_t spec_id, - const std::shared_ptr& spec) - -> Result> { - if (!data_filter_) { - return std::shared_ptr(); - } - - { - std::shared_lock lock(projected_expr_cache_mutex); - auto iter = projected_expr_cache.find(spec_id); - if (iter != projected_expr_cache.end()) { - return iter->second; - } - } + const auto cache_capacity = static_cast(specs_by_id_.size()); - std::lock_guard lock(projected_expr_cache_mutex); - auto iter = projected_expr_cache.find(spec_id); - if (iter != projected_expr_cache.end()) { - return iter->second; - } + auto get_projected_expr = internal::MemoizeLru( + [this](int32_t spec_id) -> Result> { + if (!data_filter_) { + return std::shared_ptr(); + } - auto projector = Projections::Inclusive(*spec, *schema_, case_sensitive_); - ICEBERG_ASSIGN_OR_RAISE(auto projected, projector->Project(data_filter_)); - auto [inserted_iter, _] = projected_expr_cache.emplace(spec_id, std::move(projected)); - return inserted_iter->second; - }; + auto spec_iter = specs_by_id_.find(spec_id); + ICEBERG_CHECK(spec_iter != specs_by_id_.cend(), + "Partition spec ID {} not found when projecting data filter", + spec_id); - auto get_manifest_evaluator = - [&](int32_t spec_id, const std::shared_ptr& spec, - const std::shared_ptr& filter) -> Result { - if (!filter) { - return nullptr; - } + auto projector = + Projections::Inclusive(*spec_iter->second, *schema_, case_sensitive_); + ICEBERG_ASSIGN_OR_RAISE(auto projected, projector->Project(data_filter_)); + return projected; + }, + cache_capacity); - { - std::shared_lock lock(eval_cache_mutex); - auto iter = eval_cache.find(spec_id); - if (iter != eval_cache.end()) { - return iter->second.get(); - } - } + auto get_manifest_evaluator = internal::MemoizeLru( + [this, &and_filters, &get_projected_expr]( + int32_t spec_id) -> Result> { + auto spec_iter = specs_by_id_.find(spec_id); + ICEBERG_CHECK(spec_iter != specs_by_id_.cend(), + "Partition spec ID {} not found when creating manifest evaluator", + spec_id); - std::lock_guard lock(eval_cache_mutex); - auto iter = eval_cache.find(spec_id); - if (iter != eval_cache.end()) { - return iter->second.get(); - } + ICEBERG_ASSIGN_OR_RAISE(auto projected_data_filter, get_projected_expr(spec_id)); + ICEBERG_ASSIGN_OR_RAISE(auto filter, + and_filters(partition_filter_, projected_data_filter)); + if (!filter) { + return std::shared_ptr(); + } - ICEBERG_ASSIGN_OR_RAISE(auto evaluator, ManifestEvaluator::MakePartitionFilter( - filter, spec, *schema_, case_sensitive_)); - auto [inserted_iter, _] = eval_cache.emplace(spec_id, std::move(evaluator)); - return inserted_iter->second.get(); - }; + ICEBERG_ASSIGN_OR_RAISE( + auto evaluator, ManifestEvaluator::MakePartitionFilter( + filter, spec_iter->second, *schema_, case_sensitive_)); + return std::shared_ptr(std::move(evaluator)); + }, + cache_capacity); return ParallelCollect( executor_, delete_manifests_, @@ -634,13 +615,10 @@ Result> DeleteFileIndex::Builder::LoadDeleteFiles() { const auto& spec = spec_iter->second; - ICEBERG_ASSIGN_OR_RAISE(auto projected_data_filter, - get_projected_expr(spec_id, spec)); + ICEBERG_ASSIGN_OR_RAISE(auto projected_data_filter, get_projected_expr(spec_id)); ICEBERG_ASSIGN_OR_RAISE(auto delete_partition_filter, and_filters(partition_filter_, projected_data_filter)); - ICEBERG_ASSIGN_OR_RAISE( - auto manifest_evaluator, - get_manifest_evaluator(spec_id, spec, delete_partition_filter)); + ICEBERG_ASSIGN_OR_RAISE(auto manifest_evaluator, get_manifest_evaluator(spec_id)); if (manifest_evaluator != nullptr) { ICEBERG_ASSIGN_OR_RAISE(auto should_match, manifest_evaluator->Evaluate(manifest)); diff --git a/src/iceberg/manifest/manifest_group.cc b/src/iceberg/manifest/manifest_group.cc index 3db0a1df4..02a51b113 100644 --- a/src/iceberg/manifest/manifest_group.cc +++ b/src/iceberg/manifest/manifest_group.cc @@ -21,8 +21,6 @@ #include #include -#include -#include #include #include #include @@ -42,6 +40,7 @@ #include "iceberg/schema.h" #include "iceberg/table_scan.h" #include "iceberg/type.h" +#include "iceberg/util/cache_internal.h" #include "iceberg/util/checked_cast.h" #include "iceberg/util/content_file_util.h" #include "iceberg/util/executor_util_internal.h" @@ -376,42 +375,25 @@ Result> ManifestGroup::MakeReader( Result>> ManifestGroup::ReadEntries() { - // TODO(zehua): Replace with a thread-safe LRU cache. - std::shared_mutex eval_cache_mutex; - std::unordered_map> eval_cache; - - auto get_manifest_evaluator = [&](int32_t spec_id) -> Result { - { - std::shared_lock lock(eval_cache_mutex); - auto iter = eval_cache.find(spec_id); - if (iter != eval_cache.end()) { - return iter->second.get(); - } - } - - std::lock_guard lock(eval_cache_mutex); - auto iter = eval_cache.find(spec_id); - if (iter != eval_cache.end()) { - return iter->second.get(); - } - - auto spec_iter = specs_by_id_.find(spec_id); - ICEBERG_CHECK(spec_iter != specs_by_id_.cend(), - "Cannot find partition spec for ID {}", spec_id); - - const auto& spec = spec_iter->second; - auto projector = Projections::Inclusive(*spec, *schema_, case_sensitive_); - ICEBERG_ASSIGN_OR_RAISE(auto partition_filter, projector->Project(data_filter_)); - ICEBERG_ASSIGN_OR_RAISE(partition_filter, - And::Make(partition_filter, partition_filter_)); - ICEBERG_ASSIGN_OR_RAISE( - auto manifest_evaluator, - ManifestEvaluator::MakePartitionFilter(std::move(partition_filter), spec, - *schema_, case_sensitive_)); - eval_cache[spec_id] = std::move(manifest_evaluator); - - return eval_cache[spec_id].get(); - }; + const auto cache_capacity = static_cast(specs_by_id_.size()); + auto get_manifest_evaluator = internal::MemoizeLru( + [this](int32_t spec_id) -> Result> { + auto spec_iter = specs_by_id_.find(spec_id); + ICEBERG_CHECK(spec_iter != specs_by_id_.cend(), + "Cannot find partition spec for ID {}", spec_id); + + auto projector = + Projections::Inclusive(*spec_iter->second, *schema_, case_sensitive_); + ICEBERG_ASSIGN_OR_RAISE(auto partition_filter, projector->Project(data_filter_)); + ICEBERG_ASSIGN_OR_RAISE(partition_filter, + And::Make(partition_filter, partition_filter_)); + ICEBERG_ASSIGN_OR_RAISE( + auto evaluator, ManifestEvaluator::MakePartitionFilter( + std::move(partition_filter), spec_iter->second, *schema_, + case_sensitive_)); + return std::shared_ptr(std::move(evaluator)); + }, + cache_capacity); const bool has_file_filter = file_filter_ && file_filter_->op() != Expression::Operation::kTrue; diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 1181a722e..b61f9f1e6 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -126,6 +126,7 @@ add_iceberg_test(util_test SOURCES base64_test.cc bucket_util_test.cc + cache_test.cc config_test.cc content_file_util_test.cc data_file_set_test.cc diff --git a/src/iceberg/test/cache_test.cc b/src/iceberg/test/cache_test.cc new file mode 100644 index 000000000..fbc29ab9c --- /dev/null +++ b/src/iceberg/test/cache_test.cc @@ -0,0 +1,113 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include +#include +#include +#include + +#include + +#include "iceberg/test/executor.h" +#include "iceberg/util/cache_internal.h" +#include "iceberg/util/task_group.h" + +namespace iceberg::internal { + +TEST(LruCacheTest, Eviction) { + LruCache cache(2); + + EXPECT_TRUE(cache.Replace(1, "one").first); + EXPECT_TRUE(cache.Replace(2, "two").first); + ASSERT_NE(cache.Find(1), nullptr); + + EXPECT_TRUE(cache.Replace(3, "three").first); + EXPECT_EQ(cache.Find(2), nullptr); + ASSERT_NE(cache.Find(1), nullptr); + EXPECT_EQ(*cache.Find(1), "one"); + ASSERT_NE(cache.Find(3), nullptr); + EXPECT_EQ(*cache.Find(3), "three"); +} + +TEST(MemoizeLruTest, Caches) { + std::atomic calls{0}; + auto memoized = MemoizeLru( + [&calls](int32_t key) { + ++calls; + return std::to_string(key); + }, + 2); + + EXPECT_EQ(memoized(1), "1"); + EXPECT_EQ(memoized(2), "2"); + EXPECT_EQ(memoized(1), "1"); + EXPECT_EQ(calls, 2); + + EXPECT_EQ(memoized(3), "3"); + EXPECT_EQ(memoized(2), "2"); + EXPECT_EQ(calls, 4); +} + +TEST(MemoizeLruTest, MovesKey) { + auto memoized = MemoizeLru([](const std::unique_ptr& key) { return *key; }, 2); + EXPECT_EQ(memoized(std::make_unique(1)), 1); + + auto memoized_thread_unsafe = + MemoizeLruThreadUnsafe([](const std::unique_ptr& key) { return *key; }, 2); + EXPECT_EQ(memoized_thread_unsafe(std::make_unique(2)), 2); +} + +TEST(MemoizeLruTest, ThreadUnsafe) { + int32_t calls = 0; + auto memoized = MemoizeLruThreadUnsafe( + [&calls](int32_t key) { + ++calls; + return std::to_string(key); + }, + 2); + + EXPECT_EQ(memoized(1), "1"); + EXPECT_EQ(memoized(1), "1"); + EXPECT_EQ(calls, 1); +} + +TEST(MemoizeLruTest, Threads) { + auto memoized = MemoizeLru([](int32_t key) { return key * 2; }, 4); + std::atomic mismatch{false}; + test::ThreadExecutor executor; + TaskGroup group; + group.SetExecutor(std::ref(executor)); + + for (int32_t thread_id = 0; thread_id < 4; ++thread_id) { + group.Submit([&] -> Status { + for (int32_t i = 0; i < 100; ++i) { + const int32_t key = i % 8; + if (memoized(key) != key * 2) { + mismatch.store(true, std::memory_order_relaxed); + } + } + return {}; + }); + } + + ASSERT_TRUE(std::move(group).Run().has_value()); + EXPECT_FALSE(mismatch.load(std::memory_order_relaxed)); +} + +} // namespace iceberg::internal diff --git a/src/iceberg/test/lazy_test.cc b/src/iceberg/test/lazy_test.cc index 04eac6807..59f711eef 100644 --- a/src/iceberg/test/lazy_test.cc +++ b/src/iceberg/test/lazy_test.cc @@ -66,3 +66,11 @@ TEST(LazyTest, ReusesInitializationError) { EXPECT_THAT(second, iceberg::IsError(iceberg::ErrorKind::kInvalid)); EXPECT_THAT(second, iceberg::HasErrorMessage("init failed")); } + +TEST(LazyTest, SupportsLambda) { + const iceberg::Lazy<[](int value) -> iceberg::Result { return value; }> lazy; + + auto result = lazy.Get(42); + ASSERT_THAT(result, iceberg::IsOk()); + EXPECT_EQ(result->get(), 42); +} diff --git a/src/iceberg/test/meson.build b/src/iceberg/test/meson.build index 6dde45ee9..87b317904 100644 --- a/src/iceberg/test/meson.build +++ b/src/iceberg/test/meson.build @@ -100,6 +100,7 @@ iceberg_tests = { 'sources': files( 'base64_test.cc', 'bucket_util_test.cc', + 'cache_test.cc', 'config_test.cc', 'content_file_util_test.cc', 'data_file_set_test.cc', diff --git a/src/iceberg/util/cache_internal.h b/src/iceberg/util/cache_internal.h new file mode 100644 index 000000000..d82a93b80 --- /dev/null +++ b/src/iceberg/util/cache_internal.h @@ -0,0 +1,201 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Borrowed the file from Apache Arrow: +// https://github.com/apache/arrow/blob/main/cpp/src/arrow/util/cache_internal.h + +#pragma once + +/// \file iceberg/util/cache_internal.h +/// \brief Provide internal LRU cache and memoization helpers. + +#include +#include +#include +#include +#include +#include +#include + +#include "iceberg/util/macros.h" + +namespace iceberg::internal { + +// A LRU (Least recently used) replacement cache +template +class LruCache { + public: + explicit LruCache(int32_t capacity) : capacity_(capacity) { + // The map size can temporarily exceed the cache capacity, see Replace() + map_.reserve(capacity_ + 1); + } + + LruCache(const LruCache&) = delete; + LruCache& operator=(const LruCache&) = delete; + LruCache(LruCache&&) noexcept = default; + LruCache& operator=(LruCache&&) noexcept = default; + + void Clear() { + items_.clear(); + map_.clear(); + // The C++ spec doesn't tell whether map_.clear() will shrink the map capacity + map_.reserve(capacity_ + 1); + } + + int32_t size() const { + ICEBERG_DCHECK(items_.size() == map_.size(), "Cache item and map sizes differ"); + return static_cast(items_.size()); + } + + template + Value* Find(K&& key) { + const auto it = map_.find(key); + if (it == map_.end()) { + return nullptr; + } + + // Found => move item at front of the list + auto list_it = it->second; + items_.splice(items_.begin(), items_, list_it); + return &list_it->value; + } + + template + std::pair Replace(K&& key, V&& value) { + // Try to insert temporary iterator + auto [it, inserted] = map_.emplace(std::forward(key), ListIt{}); + if (inserted) { + // Inserted => push item at front of the list, and update iterator + items_.emplace_front(&it->first, std::forward(value)); + it->second = items_.begin(); + // Did we exceed the cache capacity? If so, remove least recently used item + if (static_cast(items_.size()) > capacity_) { + const bool erased = map_.erase(*items_.back().key); + ICEBERG_DCHECK(erased, "Failed to erase least recently used cache item"); + static_cast(erased); + items_.pop_back(); + } + return {true, &it->second->value}; + } + + // Already exists => move item at front of the list, and update value + auto list_it = it->second; + items_.splice(items_.begin(), items_, list_it); + list_it->value = std::forward(value); + return {false, &list_it->value}; + } + + private: + struct Item { + // Pointer to the key inside the unordered_map + const Key* key; + Value value; + }; + using List = std::list; + using ListIt = typename List::iterator; + + const int32_t capacity_; + // In most to least recently used order + std::list items_; + std::unordered_map map_; +}; + +template +struct ThreadSafeMemoizer { + using RetType = Value; + + template + ThreadSafeMemoizer(F&& func, int32_t cache_capacity) + : func_(std::forward(func)), cache_(cache_capacity) {} + + // The memoizer can't return a pointer to the cached value, because + // the cache entry may be evicted by another thread. + template + RetType operator()(K&& key) { + std::unique_lock lock(mutex_); + if (const Value* value_ptr = cache_.Find(key); value_ptr != nullptr) { + return *value_ptr; + } + lock.unlock(); + Value value = func_(key); + lock.lock(); + return *cache_.Replace(std::forward(key), std::move(value)).second; + } + + private: + std::mutex mutex_; + Func func_; + Cache cache_; +}; + +template +struct ThreadUnsafeMemoizer { + using RetType = const Value&; + + template + ThreadUnsafeMemoizer(F&& func, int32_t cache_capacity) + : func_(std::forward(func)), cache_(cache_capacity) {} + + template + RetType operator()(K&& key) { + if (const Value* value_ptr = cache_.Find(key); value_ptr != nullptr) { + return *value_ptr; + } + Value value = func_(key); + return *cache_.Replace(std::forward(key), std::move(value)).second; + } + + private: + Func func_; + Cache cache_; +}; + +template +struct unary_traits; + +template +struct unary_traits> { + using arg = Arg; + using return_type = R; +}; + +template