datui_lib/query/
sql_group.rs

1//! What a SQL `GROUP BY` result was grouped from, read from the statement itself, so
2//! Enter on a row of the result can show the rows behind it.
3//!
4//! Only a plain grouping of the one table qualifies: `SELECT keys…, aggregates… FROM df
5//! [WHERE …] GROUP BY keys [HAVING …] [ORDER BY …] [LIMIT …]`, each key selected. Joins,
6//! subqueries, CTEs, set operations, window functions, `DISTINCT ON`, `GROUP BY ALL` and
7//! keys that are not in the select list give no plan, so drill-down stays off rather
8//! than guess.
9
10use std::ops::ControlFlow;
11
12use sqlparser::ast::{
13    Distinct, Expr, FunctionArguments, GroupByExpr, Ident, ObjectNamePart, Query, Select,
14    SelectItem, SetExpr, Statement, TableFactor, Visit, Visitor, WildcardAdditionalOptions,
15    visit_expressions_mut,
16};
17
18/// How a grouped statement relates its result to its source rows.
19#[derive(Debug, PartialEq)]
20pub struct GroupPlan {
21    /// A statement over the same table and `WHERE` that returns every source row, each
22    /// computed key added as a column named by [`KeySource::Computed`].
23    pub source_sql: String,
24    pub keys: Vec<PlanKey>,
25    /// Whether the statement says which groups come back in which order: an ORDER BY,
26    /// or a LIMIT that picks some.
27    pub ordered: bool,
28}
29
30/// One grouping key: where it sits in the result and how the source computes it.
31#[derive(Debug, PartialEq)]
32pub struct PlanKey {
33    /// Position of the key's column in the result.
34    pub result_index: usize,
35    pub source: KeySource,
36}
37
38#[derive(Debug, PartialEq)]
39pub enum KeySource {
40    /// A column of the source as it stands.
41    Column(String),
42    /// An expression, added to the source rows under this name.
43    Computed(String),
44}
45
46/// Prefix of the columns `source_sql` adds for computed keys; the drill drops them.
47const KEY_PREFIX: &str = "__datui_group_key_";
48
49/// The plan for `sql`, run against a table whose columns are `columns` and giving a
50/// result `result_width` columns wide, or `None` when the statement is not a grouping
51/// whose rows can be traced back reliably.
52pub fn plan(sql: &str, columns: &[&str], result_width: usize) -> Option<GroupPlan> {
53    let statements = super::sql_parse::statements(sql).ok()?;
54    let [Statement::Query(query)] = statements.as_slice() else {
55        return None;
56    };
57    if !plain_query(query) || !Plain::check(query) {
58        return None;
59    }
60    let SetExpr::Select(select) = query.body.as_ref() else {
61        return None;
62    };
63    if !plain_select(select) {
64        return None;
65    }
66    let GroupByExpr::Expressions(group_by, modifiers) = &select.group_by else {
67        return None;
68    };
69    if group_by.is_empty() || !modifiers.is_empty() {
70        return None;
71    }
72    let [from] = select.from.as_slice() else {
73        return None;
74    };
75    let TableFactor::Table { name, alias, .. } = &from.relation else {
76        return None;
77    };
78    // Every select item is one result column, in order; a wildcard or a
79    // multi-column alias breaks that.
80    let items: Vec<(&Expr, Option<&Ident>)> = select
81        .projection
82        .iter()
83        .map(|item| match item {
84            SelectItem::UnnamedExpr(e) => Some((e, None)),
85            SelectItem::ExprWithAlias { expr, alias } => Some((expr, Some(alias))),
86            _ => None,
87        })
88        .collect::<Option<_>>()?;
89    if items.len() != result_width {
90        return None;
91    }
92    // `df.dept` and `t.dept` name the column `dept`.
93    let qualifiers: Vec<&str> = name
94        .0
95        .last()
96        .and_then(|p| p.as_ident())
97        .into_iter()
98        .chain(alias.as_ref().map(|a| &a.name))
99        .map(|ident| ident.value.as_str())
100        .collect();
101    let normalize = |e: &Expr| normalized(e, &qualifiers);
102
103    let mut keys = Vec::with_capacity(group_by.len());
104    let mut computed = Vec::new();
105    for key in group_by {
106        let (index, expr) = resolve_key(key, &items, columns, &normalize)?;
107        let source = match normalize(expr) {
108            Expr::Identifier(ident) if columns.contains(&ident.value.as_str()) => {
109                KeySource::Column(ident.value)
110            }
111            _ => {
112                let name = format!("{KEY_PREFIX}{}", computed.len());
113                computed.push(format!("{expr} AS \"{name}\""));
114                KeySource::Computed(name)
115            }
116        };
117        keys.push(PlanKey {
118            result_index: index,
119            source,
120        });
121    }
122
123    let mut source_sql = String::from("SELECT *");
124    for column in &computed {
125        source_sql.push_str(", ");
126        source_sql.push_str(column);
127    }
128    source_sql.push_str(&format!(" FROM {from}"));
129    if let Some(selection) = &select.selection {
130        source_sql.push_str(&format!(" WHERE {selection}"));
131    }
132    let ordered = query.order_by.is_some() || query.limit_clause.is_some() || query.fetch.is_some();
133    Some(GroupPlan {
134        source_sql,
135        keys,
136        ordered,
137    })
138}
139
140/// The columns of `sql`'s result, named in `result`, that are columns of the table it
141/// reads (named in `columns`) unchanged, renamed or not: each result name with the
142/// column's. Only a plain statement's select list is read: `*`, a column, or a column
143/// `AS` a name. Anything else, or a statement that is not plain, gives none.
144pub fn passed_through(sql: &str, columns: &[&str], result: &[&str]) -> Vec<(String, String)> {
145    let Ok(statements) = super::sql_parse::statements(sql) else {
146        return Vec::new();
147    };
148    let [Statement::Query(query)] = statements.as_slice() else {
149        return Vec::new();
150    };
151    if !plain_query(query) || !Plain::check(query) {
152        return Vec::new();
153    }
154    let SetExpr::Select(select) = query.body.as_ref() else {
155        return Vec::new();
156    };
157    if !plain_select(select) {
158        return Vec::new();
159    }
160    let [from] = select.from.as_slice() else {
161        return Vec::new();
162    };
163    let TableFactor::Table { name, alias, .. } = &from.relation else {
164        return Vec::new();
165    };
166    let qualifiers: Vec<&str> = name
167        .0
168        .last()
169        .and_then(|p| p.as_ident())
170        .into_iter()
171        .chain(alias.as_ref().map(|a| &a.name))
172        .map(|ident| ident.value.as_str())
173        .collect();
174    let column = |e: &Expr| match normalized(e, &qualifiers) {
175        Expr::Identifier(ident) if columns.contains(&ident.value.as_str()) => Some(ident.value),
176        _ => None,
177    };
178    let plain = |options: &WildcardAdditionalOptions| {
179        options.opt_ilike.is_none()
180            && options.opt_exclude.is_none()
181            && options.opt_except.is_none()
182            && options.opt_replace.is_none()
183            && options.opt_rename.is_none()
184    };
185    // A name given twice is an error in Polars, so a computed column never shares a
186    // name with one kept here.
187    let mut kept: Vec<(String, String)> = Vec::new();
188    for item in &select.projection {
189        match item {
190            SelectItem::UnnamedExpr(e) => kept.extend(column(e).map(|c| (c.clone(), c))),
191            SelectItem::ExprWithAlias { expr, alias } => {
192                kept.extend(column(expr).map(|c| (alias.value.clone(), c)));
193            }
194            SelectItem::Wildcard(options) | SelectItem::QualifiedWildcard(_, options)
195                if plain(options) =>
196            {
197                kept.extend(columns.iter().map(|c| (c.to_string(), c.to_string())));
198            }
199            _ => {}
200        }
201    }
202    kept.retain(|(shown, _)| result.contains(&shown.as_str()));
203    kept
204}
205
206/// The select item a `GROUP BY` entry names, as Polars resolves it: an ordinal, then a
207/// select alias that is not also a source column, then the same expression selected.
208fn resolve_key<'a>(
209    key: &'a Expr,
210    items: &[(&'a Expr, Option<&Ident>)],
211    columns: &[&str],
212    normalize: &impl Fn(&Expr) -> Expr,
213) -> Option<(usize, &'a Expr)> {
214    if let Expr::Value(value) = key {
215        let ordinal: usize = value.to_string().parse().ok()?;
216        let (expr, _) = items.get(ordinal.checked_sub(1)?)?;
217        return Some((ordinal - 1, expr));
218    }
219    if let Expr::Identifier(ident) = key
220        && !columns.contains(&ident.value.as_str())
221        && let Some(index) = items
222            .iter()
223            .position(|(_, alias)| alias.is_some_and(|a| a.value == ident.value))
224    {
225        return Some((index, items[index].0));
226    }
227    let wanted = normalize(key);
228    let index = items.iter().position(|(e, _)| normalize(e) == wanted)?;
229    Some((index, key))
230}
231
232/// `e` spelled one way for comparing keys, keeping only what polars-sql reads: no
233/// parentheses, no quotes around names (`"dept"` is `dept`, and case always counts),
234/// no table name in front of a column, and function names in lowercase.
235fn normalized(e: &Expr, qualifiers: &[&str]) -> Expr {
236    let mut e = e.clone();
237    let _ = visit_expressions_mut(&mut e, |e| {
238        match e {
239            Expr::Nested(inner) => *e = inner.as_ref().clone(),
240            Expr::Identifier(ident) => ident.quote_style = None,
241            Expr::CompoundIdentifier(parts) => {
242                for part in parts.iter_mut() {
243                    part.quote_style = None;
244                }
245                if let [table, column] = parts.as_slice()
246                    && qualifiers.contains(&table.value.as_str())
247                {
248                    *e = Expr::Identifier(column.clone());
249                }
250            }
251            Expr::Function(f) => {
252                for part in f.name.0.iter_mut() {
253                    if let ObjectNamePart::Identifier(ident) = part {
254                        ident.value = ident.value.to_lowercase();
255                        ident.quote_style = None;
256                    }
257                }
258            }
259            _ => {}
260        }
261        ControlFlow::<()>::Continue(())
262    });
263    e
264}
265
266/// No clause beyond ORDER BY and LIMIT around the one SELECT.
267fn plain_query(query: &Query) -> bool {
268    query.with.is_none()
269        && matches!(query.body.as_ref(), SetExpr::Select(_))
270        && query.locks.is_empty()
271        && query.for_clause.is_none()
272        && query.settings.is_none()
273        && query.format_clause.is_none()
274        && query.pipe_operators.is_empty()
275}
276
277/// One table, no join, and nothing that reshapes rows before or beside the grouping.
278fn plain_select(select: &Select) -> bool {
279    let one_table = match select.from.as_slice() {
280        [from] => {
281            from.joins.is_empty()
282                && matches!(
283                    &from.relation,
284                    TableFactor::Table {
285                        alias,
286                        args: None,
287                        sample: None,
288                        version: None,
289                        with_ordinality: false,
290                        json_path: None,
291                        ..
292                    } if alias.as_ref().is_none_or(|a| a.columns.is_empty())
293                )
294        }
295        _ => false,
296    };
297    one_table
298        && matches!(
299            select.distinct,
300            None | Some(Distinct::Distinct | Distinct::All)
301        )
302        && select.top.is_none()
303        && select.exclude.is_none()
304        && select.into.is_none()
305        && select.lateral_views.is_empty()
306        && select.prewhere.is_none()
307        && select.connect_by.is_empty()
308        && select.cluster_by.is_empty()
309        && select.distribute_by.is_empty()
310        && select.sort_by.is_empty()
311        && select.named_window.is_empty()
312        && select.qualify.is_none()
313        && select.value_table_mode.is_none()
314}
315
316/// Walks the whole statement for what the checks on its clauses cannot see: a second
317/// query anywhere (a subquery), a window function, or `UNNEST`, which changes the rows
318/// before they are grouped.
319#[derive(Default)]
320struct Plain {
321    queries: usize,
322}
323
324impl Plain {
325    fn check(query: &Query) -> bool {
326        let mut plain = Plain::default();
327        query.visit(&mut plain).is_continue() && plain.queries == 1
328    }
329}
330
331impl Visitor for Plain {
332    type Break = ();
333
334    fn pre_visit_query(&mut self, _query: &Query) -> ControlFlow<()> {
335        self.queries += 1;
336        ControlFlow::Continue(())
337    }
338
339    fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<()> {
340        match expr {
341            Expr::Subquery(_) | Expr::InSubquery { .. } | Expr::Exists { .. } => {
342                ControlFlow::Break(())
343            }
344            Expr::Function(f)
345                if f.over.is_some()
346                    || matches!(f.args, FunctionArguments::Subquery(_))
347                    || f.name.to_string().eq_ignore_ascii_case("unnest") =>
348            {
349                ControlFlow::Break(())
350            }
351            _ => ControlFlow::Continue(()),
352        }
353    }
354}
355
356#[cfg(test)]
357mod tests;
datui_lib/query/sql_group.rs

datui_lib/query/
sql_group.rs
