Skip to main content

datafusion_execution/cache/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18pub mod cache_manager;
19pub mod lru_queue;
20
21pub mod default_cache;
22
23use datafusion_common::arrow::datatypes::{DataType, Schema};
24use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx};
25use datafusion_common::instant::Instant;
26use datafusion_common::{HashMap, TableReference};
27use object_store::path::Path;
28use std::collections::hash_map::DefaultHasher;
29use std::fmt::{Debug, Display, Formatter};
30use std::hash::{Hash, Hasher};
31use std::time::Duration;
32
33/// Base trait for cache implementations with common operations.
34///
35/// This trait provides the fundamental cache operations (`get`, `put`, `remove`, etc.)
36/// that all cache types share.
37///
38/// ## Thread Safety
39///
40/// Implementations must handle their own locking via internal mutability, as methods do not
41/// take mutable references and may be accessed by multiple concurrent queries.
42///
43pub trait Cache<K: CacheKey, V: CacheValue>: Send + Sync {
44    /// Get a cached entry if it exists.
45    fn get(&self, key: &K) -> Option<V>;
46
47    /// Store a value in the cache.
48    ///
49    /// Returns the previous value if one existed.
50    fn put(&self, key: &K, value: V) -> Option<V>;
51
52    /// Remove an entry from the cache, returning the value if it existed.
53    fn remove(&self, k: &K) -> Option<V>;
54
55    /// Check if the cache contains a specific key.
56    fn contains_key(&self, k: &K) -> bool;
57
58    /// Fetch the total number of cache entries.
59    fn len(&self) -> usize;
60
61    /// Check if the cache collection is empty.
62    fn is_empty(&self) -> bool {
63        self.len() == 0
64    }
65
66    /// Remove all entries from the cache.
67    fn clear(&self);
68
69    /// Return the cache name.
70    fn name(&self) -> String;
71
72    /// Current memory budget, in bytes.
73    fn cache_limit(&self) -> usize;
74
75    /// Change the memory budget in bytes.
76    fn update_cache_limit(&self, limit: usize);
77
78    /// Time-to-live applied to newly inserted entries, or `None` if entries
79    /// never expire on their own.
80    fn cache_ttl(&self) -> Option<Duration>;
81
82    /// Change the TTL applied to subsequent inserts.
83    fn update_cache_ttl(&self, _ttl: Option<Duration>);
84
85    /// Invalidate every entry associated with `table_ref`.
86    fn drop_table_entries(
87        &self,
88        table_ref: &TableReference,
89    ) -> datafusion_common::Result<()>;
90
91    /// Snapshot of all current entries with per-entry metadata (size, hits,
92    /// expiration) for diagnostics and observability.
93    fn list_entries(&self) -> HashMap<K, CacheEntryInfo<V>>;
94}
95
96/// Key type for entries stored in a [`Cache`].
97pub trait CacheKey: Clone + Eq + Hash + Send + Sync + Debug {
98    /// Size of the key in bytes, used for cache memory accounting.
99    fn size(&self) -> usize;
100
101    /// Table this key is associated with, or `None` if the key is not
102    /// table-scoped.
103    fn table_ref(&self) -> Option<&TableReference>;
104}
105
106/// Value type for entries stored in a [`Cache`].
107pub trait CacheValue: Clone + Send + Sync {
108    /// Size of the value in bytes used for cache memory accounting.
109    fn size(&self) -> usize;
110}
111
112#[derive(Clone, Debug, PartialEq, Eq)]
113pub struct CacheEntryInfo<V> {
114    pub value: V,
115    pub size_bytes: usize,
116    pub hits: usize,
117    pub expires: Option<Instant>,
118}
119
120impl<K: CacheKey, V: CacheValue> Debug for dyn Cache<K, V> {
121    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
122        write!(f, "Cache name: {} with length: {}", self.name(), self.len())
123    }
124}
125
126impl CacheKey for Path {
127    fn size(&self) -> usize {
128        self.as_ref().heap_size(&mut DFHeapSizeCtx::default())
129    }
130
131    fn table_ref(&self) -> Option<&TableReference> {
132        None
133    }
134}
135
136impl CacheKey for TableScopedPath {
137    fn size(&self) -> usize {
138        DFHeapSize::heap_size(self, &mut DFHeapSizeCtx::default())
139    }
140
141    fn table_ref(&self) -> Option<&TableReference> {
142        self.table.as_ref()
143    }
144}
145
146/// Each entry is scoped to its use within a specific table so that the cache
147/// can differentiate between identical paths in different tables, and
148/// table-level cache invalidation.
149#[derive(PartialEq, Eq, Hash, Clone, Debug)]
150pub struct TableScopedPath {
151    pub table: Option<TableReference>,
152    pub path: Path,
153}
154
155impl DFHeapSize for TableScopedPath {
156    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
157        self.path.as_ref().heap_size(ctx) + self.table.heap_size(ctx)
158    }
159}
160
161impl Display for TableScopedPath {
162    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
163        if let Some(table) = &self.table {
164            write!(f, "{}, {}", self.path, table)
165        } else {
166            write!(f, "{}", self.path)
167        }
168    }
169}
170
171/// A fingerprint of the `file_schema` used to compute a file's statistics.
172///
173/// Captures exactly the attributes that determine the layout and meaning of
174/// `Statistics::column_statistics`: each column's name, data type and
175/// nullability, in order. It deliberately excludes field/schema metadata, which
176/// cannot affect statistics — including it would needlessly fragment the cache.
177#[derive(Clone, Debug)]
178pub struct SchemaFingerprint {
179    columns: Vec<(String, DataType, bool)>,
180    /// Precomputed hash of `columns`, so hashing a key on every cache lookup is
181    /// O(1) rather than O(schema width). Computed once in `from_schema` with a
182    /// fixed-seed hasher so it is stable across keys; `PartialEq` still compares
183    /// `columns` exactly, so a hash collision can never make two different
184    /// schemas share a cache entry.
185    hash: u64,
186}
187
188impl SchemaFingerprint {
189    /// Builds a fingerprint from the `file_schema` used to compute statistics
190    /// (the schema of the columns physically read, not the full table schema —
191    /// partition columns and their statistics are handled separately).
192    pub fn from_schema(file_schema: &Schema) -> Self {
193        let columns: Vec<(String, DataType, bool)> = file_schema
194            .fields()
195            .iter()
196            .map(|f| (f.name().clone(), f.data_type().clone(), f.is_nullable()))
197            .collect();
198        let mut hasher = DefaultHasher::new();
199        columns.hash(&mut hasher);
200        Self {
201            columns,
202            hash: hasher.finish(),
203        }
204    }
205}
206
207impl PartialEq for SchemaFingerprint {
208    fn eq(&self, other: &Self) -> bool {
209        // Cheap hash gate first, then an exact comparison so collisions are safe.
210        self.hash == other.hash && self.columns == other.columns
211    }
212}
213
214impl Eq for SchemaFingerprint {}
215
216impl Hash for SchemaFingerprint {
217    fn hash<H: Hasher>(&self, state: &mut H) {
218        state.write_u64(self.hash);
219    }
220}
221
222impl DFHeapSize for SchemaFingerprint {
223    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
224        self.columns.heap_size(ctx)
225    }
226}
227
228#[cfg(test)]
229mod schema_fingerprint_tests {
230    use super::*;
231    use datafusion_common::arrow::datatypes::Field;
232
233    fn fp(fields: Vec<Field>) -> SchemaFingerprint {
234        SchemaFingerprint::from_schema(&Schema::new(fields))
235    }
236
237    /// `from_schema` must capture nullability and field order — the two
238    /// attributes most easily dropped by a wrong implementation.
239    #[test]
240    fn fingerprint_captures_nullability_and_order() {
241        assert_ne!(
242            fp(vec![Field::new("id", DataType::Int64, false)]),
243            fp(vec![Field::new("id", DataType::Int64, true)]),
244            "nullability must affect the fingerprint",
245        );
246
247        let ab = fp(vec![
248            Field::new("a", DataType::Int64, false),
249            Field::new("b", DataType::Utf8, true),
250        ]);
251        let ba = fp(vec![
252            Field::new("b", DataType::Utf8, true),
253            Field::new("a", DataType::Int64, false),
254        ]);
255        assert_ne!(ab, ba, "field order must affect the fingerprint");
256    }
257
258    /// Metadata must NOT affect the fingerprint: it cannot change column
259    /// statistics, so including it would needlessly fragment the cache.
260    #[test]
261    fn fingerprint_ignores_metadata() {
262        let plain = fp(vec![Field::new("id", DataType::Int64, false)]);
263
264        let field_md = SchemaFingerprint::from_schema(&Schema::new(vec![
265            Field::new("id", DataType::Int64, false)
266                .with_metadata([("note".to_string(), "x".to_string())].into()),
267        ]));
268        assert_eq!(plain, field_md, "field metadata must be ignored");
269
270        let schema_md = SchemaFingerprint::from_schema(
271            &Schema::new(vec![Field::new("id", DataType::Int64, false)])
272                .with_metadata([("k".to_string(), "v".to_string())].into()),
273        );
274        assert_eq!(plain, schema_md, "schema metadata must be ignored");
275    }
276}