Expose cache stats (#79)

Author: dentiny

src/include/query_condition_cache_functions.hpp

@@ -2,6 +2,7 @@
 
 #include "query_condition_cache_state.hpp"
 
+#include "duckdb/function/scalar_function.hpp"
 #include "duckdb/function/table_function.hpp"
 
 namespace duckdb {
@@ -18,5 +19,7 @@ shared_ptr<ConditionCacheEntry> BuildCacheEntry(ClientContext &context, DuckTabl
 
 TableFunction ConditionCacheBuildFunction();
 TableFunction ConditionCacheInfoFunction();
+TableFunction ConditionCacheStatsFunction();
+ScalarFunction ConditionCacheResetStatsFunction();
 
 } // namespace duckdb

src/include/query_condition_cache_state.hpp

@@ -5,6 +5,7 @@
 #include "concurrency/thread_annotation.hpp"
 
 #include "duckdb/common/array.hpp"
+#include "duckdb/common/atomic.hpp"
 #include "duckdb/common/types/hash.hpp"
 #include "duckdb/common/unordered_map.hpp"
 #include "duckdb/common/unordered_set.hpp"
@@ -131,6 +132,14 @@ struct TableFilterKeyIndex : public ObjectCacheEntry {
 	bool IsEmpty();
 	// Transfer ownership of all filter keys out. Clears the internal set.
 	unordered_set<string> Take();
+	// Return a copy of all filter keys without clearing the set.
+	unordered_set<string> Snapshot();
+};
+
+struct CacheStoreStats {
+	idx_t total_memory_bytes;
+	idx_t hit_count;
+	idx_t access_count;
 };
 
 // Stored in DuckDB's per-database ObjectCache
@@ -169,11 +178,26 @@ class ConditionCacheStore : public ObjectCacheEntry {
 	// Get or create the store from a client context
 	static shared_ptr<ConditionCacheStore> GetOrCreate(ClientContext &context);
 
+	// Record an optimizer lookup attempt.
+	void RecordAccess(bool hit);
+
+	// Reset cache access stats.
+	void ResetStats();
+
+	// Compute the sum of estimated memory used by all live cache entries.
+	idx_t ComputeTotalMemoryBytes(ClientContext &context) const;
+
+	// Return a snapshot of current stats.
+	CacheStoreStats GetStats(ClientContext &context) const;
+
 private:
-	concurrency::mutex lock;
+	mutable concurrency::mutex lock;
 	// Tracks all table OIDs that have been cached, for ClearAll
 	unordered_set<idx_t> cached_table_oids DUCKDB_GUARDED_BY(lock);
 
+	atomic<idx_t> total_accesses {0};
+	atomic<idx_t> total_hits {0};
+
 	static string MakeCacheKeyString(const CacheKey &key);
 	static string MakeFilterKeyIndexKey(idx_t table_oid);
 };

src/query_condition_cache_extension.cpp

@@ -28,6 +28,8 @@ void OnQueryConditionCacheSettingChange(ClientContext &context, SetScope scope,
 void LoadInternal(ExtensionLoader &loader) {
 	loader.RegisterFunction(ConditionCacheBuildFunction());
 	loader.RegisterFunction(ConditionCacheInfoFunction());
+	loader.RegisterFunction(ConditionCacheStatsFunction());
+	loader.RegisterFunction(ConditionCacheResetStatsFunction());
 
 	// Register the internal filter function so it survives plan serialization/verification
 	loader.RegisterFunction(ConditionCacheFilterFunction());

src/query_condition_cache_functions.cpp

@@ -11,6 +11,7 @@
 #include "duckdb/common/unordered_set.hpp"
 #include "duckdb/common/vector.hpp"
 #include "duckdb/execution/expression_executor.hpp"
+#include "duckdb/function/scalar_function.hpp"
 #include "duckdb/function/table_function.hpp"
 #include "duckdb/parallel/task_executor.hpp"
 #include "duckdb/parallel/task_scheduler.hpp"
@@ -380,4 +381,75 @@ TableFunction ConditionCacheInfoFunction() {
 	return func;
 }
 
+// ------- condition_cache_stats() -------
+// Returns global cache statistics: total memory, hit count, access count.
+
+namespace {
+
+struct ConditionCacheStatsState : public GlobalTableFunctionState {
+	bool done = false;
+};
+
+unique_ptr<FunctionData> ConditionCacheStatsBind(ClientContext &context, TableFunctionBindInput &input,
+                                                 vector<LogicalType> &return_types, vector<string> &names) {
+	names.emplace_back("total_memory_bytes");
+	return_types.emplace_back(LogicalType {LogicalTypeId::UBIGINT});
+	names.emplace_back("hit_count");
+	return_types.emplace_back(LogicalType {LogicalTypeId::UBIGINT});
+	names.emplace_back("access_count");
+	return_types.emplace_back(LogicalType {LogicalTypeId::UBIGINT});
+	return nullptr;
+}
+
+unique_ptr<GlobalTableFunctionState> ConditionCacheStatsInit(ClientContext &context, TableFunctionInitInput &input) {
+	return make_uniq<ConditionCacheStatsState>();
+}
+
+void ConditionCacheStatsExecute(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) {
+	auto &gstate = data_p.global_state->Cast<ConditionCacheStatsState>();
+	if (gstate.done) {
+		return;
+	}
+	gstate.done = true;
+
+	auto store = ConditionCacheStore::GetOrCreate(context);
+	auto stats = store->GetStats(context);
+
+	output.SetCardinality(1);
+	output.data[0].SetValue(0, Value::UBIGINT(stats.total_memory_bytes));
+	output.data[1].SetValue(0, Value::UBIGINT(stats.hit_count));
+	output.data[2].SetValue(0, Value::UBIGINT(stats.access_count));
+}
+
+} // namespace
+
+TableFunction ConditionCacheStatsFunction() {
+	TableFunction func("condition_cache_stats", {}, ConditionCacheStatsExecute, ConditionCacheStatsBind,
+	                   ConditionCacheStatsInit);
+	return func;
+}
+
+// ------- condition_cache_reset_stats() -------
+// Scalar function: resets hit_count and access_count, returns true.
+
+namespace {
+
+unique_ptr<FunctionData> ConditionCacheResetStatsBind(ClientContext &context, ScalarFunction &bound_function,
+                                                      vector<unique_ptr<Expression>> &arguments) {
+	auto store = ConditionCacheStore::GetOrCreate(context);
+	store->ResetStats();
+	return nullptr;
+}
+
+void ConditionCacheResetStatsFn(DataChunk &args, ExpressionState &state, Vector &result) {
+	result.Reference(Value::BOOLEAN(true));
+}
+
+} // namespace
+
+ScalarFunction ConditionCacheResetStatsFunction() {
+	return ScalarFunction("condition_cache_reset_stats", {}, LogicalType {LogicalTypeId::BOOLEAN},
+	                      ConditionCacheResetStatsFn, ConditionCacheResetStatsBind);
+}
+
 } // namespace duckdb

src/query_condition_cache_optimizer.cpp

@@ -101,6 +101,7 @@ void QueryConditionCacheOptimizer::PreOptimizeWalk(ClientContext &context, uniqu
 
 	auto store = ConditionCacheStore::GetOrCreate(context);
 	auto entry = store->Lookup(context, key);
+	store->RecordAccess(entry != nullptr);
 
 	if (!entry) {
 		// TODO: Consider building cache in the background and syncing later

src/query_condition_cache_state.cpp

@@ -187,6 +187,11 @@ unordered_set<string> TableFilterKeyIndex::Take() {
 	return std::move(filter_keys);
 }
 
+unordered_set<string> TableFilterKeyIndex::Snapshot() {
+	concurrency::lock_guard<concurrency::mutex> guard(lock);
+	return filter_keys;
+}
+
 // ------- CONDITION_CACHE_STORE -------
 
 string ConditionCacheStore::MakeCacheKeyString(const CacheKey &key) {
@@ -275,11 +280,59 @@ void ConditionCacheStore::ClearAll(ClientContext &context) {
 		cache.Delete(MakeFilterKeyIndexKey(table_oid));
 	}
 	cached_table_oids.clear();
+	ResetStats();
 }
 
 shared_ptr<ConditionCacheStore> ConditionCacheStore::GetOrCreate(ClientContext &context) {
 	auto &cache = ObjectCache::GetObjectCache(context);
 	return cache.GetOrCreate<ConditionCacheStore>(CACHE_KEY);
 }
 
+void ConditionCacheStore::RecordAccess(bool hit) {
+	total_accesses.fetch_add(1, std::memory_order_relaxed);
+	if (hit) {
+		total_hits.fetch_add(1, std::memory_order_relaxed);
+	}
+}
+
+void ConditionCacheStore::ResetStats() {
+	total_accesses.store(0, std::memory_order_relaxed);
+	total_hits.store(0, std::memory_order_relaxed);
+}
+
+idx_t ConditionCacheStore::ComputeTotalMemoryBytes(ClientContext &context) const {
+	auto &cache = ObjectCache::GetObjectCache(context);
+
+	unordered_set<idx_t> oid_snapshot;
+	{
+		concurrency::lock_guard<concurrency::mutex> guard(lock);
+		oid_snapshot = cached_table_oids;
+	}
+
+	idx_t total = 0;
+	for (auto table_oid : oid_snapshot) {
+		auto index = cache.Get<TableFilterKeyIndex>(MakeFilterKeyIndexKey(table_oid));
+		if (!index) {
+			continue;
+		}
+		for (const auto &filter_key : index->Snapshot()) {
+			auto entry = cache.Get<ConditionCacheEntry>(MakeCacheKeyString(CacheKey {table_oid, filter_key}));
+			if (entry) {
+				auto mem = entry->GetEstimatedCacheMemory();
+				ALWAYS_ASSERT(mem.IsValid());
+				total += mem.GetIndex();
+			}
+		}
+	}
+	return total;
+}
+
+CacheStoreStats ConditionCacheStore::GetStats(ClientContext &context) const {
+	return CacheStoreStats {
+	    .total_memory_bytes = ComputeTotalMemoryBytes(context),
+	    .hit_count = total_hits.load(std::memory_order_relaxed),
+	    .access_count = total_accesses.load(std::memory_order_relaxed),
+	};
+}
+
 } // namespace duckdb

test/sql/condition_cache_stats.test

@@ -0,0 +1,95 @@
+# name: test/sql/condition_cache_stats.test
+# description: Test predicate cache stats
+# group: [sql]
+
+require query_condition_cache
+
+# Before any cache activity: all stats are zero
+query III
+SELECT total_memory_bytes, hit_count, access_count FROM condition_cache_stats();
+----
+0	0	0
+
+statement ok
+CREATE TABLE t AS SELECT i AS id, i % 100 AS val FROM range(500000) t(i);
+
+statement ok
+SELECT * FROM condition_cache_build('t', 'val = 42');
+
+query I
+SELECT total_memory_bytes > 0 FROM condition_cache_stats();
+----
+true
+
+query II
+SELECT hit_count, access_count FROM condition_cache_stats();
+----
+0	0
+
+# First auto query: optimizer finds the pre-built entry (hit)
+statement ok
+SELECT count(*) FROM t WHERE val = 42;
+
+query II
+SELECT hit_count, access_count FROM condition_cache_stats();
+----
+1	1
+
+# Second auto query: another hit
+statement ok
+SELECT count(*) FROM t WHERE val = 42;
+
+query II
+SELECT hit_count, access_count FROM condition_cache_stats();
+----
+2	2
+
+# Capture memory before reset to verify it is unchanged afterwards
+statement ok
+CREATE TABLE mem_before AS SELECT total_memory_bytes FROM condition_cache_stats();
+
+# Reset stats: clears counters but does NOT clear cache entries
+query I
+SELECT condition_cache_reset_stats();
+----
+true
+
+query II
+SELECT hit_count, access_count FROM condition_cache_stats();
+----
+0	0
+
+# Memory is unchanged after reset: must equal the value captured before reset
+query I
+SELECT (SELECT total_memory_bytes FROM condition_cache_stats()) =
+       (SELECT total_memory_bytes FROM mem_before);
+----
+true
+
+# First query on a new predicate (not pre-built): optimizer misses, builds entry
+statement ok
+SELECT count(*) FROM t WHERE val = 99;
+
+# access_count = 1, hit_count = 0 (first lookup was a miss)
+query II
+SELECT hit_count, access_count FROM condition_cache_stats();
+----
+0	1
+
+# Second query on the same new predicate: optimizer finds the entry
+statement ok
+SELECT count(*) FROM t WHERE val = 99;
+
+query II
+SELECT hit_count, access_count FROM condition_cache_stats();
+----
+1	2
+
+# Disabling the setting calls ClearAll, which drops all cache entries and resets all stats
+statement ok
+SET use_query_condition_cache = false;
+
+query III
+SELECT total_memory_bytes, hit_count, access_count FROM condition_cache_stats();
+----
+0	0	0