polars_utils/
regex_cache.rs

1use std::cell::RefCell;
2
3use regex::bytes::{Regex as BytesRegex, RegexBuilder as BytesRegexBuilder};
4use regex::{Regex, RegexBuilder};
5
6use crate::cache::LruCache;
7
8fn get_size_limit() -> Option<usize> {
9    Some(
10        std::env::var("POLARS_REGEX_SIZE_LIMIT")
11            .ok()
12            .filter(|l| !l.is_empty())?
13            .parse()
14            .expect("invalid POLARS_REGEX_SIZE_LIMIT"),
15    )
16}
17
18// Regex compilation is really heavy, and the resulting regexes can be large as
19// well, so we should have a good caching scheme.
20//
21// TODO: add larger global cache which has time-based flush.
22
23/// A cache for compiled regular expressions.
24pub struct RegexCache {
25    cache: LruCache<String, Regex>,
26    bytes_cache: LruCache<String, BytesRegex>,
27    size_limit: Option<usize>,
28}
29
30impl RegexCache {
31    fn new() -> Self {
32        Self {
33            cache: LruCache::with_capacity(32),
34            bytes_cache: LruCache::with_capacity(32),
35            size_limit: get_size_limit(),
36        }
37    }
38
39    pub fn compile(&mut self, re: &str) -> Result<&Regex, regex::Error> {
40        let size_limit = &mut self.size_limit;
41        let r = self.cache.try_get_or_insert_with(re, |re| {
42            build_within_size_limit(size_limit, |limit| {
43                let mut builder = RegexBuilder::new(re);
44                if let Some(bytes) = limit {
45                    builder.size_limit(bytes);
46                }
47                builder.build()
48            })
49        });
50        Ok(&*r?)
51    }
52
53    pub fn compile_bytes(&mut self, re: &str) -> Result<&BytesRegex, regex::Error> {
54        let size_limit = &mut self.size_limit;
55        let r = self.bytes_cache.try_get_or_insert_with(re, |re| {
56            build_within_size_limit(size_limit, |limit| {
57                let mut builder = BytesRegexBuilder::new(re);
58                if let Some(bytes) = limit {
59                    builder.size_limit(bytes);
60                }
61                builder.build()
62            })
63        });
64        Ok(&*r?)
65    }
66}
67
68// We do this little loop to only check POLARS_REGEX_SIZE_LIMIT when a regex
69// fails to compile due to the size limit.
70fn build_within_size_limit<R>(
71    size_limit: &mut Option<usize>,
72    build: impl Fn(Option<usize>) -> Result<R, regex::Error>,
73) -> Result<R, regex::Error> {
74    loop {
75        match build(*size_limit) {
76            err @ Err(regex::Error::CompiledTooBig(_)) => {
77                let new_size_limit = get_size_limit();
78                if new_size_limit != *size_limit {
79                    *size_limit = new_size_limit;
80                    continue; // Try to compile again.
81                }
82                break err;
83            },
84            r => break r,
85        }
86    }
87}
88
89thread_local! {
90    static LOCAL_REGEX_CACHE: RefCell<RegexCache> = RefCell::new(RegexCache::new());
91}
92
93pub fn compile_regex(re: &str) -> Result<Regex, regex::Error> {
94    LOCAL_REGEX_CACHE.with_borrow_mut(|cache| cache.compile(re).cloned())
95}
96
97pub fn compile_bytes_regex(re: &str) -> Result<BytesRegex, regex::Error> {
98    LOCAL_REGEX_CACHE.with_borrow_mut(|cache| cache.compile_bytes(re).cloned())
99}
100
101pub fn with_regex_cache<R, F: FnOnce(&mut RegexCache) -> R>(f: F) -> R {
102    LOCAL_REGEX_CACHE.with_borrow_mut(f)
103}
104
105#[macro_export]
106macro_rules! cached_regex {
107    () => {};
108
109    ($vis:vis static $name:ident = $regex:expr; $($rest:tt)*) => {
110        #[allow(clippy::disallowed_methods)]
111        $vis static $name: std::sync::LazyLock<regex::Regex> = std::sync::LazyLock::new(|| regex::Regex::new($regex).unwrap());
112        $crate::regex_cache::cached_regex!($($rest)*);
113    };
114}
115pub use cached_regex;
polars_utils/regex_cache.rs

polars_utils/
regex_cache.rs
