1use 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#[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 of: Name,
38 #[serde(default)]
41 ids: Option<Ids>,
42 #[serde(default)]
44 filter: Option<Expr>,
45 #[serde(default)]
48 sort: Vec<(Expr, Order)>,
49 #[serde(default)]
51 offset: usize,
52 limit: Option<usize>,
54 #[serde(default)]
56 select: Ordered<Expr>,
57 at: Option<f64>,
59}
60
61#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, serde::Deserialize, schemars::JsonSchema)]
63#[serde(rename_all = "lowercase")]
64pub enum Order {
65 #[default]
67 Ascending,
68 Descending,
70}
71
72impl Query {
73 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 #[must_use]
89 pub fn ids(mut self, ids: impl Into<Ids>) -> Query {
90 self.ids = Some(ids.into());
91 self
92 }
93
94 #[must_use]
96 pub fn filter(mut self, keep: impl Into<Expr>) -> Query {
97 self.filter = Some(keep.into());
98 self
99 }
100
101 #[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 #[must_use]
110 pub fn offset(mut self, offset: usize) -> Query {
111 self.offset = offset;
112 self
113 }
114
115 #[must_use]
117 pub fn limit(mut self, limit: usize) -> Query {
118 self.limit = Some(limit);
119 self
120 }
121
122 #[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 #[must_use]
132 pub fn at(mut self, at: f64) -> Query {
133 self.at = Some(at);
134 self
135 }
136}
137
138impl Data {
139 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
209fn 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
220fn 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}