Skip to main content

latkit/data/
query.rs

1//! Reading data back: the rows of a table a query keeps, in the order it sorts them, a page of them,
2//! with their fields or what expressions make of them. Queries run on the CPU through the same
3//! expressions marks draw with, a plain field read page by page, and return a table, so what comes
4//! back has the shape of what goes in.
5
6use std::cmp::Ordering;
7
8use super::ids::Ids;
9use super::name::Name;
10use super::ordered::Ordered;
11use super::table::{Field, Table};
12use super::values::Values;
13use super::{Data, column::Column};
14use crate::expr::Expr;
15use crate::expr::cpu::Rows;
16use crate::{Error, Result};
17
18/// The rows of a table to read back: of the rows `ids` names (every row unless set), those `filter`
19/// keeps, sorted by `sort`, a page of them from `offset` of at most `limit`, with the fields `select`
20/// names or computes; sampled fields read at axis coordinate `at`, their last frame by default. Plain
21/// data, like a scene.
22///
23/// ```
24/// use latkit::{Data, Ids, Order, Query, Table, field};
25/// let mut data = Data::new("grid");
26/// data.apply(Table::new("Branch", Ids::count(3)).field("loading", vec![40.0, 95.0, 120.0]).field("name", vec!["a", "b", "c"]))?;
27/// let (hot, kept) = data.query(&Query::of("Branch").filter(field("loading").gt(90.0)).sort(field("loading"), Order::Descending).limit(1))?;
28/// assert_eq!(hot.get(2_u64, "name"), Some(latkit::Value::Text("c".into())));
29/// assert_eq!((hot.len(), kept), (1, 2));
30/// # Ok::<(), latkit::Error>(())
31/// ```
32#[derive(Clone, PartialEq, Debug, serde::Deserialize, schemars::JsonSchema)]
33#[serde(deny_unknown_fields)]
34#[schemars(description = "The rows of a table to read back, run through the expressions marks draw with.")]
35pub struct Query {
36    /// The table read.
37    of: Name,
38    /// The rows read, in this order unless sorted; every row unless set. Ids that name no row are
39    /// passed over.
40    #[serde(default)]
41    ids: Option<Ids>,
42    /// Keeps the rows it is true for.
43    #[serde(default)]
44    filter: Option<Expr>,
45    /// Keys to sort by, each after those before it; rows a key reads nothing for come last, and rows
46    /// all keys order alike keep their order.
47    #[serde(default)]
48    sort: Vec<(Expr, Order)>,
49    /// The kept rows skipped first.
50    #[serde(default)]
51    offset: usize,
52    /// The most rows read; a limit picks them in linear time before sorting only those.
53    limit: Option<usize>,
54    /// Fields to read by name: a field as it is, anything else as numbers. Every field unless set.
55    #[serde(default)]
56    select: Ordered<Expr>,
57    /// The axis coordinate sampled fields are read at; their last frame unless set.
58    at: Option<f64>,
59}
60
61/// Which way a sort runs; rows a key reads nothing for come last either way.
62#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, serde::Deserialize, schemars::JsonSchema)]
63#[serde(rename_all = "lowercase")]
64pub enum Order {
65    /// Least first.
66    #[default]
67    Ascending,
68    /// Greatest first.
69    Descending,
70}
71
72impl Query {
73    /// Every row of table `table`, every field.
74    pub fn of(table: impl Into<Name>) -> Query {
75        Query {
76            of: table.into(),
77            ids: None,
78            filter: None,
79            sort: Vec::new(),
80            offset: 0,
81            limit: None,
82            select: Ordered::default(),
83            at: None,
84        }
85    }
86
87    /// Reads only the rows `ids` names, in that order unless sorted.
88    #[must_use]
89    pub fn ids(mut self, ids: impl Into<Ids>) -> Query {
90        self.ids = Some(ids.into());
91        self
92    }
93
94    /// Keeps the rows `keep` is true for.
95    #[must_use]
96    pub fn filter(mut self, keep: impl Into<Expr>) -> Query {
97        self.filter = Some(keep.into());
98        self
99    }
100
101    /// Sorts by `key`, after the keys given before; rows it orders alike keep their order.
102    #[must_use]
103    pub fn sort(mut self, key: impl Into<Expr>, order: Order) -> Query {
104        self.sort.push((key.into(), order));
105        self
106    }
107
108    /// Skips the first `offset` rows kept.
109    #[must_use]
110    pub fn offset(mut self, offset: usize) -> Query {
111        self.offset = offset;
112        self
113    }
114
115    /// Reads at most `limit` rows.
116    #[must_use]
117    pub fn limit(mut self, limit: usize) -> Query {
118        self.limit = Some(limit);
119        self
120    }
121
122    /// Reads field `name` as `value` makes it: a field as it is, anything else as numbers. With none
123    /// selected, every field is read.
124    #[must_use]
125    pub fn select(mut self, name: impl Into<Name>, value: impl Into<Expr>) -> Query {
126        self.select.0.push((name.into(), value.into()));
127        self
128    }
129
130    /// Reads sampled fields at axis coordinate `at`.
131    #[must_use]
132    pub fn at(mut self, at: f64) -> Query {
133        self.at = Some(at);
134        self
135    }
136}
137
138impl Data {
139    /// The rows of a table `query` asks for, as a table of their own: their ids, and their fields or
140    /// what the query computes, sampled fields read at its coordinate; and how many rows its filter
141    /// kept, before the page. A limit picks its rows in linear time before sorting only them.
142    pub fn query(&self, query: &Query) -> Result<(Table, usize)> {
143        let table = self.table(&query.of).ok_or_else(|| Error::invalid("no such table").at(query.of.as_str()))?;
144        let at = query.at.unwrap_or(f64::INFINITY);
145        let all = std::slice::from_ref(self);
146        let rows = Rows::new(table, all).at(at);
147        let named: Vec<usize> = match &query.ids {
148            Some(ids) => (0..ids.len()).filter_map(|k| table.row(ids.key(k)?)).collect(),
149            None => (0..table.len()).collect(),
150        };
151        let mut picked: Vec<usize> = match &query.filter {
152            Some(keep) => {
153                let kept = rows.numbers(keep, |x| x > 0.5);
154                named.into_iter().filter(|&r| kept[r]).collect()
155            }
156            None => named,
157        };
158        let total = picked.len();
159        let end = query.limit.map_or(picked.len(), |limit| query.offset.saturating_add(limit).min(picked.len()));
160        if !query.sort.is_empty() {
161            let keys: Vec<(Vec<f64>, Order)> =
162                query.sort.iter().map(|(key, order)| (rows.numbers(key, |x| x), *order)).collect();
163            let compare = |a: &usize, b: &usize| {
164                keys.iter()
165                    .map(|(values, order)| compare(values[*a], values[*b], *order))
166                    .find(|o| o.is_ne())
167                    .unwrap_or(Ordering::Equal)
168                    .then(a.cmp(b))
169            };
170            if end < picked.len() {
171                picked.select_nth_unstable_by(end, compare);
172                picked.truncate(end);
173            }
174            picked.sort_unstable_by(compare);
175        }
176        let page = &picked[query.offset.min(end)..end];
177        let version = super::name::Version::next();
178        let mut out = Table {
179            name: table.name.clone(),
180            ids: table.ids.gather(page),
181            index: std::sync::Arc::default(),
182            fields: Vec::new(),
183            structure: version,
184            order: version,
185        };
186        if query.select.0.is_empty() {
187            for (name, field) in &table.fields {
188                out.insert(name.clone(), plain(field, at, page));
189            }
190            return Ok((out, total));
191        }
192        for (name, value) in &query.select.0 {
193            let field = if let Some(f) = value.field_name() {
194                let Some(field) = table.field_named(f) else {
195                    return Err(Error::invalid(format!("no field {f}")).at(name.as_str()).at("select"));
196                };
197                plain(field, at, page)
198            } else {
199                let numbers = rows.numbers(value, |x| x.is_finite().then_some(x));
200                let values = Values::from(page.iter().map(|&r| numbers[r]).collect::<Vec<_>>());
201                Field::Values(Column::new(&values), super::name::Version::next())
202            };
203            out.insert(name.clone(), field);
204        }
205        Ok((out, total))
206    }
207}
208
209/// How two sort keys order rows: missing last, whichever way the sort runs.
210fn compare(a: f64, b: f64, order: Order) -> Ordering {
211    match (a.is_finite(), b.is_finite()) {
212        (true, true) if order == Order::Descending => b.total_cmp(&a),
213        (true, true) => a.total_cmp(&b),
214        (true, false) => Ordering::Less,
215        (false, true) => Ordering::Greater,
216        (false, false) => Ordering::Equal,
217    }
218}
219
220/// `field` of rows `rows` as plain values: a sampled field read at `at`, one frame, not its history.
221fn plain(field: &Field, at: f64, rows: &[usize]) -> Field {
222    match field {
223        Field::Sampled(samples, _) => {
224            let frame = samples.frame_at(at);
225            let values: Vec<Option<f64>> = rows.iter().map(|&row| frame.and_then(|f| samples.value(row, f))).collect();
226            Field::Values(Column::new(&Values::from(values)), super::name::Version::next())
227        }
228        other => other.gather(rows),
229    }
230}