polars_utils/
regex_cache.rs

1use std::cell::RefCell;
2use std::sync::Arc;
3
4use regex::bytes::Regex as BytesRegex;
5use regex::{Regex, RegexBuilder};
6
7use crate::cache::LruCache;
8
9fn get_size_limit() -> Option<usize> {
10    Some(
11        std::env::var("POLARS_REGEX_SIZE_LIMIT")
12            .ok()
13            .filter(|l| !l.is_empty())?
14            .parse()
15            .expect("invalid POLARS_REGEX_SIZE_LIMIT"),
16    )
17}
18
19// Regex compilation is really heavy, and the resulting regexes can be large as
20// well, so we should have a good caching scheme.
21//
22// TODO: add larger global cache which has time-based flush.
23
24/// A cache for compiled regular expressions.
25pub struct RegexCache {
26    cache: LruCache<String, Regex>,
27    bytes_cache: LruCache<usize, (Arc<BytesRegex>, BytesRegex)>,
28    size_limit: Option<usize>,
29}
30
31impl RegexCache {
32    fn new() -> Self {
33        Self {
34            cache: LruCache::with_capacity(32),
35            bytes_cache: LruCache::with_capacity(32),
36            size_limit: get_size_limit(),
37        }
38    }
39
40    pub fn compile(&mut self, re: &str) -> Result<&Regex, regex::Error> {
41        let size_limit = &mut self.size_limit;
42        let r = self.cache.try_get_or_insert_with(re, |re| {
43            build_within_size_limit(size_limit, |limit| {
44                let mut builder = RegexBuilder::new(re);
45                if let Some(bytes) = limit {
46                    builder.size_limit(bytes);
47                }
48                builder.build()
49            })
50        });
51        Ok(&*r?)
52    }
53
54    /// Borrows a cached clone of the supplied regex, keyed by its Arc pointer.
55    pub fn get_or_insert_bytes(&mut self, re: &Arc<BytesRegex>) -> &BytesRegex {
56        // Retain the original Arc so its address cannot be reused while cached.
57        &self
58            .bytes_cache
59            .get_or_insert_with(&(Arc::as_ptr(re) as usize), |_| {
60                (re.clone(), re.as_ref().clone())
61            })
62            .1
63    }
64}
65
66// We do this little loop to only check POLARS_REGEX_SIZE_LIMIT when a regex
67// fails to compile due to the size limit.
68fn build_within_size_limit<R>(
69    size_limit: &mut Option<usize>,
70    build: impl Fn(Option<usize>) -> Result<R, regex::Error>,
71) -> Result<R, regex::Error> {
72    loop {
73        match build(*size_limit) {
74            err @ Err(regex::Error::CompiledTooBig(_)) => {
75                let new_size_limit = get_size_limit();
76                if new_size_limit != *size_limit {
77                    *size_limit = new_size_limit;
78                    continue; // Try to compile again.
79                }
80                break err;
81            },
82            r => break r,
83        }
84    }
85}
86
87thread_local! {
88    static LOCAL_REGEX_CACHE: RefCell<RegexCache> = RefCell::new(RegexCache::new());
89}
90
91pub fn compile_regex(re: &str) -> Result<Regex, regex::Error> {
92    LOCAL_REGEX_CACHE.with_borrow_mut(|cache| cache.compile(re).cloned())
93}
94
95pub fn with_regex_cache<R, F: FnOnce(&mut RegexCache) -> R>(f: F) -> R {
96    LOCAL_REGEX_CACHE.with_borrow_mut(f)
97}
98
99#[macro_export]
100macro_rules! cached_regex {
101    () => {};
102
103    ($vis:vis static $name:ident = $regex:expr; $($rest:tt)*) => {
104        #[allow(clippy::disallowed_methods)]
105        $vis static $name: std::sync::LazyLock<regex::Regex> = std::sync::LazyLock::new(|| regex::Regex::new($regex).unwrap());
106        $crate::regex_cache::cached_regex!($($rest)*);
107    };
108}
109pub use cached_regex;
110
111#[cfg(test)]
112mod tests {
113    use regex::bytes::RegexBuilder as BytesRegexBuilder;
114
115    use super::*;
116
117    #[test]
118    fn bytes_cache_preserves_builder_options() {
119        let mut cache = RegexCache::new();
120        let sensitive = Arc::new(BytesRegexBuilder::new("abc").build().unwrap());
121        let insensitive = Arc::new(
122            BytesRegexBuilder::new("abc")
123                .case_insensitive(true)
124                .build()
125                .unwrap(),
126        );
127
128        for _ in 0..2 {
129            assert!(!cache.get_or_insert_bytes(&sensitive).is_match(b"ABC"));
130            assert!(cache.get_or_insert_bytes(&insensitive).is_match(b"ABC"));
131        }
132    }
133
134    #[test]
135    fn bytes_cache_reuses_cloned_arc() {
136        let mut cache = RegexCache::new();
137        let re = Arc::new(BytesRegexBuilder::new("abc").build().unwrap());
138        let cloned = re.clone();
139        let cached = std::ptr::from_ref(cache.get_or_insert_bytes(&re));
140
141        assert_eq!(
142            cached,
143            std::ptr::from_ref(cache.get_or_insert_bytes(&cloned))
144        );
145        assert_ne!(cached, Arc::as_ptr(&re));
146    }
147
148    #[test]
149    fn bytes_cache_retains_arc_until_eviction() {
150        let mut cache = RegexCache::new();
151        let regex = BytesRegexBuilder::new("abc").build().unwrap();
152        let re = Arc::new(regex.clone());
153        let weak = Arc::downgrade(&re);
154        cache.get_or_insert_bytes(&re);
155        drop(re);
156        assert!(weak.upgrade().is_some());
157
158        for _ in 0..32 {
159            let re = Arc::new(regex.clone());
160            cache.get_or_insert_bytes(&re);
161        }
162        assert!(weak.upgrade().is_none());
163    }
164}
polars_utils/regex_cache.rs

polars_utils/
regex_cache.rs
