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. They are not kept in the `.ursa` file: the
5//! catalog stores definitions, and a setting is a property of the process
6//! using the database, as DuckDB's `memory_limit` is.
7
8use std::fmt;
9
10use datafusion::common::extensions_options;
11use datafusion::config::{ConfigExtension, ConfigField, Visit};
12use datafusion::error::{DataFusionError, Result};
13use datafusion::prelude::SessionConfig;
14
15/// The default `ursa.index_storage_limit`: 4 GiB.
16pub const DEFAULT_INDEX_STORAGE_LIMIT: u64 = 4 << 30;
17
18/// A number of bytes, written as plain bytes (`1048576`) or with a binary
19/// unit (`512KiB`, `64MiB`, `8GiB`, `1TiB`; `B` for bytes).
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub struct ByteSize(pub u64);
22
23const UNITS: [(&str, u64); 5] = [
24    ("TiB", 1 << 40),
25    ("GiB", 1 << 30),
26    ("MiB", 1 << 20),
27    ("KiB", 1 << 10),
28    ("B", 1),
29];
30
31impl ByteSize {
32    /// Parse `512KiB`, `8 GiB`, `1.5GiB` or `1048576`. Units are binary and
33    /// matched without regard to case.
34    pub fn parse(text: &str) -> Result<Self> {
35        let invalid = || {
36            DataFusionError::Configuration(format!(
37                "'{text}' is not a size; write bytes (1048576) or a number with KiB, MiB, GiB \
38                 or TiB (8GiB)"
39            ))
40        };
41        let text = text.trim();
42        let split = text
43            .find(|c: char| !(c.is_ascii_digit() || c == '.'))
44            .unwrap_or(text.len());
45        let (number, unit) = (text[..split].trim(), text[split..].trim());
46        let scale = match unit {
47            "" => 1,
48            unit => {
49                UNITS
50                    .iter()
51                    .find(|(name, _)| name.eq_ignore_ascii_case(unit))
52                    .ok_or_else(invalid)?
53                    .1
54            }
55        };
56        if let Ok(n) = number.parse::<u64>() {
57            return n.checked_mul(scale).map(ByteSize).ok_or_else(invalid);
58        }
59        let n: f64 = number.parse().map_err(|_| invalid())?;
60        let bytes = n * scale as f64;
61        if !bytes.is_finite() || bytes < 0.0 || bytes >= u64::MAX as f64 {
62            return Err(invalid());
63        }
64        Ok(ByteSize(bytes as u64))
65    }
66}
67
68impl fmt::Display for ByteSize {
69    /// In the largest binary unit that divides it, as `SHOW` prints it.
70    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71        let (name, scale) = UNITS
72            .iter()
73            .find(|(_, scale)| self.0 != 0 && self.0.is_multiple_of(*scale))
74            .copied()
75            .unwrap_or(("B", 1));
76        if name == "B" {
77            write!(f, "{}", self.0)
78        } else {
79            write!(f, "{}{name}", self.0 / scale)
80        }
81    }
82}
83
84impl ConfigField for ByteSize {
85    fn visit<V: Visit>(&self, v: &mut V, key: &str, description: &'static str) {
86        v.some(key, self, description)
87    }
88
89    fn set(&mut self, _key: &str, value: &str) -> Result<()> {
90        *self = ByteSize::parse(value)?;
91        Ok(())
92    }
93}
94
95extensions_options! {
96    /// Ursa's settings (`SET ursa.<name> = …`).
97    pub struct UrsaOptions {
98        /// The most bytes of graph indexes a `.ursa` database keeps on disk.
99        /// Past it, the least recently used are removed. Default 4GiB.
100        pub index_storage_limit: ByteSize, default = ByteSize(DEFAULT_INDEX_STORAGE_LIMIT)
101    }
102}
103
104impl ConfigExtension for UrsaOptions {
105    const PREFIX: &'static str = "ursa";
106}
107
108/// `config` with Ursa's settings, unless it has them.
109pub(crate) fn with_options(mut config: SessionConfig) -> SessionConfig {
110    if config.options().extensions.get::<UrsaOptions>().is_none() {
111        config
112            .options_mut()
113            .extensions
114            .insert(UrsaOptions::default());
115    }
116    config
117}
118
119/// The session's `ursa.index_storage_limit`, in bytes.
120pub fn index_storage_limit(config: &SessionConfig) -> u64 {
121    config
122        .options()
123        .extensions
124        .get::<UrsaOptions>()
125        .map_or(DEFAULT_INDEX_STORAGE_LIMIT, |o| o.index_storage_limit.0)
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    #[test]
133    fn sizes_parse_and_print() {
134        for (text, bytes) in [
135            ("0", 0),
136            ("1048576", 1 << 20),
137            ("512KiB", 512 << 10),
138            ("64 MiB", 64 << 20),
139            ("8GiB", 8 << 30),
140            ("8gib", 8 << 30),
141            ("1.5GiB", 3 << 29),
142            ("2TiB", 2 << 40),
143            ("100B", 100),
144        ] {
145            assert_eq!(ByteSize::parse(text).unwrap(), ByteSize(bytes), "{text}");
146        }
147        for bad in [
148            "",
149            "GiB",
150            "8GB",
151            "-1",
152            "8 gigabytes",
153            "1e3",
154            "99999999999TiB",
155        ] {
156            assert!(ByteSize::parse(bad).is_err(), "{bad}");
157        }
158        assert_eq!(ByteSize(4 << 30).to_string(), "4GiB");
159        assert_eq!(ByteSize(1536 << 20).to_string(), "1536MiB");
160        assert_eq!(ByteSize(1000).to_string(), "1000");
161        assert_eq!(ByteSize(0).to_string(), "0");
162    }
163}
ursa_engine/options.rs

ursa_engine/
options.rs
