ursa_engine/
options.rs

1//! Ursa's session settings, set with `SET ursa.<name> = …` and shown with
2//! `SHOW ursa.<name>`: a DataFusion config extension with the prefix `ursa`.
3//!
4//! Settings last for the session, with one exception: in a writable `.ursa`
5//! database, `SET ursa.index_storage_limit` (and `RESET`) is stored in the
6//! catalog file and applied again when the database is opened, because it
7//! governs what the database keeps on disk (#327). The database's stored
8//! indexes obey the stored limit (or the default), never a session's own: a
9//! read-only or in-memory session's `SET` changes only what `SHOW` says.
10
11use std::fmt;
12
13use datafusion::common::extensions_options;
14use datafusion::config::{ConfigExtension, ConfigField, Visit};
15use datafusion::error::{DataFusionError, Result};
16use datafusion::prelude::SessionConfig;
17
18/// The default `ursa.index_storage_limit`: 4 GiB.
19pub const DEFAULT_INDEX_STORAGE_LIMIT: u64 = 4 << 30;
20
21/// A number of bytes, written as plain bytes (`1048576`) or with a binary
22/// unit (`512KiB`, `64MiB`, `8GiB`, `1TiB`; `B` for bytes).
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub struct ByteSize(pub u64);
25
26const UNITS: [(&str, u64); 5] = [
27    ("TiB", 1 << 40),
28    ("GiB", 1 << 30),
29    ("MiB", 1 << 20),
30    ("KiB", 1 << 10),
31    ("B", 1),
32];
33
34impl ByteSize {
35    /// Parse `512KiB`, `8 GiB`, `1.5GiB` or `1048576`. Units are binary and
36    /// matched without regard to case.
37    pub fn parse(text: &str) -> Result<Self> {
38        let invalid = || {
39            DataFusionError::Configuration(format!(
40                "'{text}' is not a size; write bytes (1048576) or a number with KiB, MiB, GiB \
41                 or TiB (8GiB)"
42            ))
43        };
44        let text = text.trim();
45        let split = text
46            .find(|c: char| !(c.is_ascii_digit() || c == '.'))
47            .unwrap_or(text.len());
48        let (number, unit) = (text[..split].trim(), text[split..].trim());
49        let scale = match unit {
50            "" => 1,
51            unit => {
52                UNITS
53                    .iter()
54                    .find(|(name, _)| name.eq_ignore_ascii_case(unit))
55                    .ok_or_else(invalid)?
56                    .1
57            }
58        };
59        if let Ok(n) = number.parse::<u64>() {
60            return n.checked_mul(scale).map(ByteSize).ok_or_else(invalid);
61        }
62        let n: f64 = number.parse().map_err(|_| invalid())?;
63        let bytes = n * scale as f64;
64        if !bytes.is_finite() || bytes < 0.0 || bytes >= u64::MAX as f64 {
65            return Err(invalid());
66        }
67        Ok(ByteSize(bytes as u64))
68    }
69}
70
71impl fmt::Display for ByteSize {
72    /// In the largest binary unit that divides it, as `SHOW` prints it.
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74        let (name, scale) = UNITS
75            .iter()
76            .find(|(_, scale)| self.0 != 0 && self.0.is_multiple_of(*scale))
77            .copied()
78            .unwrap_or(("B", 1));
79        if name == "B" {
80            write!(f, "{}", self.0)
81        } else {
82            write!(f, "{}{name}", self.0 / scale)
83        }
84    }
85}
86
87impl ConfigField for ByteSize {
88    fn visit<V: Visit>(&self, v: &mut V, key: &str, description: &'static str) {
89        v.some(key, self, description)
90    }
91
92    fn set(&mut self, _key: &str, value: &str) -> Result<()> {
93        *self = ByteSize::parse(value)?;
94        Ok(())
95    }
96}
97
98/// Which notices a session prints (`ursa.notices`): one line each, to the
99/// sink the host installed (`crate::notices`; the CLI's is stderr).
100#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
101pub enum NoticeLevel {
102    /// None.
103    Off,
104    /// Index builds that take more than a second, and every stored-index
105    /// decision that loses work: an index not stored, an index evicted, a
106    /// failed write.
107    #[default]
108    On,
109    /// Every index build, load and store, however quick.
110    All,
111}
112
113impl NoticeLevel {
114    /// `on`, `off` or `all` (`true` and `false` too), in any case.
115    pub fn parse(text: &str) -> Result<Self> {
116        match text.trim().to_ascii_lowercase().as_str() {
117            "off" | "false" => Ok(Self::Off),
118            "on" | "true" => Ok(Self::On),
119            "all" => Ok(Self::All),
120            _ => Err(DataFusionError::Configuration(format!(
121                "'{text}' is not a notice level; use 'on', 'off' or 'all'"
122            ))),
123        }
124    }
125}
126
127impl fmt::Display for NoticeLevel {
128    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
129        f.write_str(match self {
130            Self::Off => "off",
131            Self::On => "on",
132            Self::All => "all",
133        })
134    }
135}
136
137impl ConfigField for NoticeLevel {
138    fn visit<V: Visit>(&self, v: &mut V, key: &str, description: &'static str) {
139        v.some(key, self, description)
140    }
141
142    fn set(&mut self, _key: &str, value: &str) -> Result<()> {
143        *self = NoticeLevel::parse(value)?;
144        Ok(())
145    }
146}
147
148extensions_options! {
149    /// Ursa's settings (`SET ursa.<name> = …`).
150    pub struct UrsaOptions {
151        /// The most bytes of graph indexes a `.ursa` database keeps on disk.
152        /// Past it, the least recently used are removed to make room for a
153        /// new one. Default 4GiB. SET in a writable database, it is stored
154        /// in the database.
155        pub index_storage_limit: ByteSize, default = ByteSize(DEFAULT_INDEX_STORAGE_LIMIT)
156        /// Which notices to print on stderr: 'on' (slow index builds, and
157        /// indexes not stored or evicted), 'all' (every index build, load
158        /// and store) or 'off'. Default on.
159        pub notices: NoticeLevel, default = NoticeLevel::On
160    }
161}
162
163impl ConfigExtension for UrsaOptions {
164    const PREFIX: &'static str = "ursa";
165}
166
167/// `config` with Ursa's settings, unless it has them.
168pub(crate) fn with_options(mut config: SessionConfig) -> SessionConfig {
169    if config.options().extensions.get::<UrsaOptions>().is_none() {
170        config
171            .options_mut()
172            .extensions
173            .insert(UrsaOptions::default());
174    }
175    config
176}
177
178/// The session's `ursa.notices`.
179pub fn notice_level(config: &SessionConfig) -> NoticeLevel {
180    config
181        .options()
182        .extensions
183        .get::<UrsaOptions>()
184        .map_or(NoticeLevel::default(), |o| o.notices)
185}
186
187/// The session's `ursa.index_storage_limit`, in bytes: what `SHOW` says. The
188/// database's stored indexes obey the database's limit
189/// (`Database::index_storage_limit`), which the SQL layer keeps in step with
190/// this in a writable database.
191pub fn index_storage_limit(config: &SessionConfig) -> u64 {
192    config
193        .options()
194        .extensions
195        .get::<UrsaOptions>()
196        .map_or(DEFAULT_INDEX_STORAGE_LIMIT, |o| o.index_storage_limit.0)
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    #[test]
204    fn sizes_parse_and_print() {
205        for (text, bytes) in [
206            ("0", 0),
207            ("1048576", 1 << 20),
208            ("512KiB", 512 << 10),
209            ("64 MiB", 64 << 20),
210            ("8GiB", 8 << 30),
211            ("8gib", 8 << 30),
212            ("1.5GiB", 3 << 29),
213            ("2TiB", 2 << 40),
214            ("100B", 100),
215        ] {
216            assert_eq!(ByteSize::parse(text).unwrap(), ByteSize(bytes), "{text}");
217        }
218        for bad in [
219            "",
220            "GiB",
221            "8GB",
222            "-1",
223            "8 gigabytes",
224            "1e3",
225            "99999999999TiB",
226        ] {
227            assert!(ByteSize::parse(bad).is_err(), "{bad}");
228        }
229        assert_eq!(ByteSize(4 << 30).to_string(), "4GiB");
230        assert_eq!(ByteSize(1536 << 20).to_string(), "1536MiB");
231        assert_eq!(ByteSize(1000).to_string(), "1000");
232        assert_eq!(ByteSize(0).to_string(), "0");
233    }
234}
ursa_engine/options.rs

ursa_engine/
options.rs
