Skip to main content

datui_lib/query/
mod.rs

1pub(crate) mod query_prompt;
2pub(crate) mod sql_assist;
3// Public for the `sql_group_plan` fuzz target.
4#[cfg(feature = "sql")]
5pub mod sql_group;
6#[cfg(feature = "sql")]
7pub(crate) mod sql_parse;
8#[cfg(feature = "sql")]
9pub(crate) mod sql_plan;
10
11use polars::prelude::StrptimeOptions;
12use polars::prelude::*;
13use std::ops::{Add, Mul, Rem, Sub};
14
15#[derive(Debug, Clone, PartialEq)]
16enum Token {
17    Identifier(String),
18    Number(f64),
19    String(String),
20    /// Date literal in YYYY.MM.DD format, stored as ISO "YYYY-MM-DD" for Polars
21    DateLiteral(String),
22    /// Timestamp literal YYYY.MM.DDTHH:MM:SS[.fff...]
23    TimestampLiteral {
24        iso: String,
25        format_str: String,
26        time_unit: TimeUnit,
27    },
28    Op(String),
29    LParen,
30    RParen,
31    LBracket,
32    RBracket,
33    Comma,
34    Colon,
35    Dot,
36    Select,
37    Where,
38    By,
39}
40
41/// Parse YYYY.MM.DDTHH:MM:SS[.fff...] timestamp. Consumes from chars. Returns (iso_string, format, time_unit) or None.
42fn parse_timestamp_literal(
43    date_part: &str,
44    chars: &mut std::iter::Peekable<std::str::Chars<'_>>,
45) -> Option<(String, String, TimeUnit)> {
46    if chars.peek() != Some(&'T') {
47        return None;
48    }
49    chars.next(); // consume 'T'
50    let mut time_part = String::new();
51    while let Some(&c) = chars.peek() {
52        if c.is_ascii_digit() || c == ':' || c == '.' {
53            time_part.push(c);
54            chars.next();
55        } else {
56            break;
57        }
58    }
59    let parts: Vec<&str> = time_part.split(':').collect();
60    if parts.len() != 3 {
61        return None;
62    }
63    let (h, m, s) = (parts[0], parts[1], parts[2]);
64    if h.len() != 2 || m.len() != 2 || s.len() < 2 {
65        return None;
66    }
67    let (sec_part, frac) = match s.split_once('.') {
68        Some((a, f)) => (a, f),
69        None => (s, ""),
70    };
71    let (time_unit, format_str) = match frac.len() {
72        0 => (TimeUnit::Microseconds, "%Y-%m-%dT%H:%M:%S".to_string()),
73        1..=3 => (TimeUnit::Milliseconds, "%Y-%m-%dT%H:%M:%S%.3f".to_string()),
74        4..=6 => (TimeUnit::Microseconds, "%Y-%m-%dT%H:%M:%S%.6f".to_string()),
75        7..=9 => (TimeUnit::Nanoseconds, "%Y-%m-%dT%H:%M:%S%.9f".to_string()),
76        _ => (TimeUnit::Nanoseconds, "%Y-%m-%dT%H:%M:%S%.9f".to_string()),
77    };
78    let iso_date = parse_date_literal(date_part)?;
79    let frac_padded = match time_unit {
80        TimeUnit::Milliseconds => format!("{:0<3}", frac),
81        TimeUnit::Microseconds => format!("{:0<6}", frac),
82        TimeUnit::Nanoseconds => format!("{:0<9}", frac),
83    };
84    let iso = if frac.is_empty() {
85        format!("{}T{}:{}:{}", iso_date, h, m, sec_part)
86    } else {
87        format!("{}T{}:{}:{}.{}", iso_date, h, m, sec_part, frac_padded)
88    };
89    Some((iso, format_str, time_unit))
90}
91
92/// Parse YYYY.MM.DD date literal (e.g. 2021.01.01). Returns ISO string "YYYY-MM-DD" or None.
93fn parse_date_literal(s: &str) -> Option<String> {
94    let parts: Vec<&str> = s.split('.').collect();
95    if parts.len() != 3 {
96        return None;
97    }
98    let year: u32 = parts[0].parse().ok()?;
99    let month: u32 = parts[1].parse().ok()?;
100    let day: u32 = parts[2].parse().ok()?;
101    if parts[0].len() != 4 || !(1000..=9999).contains(&year) {
102        return None;
103    }
104    if !(1..=12).contains(&month) || !(1..=31).contains(&day) {
105        return None;
106    }
107    Some(format!("{:04}-{:02}-{:02}", year, month, day))
108}
109
110fn tokenize(input: &str) -> Result<Vec<Token>, String> {
111    let mut tokens = Vec::new();
112    let mut chars = input.chars().peekable();
113
114    while let Some(&c) = chars.peek() {
115        match c {
116            ' ' | '\t' | '\n' | '\r' => {
117                chars.next();
118            }
119            ',' => {
120                tokens.push(Token::Comma);
121                chars.next();
122            }
123            ':' => {
124                tokens.push(Token::Colon);
125                chars.next();
126            }
127            '(' => {
128                tokens.push(Token::LParen);
129                chars.next();
130            }
131            ')' => {
132                tokens.push(Token::RParen);
133                chars.next();
134            }
135            '[' => {
136                tokens.push(Token::LBracket);
137                chars.next();
138            }
139            ']' => {
140                tokens.push(Token::RBracket);
141                chars.next();
142            }
143            '"' => {
144                chars.next(); // consume opening quote
145                let mut string_val = String::new();
146                let mut found_closing_quote = false;
147                while let Some(&c) = chars.peek() {
148                    if c == '\\' {
149                        chars.next(); // consume backslash
150                        if let Some(&next_c) = chars.peek() {
151                            match next_c {
152                                'n' => {
153                                    string_val.push('\n');
154                                    chars.next();
155                                }
156                                't' => {
157                                    string_val.push('\t');
158                                    chars.next();
159                                }
160                                'r' => {
161                                    string_val.push('\r');
162                                    chars.next();
163                                }
164                                '\\' => {
165                                    string_val.push('\\');
166                                    chars.next();
167                                }
168                                '"' => {
169                                    string_val.push('"');
170                                    chars.next();
171                                }
172                                _ => {
173                                    // Unknown escape, just include the backslash and next char
174                                    string_val.push('\\');
175                                    string_val.push(next_c);
176                                    chars.next();
177                                }
178                            }
179                        } else {
180                            return Err("Unterminated escape sequence in string".to_string());
181                        }
182                    } else if c == '"' {
183                        chars.next(); // consume closing quote
184                        found_closing_quote = true;
185                        break;
186                    } else {
187                        string_val.push(c);
188                        chars.next();
189                    }
190                }
191                if !found_closing_quote {
192                    return Err("Unterminated string literal".to_string());
193                }
194                tokens.push(Token::String(string_val));
195            }
196            '&' | '|' => {
197                chars.next();
198                if let Some(&next) = chars.peek()
199                    && (next == '&' || next == '|')
200                {
201                    let fix = if next == c {
202                        format!("use `{c}`")
203                    } else {
204                        "use `&` or `|`".to_string()
205                    };
206                    return Err(format!("`{c}{next}` is not q: {fix}"));
207                }
208                tokens.push(Token::Op(c.to_string()));
209            }
210            '^' => {
211                tokens.push(Token::Op("^".to_string()));
212                chars.next();
213            }
214            '+' | '-' | '*' | '%' | '/' | '=' | '<' | '>' | '!' => {
215                let mut op = c.to_string();
216                chars.next();
217                if let Some(&next_c) = chars.peek()
218                    && ((c == '<' && (next_c == '=' || next_c == '>'))
219                        || (c == '>' && next_c == '=')
220                        || (c == '!' && next_c == '='))
221                {
222                    op.push(next_c);
223                    chars.next();
224                }
225                tokens.push(Token::Op(op));
226            }
227            '.' => {
228                chars.next();
229                if chars.peek().is_some_and(|nc| nc.is_ascii_digit()) {
230                    let mut num_str = String::from('.');
231                    while let Some(&nc) = chars.peek() {
232                        if nc.is_ascii_digit() {
233                            num_str.push(nc);
234                            chars.next();
235                        } else {
236                            break;
237                        }
238                    }
239                    if let Ok(n) = num_str.parse::<f64>() {
240                        tokens.push(Token::Number(n));
241                    } else {
242                        return Err(format!("Invalid number: {}", num_str));
243                    }
244                } else {
245                    tokens.push(Token::Dot);
246                }
247            }
248            '0'..='9' => {
249                let mut num_str = String::new();
250                while let Some(&nc) = chars.peek() {
251                    if nc.is_ascii_digit() || nc == '.' {
252                        num_str.push(nc);
253                        chars.next();
254                    } else {
255                        break;
256                    }
257                }
258                // Check for YYYY.MM.DDTHH:MM:SS timestamp literal (peek for 'T' before consuming)
259                let is_timestamp =
260                    parse_date_literal(&num_str).is_some() && chars.peek() == Some(&'T');
261                if is_timestamp
262                    && let Some((iso, format_str, time_unit)) =
263                        parse_timestamp_literal(&num_str, &mut chars)
264                {
265                    tokens.push(Token::TimestampLiteral {
266                        iso,
267                        format_str,
268                        time_unit,
269                    });
270                    continue;
271                }
272                // Check for YYYY.MM.DD date literal
273                if let Some(iso) = parse_date_literal(&num_str) {
274                    tokens.push(Token::DateLiteral(iso));
275                } else if let Ok(n) = num_str.parse::<f64>() {
276                    tokens.push(Token::Number(n));
277                } else {
278                    return Err(format!("Invalid number: {}", num_str));
279                }
280            }
281            _ if c.is_alphabetic() || c == '_' => {
282                let mut ident = String::new();
283                while let Some(&nc) = chars.peek() {
284                    if nc.is_alphanumeric() || nc == '_' {
285                        ident.push(nc);
286                        chars.next();
287                    } else {
288                        break;
289                    }
290                }
291                match ident.as_str() {
292                    "select" => tokens.push(Token::Select),
293                    "where" => tokens.push(Token::Where),
294                    "by" => tokens.push(Token::By),
295                    _ => tokens.push(Token::Identifier(ident)),
296                }
297            }
298            _ => return Err(format!("Unexpected character: {}", c)),
299        }
300    }
301    Ok(tokens)
302}
303
304fn split_tokens(tokens: &[Token], delimiter: &Token) -> Vec<Vec<Token>> {
305    let mut result = Vec::new();
306    let mut current = Vec::new();
307    let mut depth = 0;
308    let mut bracket_depth = 0;
309
310    for token in tokens {
311        match token {
312            Token::LParen => depth += 1,
313            Token::RParen => depth -= 1,
314            Token::LBracket => bracket_depth += 1,
315            Token::RBracket => bracket_depth -= 1,
316            _ => {}
317        }
318
319        if depth == 0 && bracket_depth == 0 && token == delimiter {
320            result.push(current);
321            current = Vec::new();
322        } else {
323            current.push(token.clone());
324        }
325    }
326    result.push(current);
327    result
328}
329
330/// The token as the user typed it, for error messages.
331fn token_text(token: &Token) -> String {
332    match token {
333        Token::Identifier(s) => s.clone(),
334        Token::Number(n) => n.to_string(),
335        Token::String(s) => format!("\"{}\"", s),
336        Token::DateLiteral(iso) => iso.clone(),
337        Token::TimestampLiteral { iso, .. } => iso.clone(),
338        Token::Op(op) => op.clone(),
339        Token::LParen => "(".to_string(),
340        Token::RParen => ")".to_string(),
341        Token::LBracket => "[".to_string(),
342        Token::RBracket => "]".to_string(),
343        Token::Comma => ",".to_string(),
344        Token::Colon => ":".to_string(),
345        Token::Dot => ".".to_string(),
346        Token::Select => "select".to_string(),
347        Token::Where => "where".to_string(),
348        Token::By => "by".to_string(),
349    }
350}
351
352/// Remedy shown when a clause keyword turns up out of place.
353const CLAUSE_ORDER: &str = "clause order is select [by group] [where conditions]";
354
355/// [`CLAUSE_ORDER`] with q's optional `from df`, for errors about it.
356const FROM_ORDER: &str = "clause order is select [by group] [from df] [where conditions]";
357
358/// The one table q reads: the one on screen, named as SQL names it.
359const TABLE: &str = "df";
360
361/// The table name after a `from` at `tokens[i]`: an identifier, or a dotted
362/// path like `data.csv`, that is not an operator word. Returns it and the index
363/// past it.
364fn from_table_at(tokens: &[Token], i: usize) -> Option<(String, usize)> {
365    if tokens.get(i) != Some(&Token::Identifier("from".to_string())) {
366        return None;
367    }
368    let mut name = match tokens.get(i + 1) {
369        Some(Token::Identifier(n)) if !WORD_OPS.contains(&n.as_str()) => n.clone(),
370        _ => return None,
371    };
372    let mut end = i + 2;
373    while let (Some(Token::Dot), Some(Token::Identifier(part))) =
374        (tokens.get(end), tokens.get(end + 1))
375    {
376        name.push('.');
377        name.push_str(part);
378        end += 2;
379    }
380    Some((name, end))
381}
382
383/// The query body without q's `from df` (after the select list and `by`, before
384/// `where`). `from` is the clause only where a column cannot be: followed by a table
385/// name, then `where`, `by` or the end, outside brackets.
386fn strip_from(body: &[Token]) -> Result<Vec<Token>, String> {
387    let mut depth = 0i32;
388    let mut found: Option<(usize, usize)> = None;
389    for (i, token) in body.iter().enumerate() {
390        match token {
391            Token::LParen | Token::LBracket => depth += 1,
392            Token::RParen | Token::RBracket => depth -= 1,
393            _ => {}
394        }
395        if depth != 0 {
396            continue;
397        }
398        let Some((name, end)) = from_table_at(body, i) else {
399            continue;
400        };
401        let after_where = body[..i].contains(&Token::Where);
402        match body.get(end) {
403            None | Some(Token::Where) | Some(Token::By) => {}
404            _ => continue,
405        }
406        if name != TABLE {
407            return Err("q reads the table on screen, named df: … from df …".to_string());
408        }
409        if after_where {
410            return Err(format!(
411                "Unexpected 'from df' after the where clause: {FROM_ORDER}"
412            ));
413        }
414        if body.get(end) == Some(&Token::By) {
415            return Err(format!("Unexpected 'by' after 'from df': {FROM_ORDER}"));
416        }
417        found = Some((i, end));
418    }
419    let mut body = body.to_vec();
420    if let Some((start, end)) = found {
421        body.drain(start..end);
422    }
423    Ok(body)
424}
425
426/// Infix operators spelled as words (q's); ordinary identifiers elsewhere, so a column
427/// named `in` or `mod` still works at an expression's start or after `.`. `and` and
428/// `or` are `&` and `|`.
429const WORD_OPS: [&str; 7] = ["in", "like", "xbar", "mod", "wavg", "and", "or"];
430
431/// The infix operator at `tokens[i]`, if there is one: a symbol, or an operator
432/// word that has an operand before it.
433fn infix_op_at(tokens: &[Token], i: usize) -> Option<&str> {
434    match tokens.get(i)? {
435        Token::Op(op) => Some(op.as_str()),
436        Token::Identifier(word)
437            if i > 0 && tokens[i - 1] != Token::Dot && WORD_OPS.contains(&word.as_str()) =>
438        {
439            Some(word.as_str())
440        }
441        _ => None,
442    }
443}
444
445/// A parsed expression, before becoming a Polars expression ([`Node::to_expr`]) or
446/// Python ([`Node::python`]): one parse, so "Copy as Python" is the query datui ran.
447#[derive(Debug, Clone, PartialEq)]
448pub(crate) enum Node {
449    Col(String),
450    /// A number as typed: a float literal.
451    Num(f64),
452    /// A whole number where the operator keeps integers whole (`mod`, `xbar`).
453    Int(i64),
454    Str(String),
455    Bool(bool),
456    Null,
457    /// A `YYYY.MM.DD` literal, as ISO `YYYY-MM-DD`.
458    Date(String),
459    /// A `YYYY.MM.DDTHH:MM:SS[.fff]` literal.
460    Timestamp {
461        iso: String,
462        format: String,
463        unit: TimeUnit,
464        /// The zone of the column it meets, read as a clock there; none for a column
465        /// without one. See [`Node::resolve_time_zones`].
466        zone: Option<String>,
467    },
468    Bin(BinOp, Box<Node>, Box<Node>),
469    Coalesce(Box<Node>, Box<Node>),
470    /// The values of the first where the second holds.
471    Filter(Box<Node>, Box<Node>),
472    /// `when(condition).then(value).otherwise(other)`.
473    When(Box<Node>, Box<Node>, Box<Node>),
474    Op(Box<Node>, Op),
475    Alias(Box<Node>, String),
476}
477
478#[derive(Debug, Clone, Copy, PartialEq, Eq)]
479pub(crate) enum BinOp {
480    Add,
481    Sub,
482    Mul,
483    /// Division, always to a float, as q's `%`: `7 % 2` is 3.5.
484    TrueDiv,
485    /// Floor division, for `xbar`'s buckets.
486    FloorDiv,
487    Rem,
488    Eq,
489    Neq,
490    Lt,
491    Gt,
492    LtEq,
493    GtEq,
494    /// q's `&` on two booleans. Any value exceeds a null in q, so a null on either
495    /// side is null (`false & null` too), unlike Polars' Kleene `and`. A where
496    /// condition lowers it to Polars' `and` ([`Node::to_predicate`]).
497    And,
498    /// q's `|` on two booleans: a null on one side gives the other side.
499    Or,
500    /// q's `&` (Lesser) on anything else: the smaller, null when either is null.
501    /// See [`Node::resolve_types`].
502    Lesser,
503    /// q's `|` (Greater) on anything else: the larger, so a null gives the other side.
504    Greater,
505    /// Polars' Kleene `or`, for `in`'s equalities.
506    Any,
507}
508
509/// A method applied to one expression.
510#[derive(Debug, Clone, PartialEq)]
511pub(crate) enum Op {
512    Mean,
513    Min,
514    Max,
515    Count,
516    Std,
517    Var,
518    Median,
519    Sum,
520    First,
521    Last,
522    NUnique,
523    Not,
524    IsNull,
525    IsNotNull,
526    LenChars,
527    Upper,
528    Lower,
529    Abs,
530    Floor,
531    Ceil,
532    Sqrt,
533    Ln,
534    Exp,
535    Date,
536    Time,
537    Year,
538    Quarter,
539    Month,
540    Week,
541    Day,
542    OrdinalDay,
543    Weekday,
544    Hour,
545    Minute,
546    Second,
547    MonthStart,
548    MonthEnd,
549    DtFormat(String),
550    StartsWith(String),
551    EndsWith(String),
552    ContainsLiteral(String),
553    /// A regex match, strict.
554    ContainsRegex(String),
555    /// Split on the text and take the piece at the index, null past the last.
556    Part(String, i64),
557    Slice(i64, Option<u64>),
558    ReplaceAll(String, String),
559    Strip,
560    ToDate(Option<String>),
561    ToDatetime(Option<String>),
562    Round(u32),
563    /// A non-strict cast.
564    Cast(CastTo),
565}
566
567#[derive(Debug, Clone, Copy, PartialEq, Eq)]
568pub(crate) enum CastTo {
569    Int8,
570    Int16,
571    Int32,
572    Int64,
573    UInt8,
574    UInt16,
575    UInt32,
576    UInt64,
577    Float32,
578    Float64,
579    String,
580}
581
582impl CastTo {
583    fn dtype(self) -> DataType {
584        match self {
585            CastTo::Int8 => DataType::Int8,
586            CastTo::Int16 => DataType::Int16,
587            CastTo::Int32 => DataType::Int32,
588            CastTo::Int64 => DataType::Int64,
589            CastTo::UInt8 => DataType::UInt8,
590            CastTo::UInt16 => DataType::UInt16,
591            CastTo::UInt32 => DataType::UInt32,
592            CastTo::UInt64 => DataType::UInt64,
593            CastTo::Float32 => DataType::Float32,
594            CastTo::Float64 => DataType::Float64,
595            CastTo::String => DataType::String,
596        }
597    }
598
599    /// The number type `dtype` is, as a cast.
600    fn number(dtype: &DataType) -> Option<CastTo> {
601        Some(match dtype {
602            DataType::Int8 => CastTo::Int8,
603            DataType::Int16 => CastTo::Int16,
604            DataType::Int32 => CastTo::Int32,
605            DataType::Int64 => CastTo::Int64,
606            DataType::UInt8 => CastTo::UInt8,
607            DataType::UInt16 => CastTo::UInt16,
608            DataType::UInt32 => CastTo::UInt32,
609            DataType::UInt64 => CastTo::UInt64,
610            DataType::Float32 => CastTo::Float32,
611            DataType::Float64 => CastTo::Float64,
612            _ => return None,
613        })
614    }
615
616    /// The Python Polars type.
617    fn python(self) -> &'static str {
618        match self {
619            CastTo::Int8 => "pl.Int8",
620            CastTo::Int16 => "pl.Int16",
621            CastTo::Int32 => "pl.Int32",
622            CastTo::Int64 => "pl.Int64",
623            CastTo::UInt8 => "pl.UInt8",
624            CastTo::UInt16 => "pl.UInt16",
625            CastTo::UInt32 => "pl.UInt32",
626            CastTo::UInt64 => "pl.UInt64",
627            CastTo::Float32 => "pl.Float32",
628            CastTo::Float64 => "pl.Float64",
629            CastTo::String => "pl.String",
630        }
631    }
632}
633
634/// The most nodes an expression may grow to: `wavg`, `xbar` and `in` repeat operands,
635/// so nesting multiplies (found by the `parse_query` fuzz target).
636const MAX_EXPR_NODES: usize = 10_000;
637
638/// Err when `copies` of `node` would pass [`MAX_EXPR_NODES`].
639fn check_copies(node: &Node, copies: usize) -> Result<(), String> {
640    if node.size().saturating_mul(copies) > MAX_EXPR_NODES {
641        return Err(
642            "Expression is too large: nested wavg, xbar, in or & repeat what they are \
643                    given. Simplify it or split it into steps."
644                .to_string(),
645        );
646    }
647    Ok(())
648}
649
650impl Node {
651    /// How many nodes the tree has.
652    fn size(&self) -> usize {
653        1 + match self {
654            Node::Col(_)
655            | Node::Num(_)
656            | Node::Int(_)
657            | Node::Str(_)
658            | Node::Bool(_)
659            | Node::Null
660            | Node::Date(_)
661            | Node::Timestamp { .. } => 0,
662            // Lowered, `&` on numbers repeats both sides (null where either is).
663            Node::Bin(BinOp::Lesser, a, b) => 2 * (a.size() + b.size()),
664            Node::Bin(_, a, b) | Node::Coalesce(a, b) | Node::Filter(a, b) => a.size() + b.size(),
665            Node::When(a, b, c) => a.size() + b.size() + c.size(),
666            Node::Op(a, _) | Node::Alias(a, _) => a.size(),
667        }
668    }
669
670    fn op(self, op: Op) -> Node {
671        Node::Op(Box::new(self), op)
672    }
673
674    fn bin(self, op: BinOp, right: Node) -> Node {
675        Node::Bin(op, Box::new(self), Box::new(right))
676    }
677
678    fn alias(self, name: impl Into<String>) -> Node {
679        Node::Alias(Box::new(self), name.into())
680    }
681
682    fn cast_text(self) -> Node {
683        self.op(Op::Cast(CastTo::String))
684    }
685
686    /// The Polars expression.
687    pub(crate) fn to_expr(&self) -> Expr {
688        match self {
689            Node::Col(name) => col(name),
690            Node::Num(n) => lit(*n),
691            Node::Int(n) => lit(*n),
692            Node::Str(s) => lit(s.as_str()),
693            Node::Bool(b) => lit(*b),
694            Node::Null => lit(NULL),
695            Node::Date(iso) => {
696                let opts = StrptimeOptions {
697                    format: Some("%Y-%m-%d".into()),
698                    ..Default::default()
699                };
700                lit(iso.as_str()).str().to_date(opts)
701            }
702            Node::Timestamp {
703                iso,
704                format,
705                unit,
706                zone,
707            } => {
708                let opts = StrptimeOptions {
709                    format: Some(format.as_str().into()),
710                    ..Default::default()
711                };
712                // Set only from a column's own dtype, so it parses.
713                let zone = TimeZone::opt_try_new(zone.as_deref()).ok().flatten();
714                // A clock time a fall back repeats is its first instant.
715                lit(iso.as_str())
716                    .str()
717                    .to_datetime(Some(*unit), zone, opts, lit("earliest"))
718            }
719            Node::Bin(op, left, right) => {
720                let (left, right) = (left.to_expr(), right.to_expr());
721                match op {
722                    BinOp::Add => left.add(right),
723                    BinOp::Sub => left.sub(right),
724                    BinOp::Mul => left.mul(right),
725                    BinOp::TrueDiv => left.true_div(right),
726                    BinOp::FloorDiv => left.floor_div(right),
727                    BinOp::Rem => left.rem(right),
728                    BinOp::Eq => left.eq(right),
729                    BinOp::Neq => left.neq(right),
730                    BinOp::Lt => left.lt(right),
731                    BinOp::Gt => left.gt(right),
732                    BinOp::LtEq => left.lt_eq(right),
733                    BinOp::GtEq => left.gt_eq(right),
734                    BinOp::Any => left.or(right),
735                    // q: any value exceeds a null. The smaller of a null and anything
736                    // is null; the larger is the other side, as max_horizontal skips
737                    // nulls. On booleans the smaller is `and` and the larger `or`.
738                    // 1 × 1 is 1 and a null spreads, without repeating either side as
739                    // a null test would.
740                    BinOp::And => (left.cast(DataType::UInt8) * right.cast(DataType::UInt8))
741                        .cast(DataType::Boolean),
742                    BinOp::Lesser => {
743                        null_if_either(&left, &right).otherwise(min_max(true, left, right))
744                    }
745                    BinOp::Or | BinOp::Greater => min_max(false, left, right),
746                }
747            }
748            Node::Coalesce(left, right) => coalesce(&[left.to_expr(), right.to_expr()]),
749            Node::Filter(values, predicate) => values.to_expr().filter(predicate.to_expr()),
750            Node::When(condition, then, otherwise) => when(condition.to_expr())
751                .then(then.to_expr())
752                .otherwise(otherwise.to_expr()),
753            Node::Op(inner, op) => apply_op_expr(inner.to_expr(), op),
754            Node::Alias(inner, name) => inner.to_expr().alias(name.as_str()),
755        }
756    }
757
758    /// The expression as a where condition: only its true rows matter, so q's `&` and
759    /// `|` on booleans at its top (outside `not`) are Polars' `and` and `or`, which are
760    /// true on the same rows and let Polars push the condition down to the scan.
761    pub(crate) fn to_predicate(&self) -> Expr {
762        match self {
763            Node::Bin(BinOp::And, left, right) => left.to_predicate().and(right.to_predicate()),
764            Node::Bin(BinOp::Or, left, right) => left.to_predicate().or(right.to_predicate()),
765            _ => self.to_expr(),
766        }
767    }
768
769    /// [`Self::to_predicate`] as Python.
770    pub(crate) fn python_predicate(&self) -> String {
771        let operand = |node: &Node| match node {
772            Node::Bin(BinOp::And | BinOp::Or, ..) => format!("({})", node.python_predicate()),
773            _ => node.python_operand(),
774        };
775        match self {
776            Node::Bin(op @ (BinOp::And | BinOp::Or), left, right) => {
777                let symbol = if *op == BinOp::And { "&" } else { "|" };
778                format!("{} {symbol} {}", operand(left), operand(right))
779            }
780            _ => self.python(),
781        }
782    }
783
784    /// The same node with every alias inside it taken off, as Polars' `undo_aliases`.
785    pub(crate) fn without_aliases(&self) -> Node {
786        let strip = |n: &Node| Box::new(n.without_aliases());
787        match self {
788            Node::Alias(inner, _) => inner.without_aliases(),
789            Node::Bin(op, l, r) => Node::Bin(*op, strip(l), strip(r)),
790            Node::Coalesce(l, r) => Node::Coalesce(strip(l), strip(r)),
791            Node::Filter(v, p) => Node::Filter(strip(v), strip(p)),
792            Node::When(c, t, o) => Node::When(strip(c), strip(t), strip(o)),
793            Node::Op(inner, op) => Node::Op(strip(inner), op.clone()),
794            leaf => leaf.clone(),
795        }
796    }
797
798    /// Each timestamp literal meeting a zoned column (comparison, arithmetic, `^`, `?`
799    /// branches) takes that zone; Polars refuses zoned-naive comparisons.
800    pub(crate) fn resolve_time_zones(&mut self, schema: &Schema) {
801        match self {
802            Node::Bin(_, left, right) | Node::Coalesce(left, right) => {
803                left.resolve_time_zones(schema);
804                right.resolve_time_zones(schema);
805                Self::share_zone(left, right, schema);
806            }
807            Node::Filter(values, predicate) => {
808                values.resolve_time_zones(schema);
809                predicate.resolve_time_zones(schema);
810            }
811            Node::When(c, t, o) => {
812                c.resolve_time_zones(schema);
813                t.resolve_time_zones(schema);
814                o.resolve_time_zones(schema);
815                Self::share_zone(t, o, schema);
816            }
817            Node::Op(inner, _) | Node::Alias(inner, _) => inner.resolve_time_zones(schema),
818            _ => {}
819        }
820    }
821
822    /// An error for the first comparison of a typed temporal column with quoted text (a
823    /// string in q, never a date), worded in q's terms.
824    fn check_quoted_temporal(&self, schema: &Schema) -> Result<(), String> {
825        match self {
826            Node::Bin(op, left, right) => {
827                left.check_quoted_temporal(schema)?;
828                right.check_quoted_temporal(schema)?;
829                let compares = matches!(
830                    op,
831                    BinOp::Eq | BinOp::Neq | BinOp::Lt | BinOp::Gt | BinOp::LtEq | BinOp::GtEq
832                );
833                if compares
834                    && let Some(err) = quoted_temporal(left, right, schema)
835                        .or_else(|| quoted_temporal(right, left, schema))
836                {
837                    return Err(err);
838                }
839            }
840            Node::Coalesce(left, right) | Node::Filter(left, right) => {
841                left.check_quoted_temporal(schema)?;
842                right.check_quoted_temporal(schema)?;
843            }
844            Node::When(c, t, o) => {
845                c.check_quoted_temporal(schema)?;
846                t.check_quoted_temporal(schema)?;
847                o.check_quoted_temporal(schema)?;
848            }
849            Node::Op(inner, _) | Node::Alias(inner, _) => inner.check_quoted_temporal(schema)?,
850            _ => {}
851        }
852        Ok(())
853    }
854
855    /// Give a zoneless timestamp literal on one side the zone of the other side's type.
856    fn share_zone(a: &mut Node, b: &mut Node, schema: &Schema) {
857        if !Self::take_zone(a, b, schema) {
858            Self::take_zone(b, a, schema);
859        }
860    }
861
862    /// Whether `literal`, a zoneless timestamp literal, took the zone of `other`'s type.
863    fn take_zone(literal: &mut Node, other: &Node, schema: &Schema) -> bool {
864        if let Node::Timestamp {
865            zone: zone @ None, ..
866        } = literal
867            && let Some(DataType::Datetime(_, Some(tz))) = other.dtype(schema)
868        {
869            *zone = Some(tz.to_string());
870            return true;
871        }
872        false
873    }
874
875    /// Whether the node is boolean by its form, whatever the data: a test, a logical
876    /// operation or a boolean literal.
877    fn is_boolean(&self) -> bool {
878        match self {
879            Node::Bool(_) => true,
880            Node::Bin(op, ..) => matches!(
881                op,
882                BinOp::Eq
883                    | BinOp::Neq
884                    | BinOp::Lt
885                    | BinOp::Gt
886                    | BinOp::LtEq
887                    | BinOp::GtEq
888                    | BinOp::And
889                    | BinOp::Or
890            ),
891            Node::Op(_, op) => matches!(
892                op,
893                Op::Not
894                    | Op::IsNull
895                    | Op::IsNotNull
896                    | Op::StartsWith(_)
897                    | Op::EndsWith(_)
898                    | Op::ContainsLiteral(_)
899                    | Op::ContainsRegex(_)
900            ),
901            Node::Alias(inner, _) => inner.is_boolean(),
902            _ => false,
903        }
904    }
905
906    /// Settle what the parser could not tell by form, from the types over `schema`,
907    /// as q does. Inner operations settle first, so an outer one sees their types.
908    ///
909    /// - `&` and `|` ([`BinOp::Lesser`], [`BinOp::Greater`]) on two booleans are
910    ///   [`BinOp::And`] and [`BinOp::Or`]; a boolean with a number counts as 0/1 in the
911    ///   number's type; a string or a date with anything else is an error.
912    /// - A whole number typed next to an integer (`a & 5`, `a - 5`) is an integer, so
913    ///   the result stays one as in q.
914    /// - `not` of a number is `x = 0`, not Polars' bitwise not.
915    pub(crate) fn resolve_types(&mut self, schema: &Schema) -> Result<(), String> {
916        match self {
917            Node::Bin(_, left, right) | Node::Coalesce(left, right) | Node::Filter(left, right) => {
918                left.resolve_types(schema)?;
919                right.resolve_types(schema)?;
920            }
921            Node::When(c, t, o) => {
922                c.resolve_types(schema)?;
923                t.resolve_types(schema)?;
924                o.resolve_types(schema)?;
925            }
926            Node::Op(inner, _) | Node::Alias(inner, _) => inner.resolve_types(schema)?,
927            _ => {}
928        }
929        match self {
930            Node::Op(inner, Op::Not) if !inner.is_boolean() => {
931                let zero = match inner.dtype(schema) {
932                    Some(t) if t.is_integer() => Node::Int(0),
933                    Some(t) if t.is_float() => Node::Num(0.0),
934                    _ => return Ok(()),
935                };
936                let inner = std::mem::replace(inner.as_mut(), Node::Null);
937                *self = inner.bin(BinOp::Eq, zero);
938            }
939            Node::Bin(BinOp::Add | BinOp::Sub | BinOp::Mul, left, right) => {
940                whole_next_to_integer(left, right, schema);
941            }
942            Node::Bin(op @ (BinOp::Lesser | BinOp::Greater), left, right) => {
943                resolve_min_max(op, left, right, schema)?;
944            }
945            _ => {}
946        }
947        Ok(())
948    }
949
950    /// The type the expression has over `schema`, when Polars can say.
951    fn dtype(&self, schema: &Schema) -> Option<DataType> {
952        DataFrame::empty_with_schema(schema)
953            .lazy()
954            .select([self.to_expr()])
955            .collect_schema()
956            .ok()
957            .and_then(|s| s.get_at_index(0).map(|(_, dtype)| dtype.clone()))
958    }
959
960    /// The literal as Python, bare: `1.0`, `"a"`, `True`, `None`.
961    fn python_literal(&self) -> Option<String> {
962        Some(match self {
963            Node::Num(n) => crate::export::python_script::py_float(*n),
964            Node::Int(n) => n.to_string(),
965            Node::Str(s) => crate::export::python_script::py_str(s),
966            Node::Bool(b) => crate::export::python_script::py_bool(*b).to_string(),
967            Node::Null => "None".to_string(),
968            _ => return None,
969        })
970    }
971
972    /// Python Polars code for the expression.
973    pub(crate) fn python(&self) -> String {
974        use crate::export::python_script::py_str;
975        if let Some(literal) = self.python_literal() {
976            return format!("pl.lit({literal})");
977        }
978        match self {
979            Node::Col(name) => format!("pl.col({})", py_str(name)),
980            Node::Date(iso) => {
981                let parts: Vec<String> = iso
982                    .split('-')
983                    .map(|p| p.trim_start_matches('0').to_string())
984                    .map(|p| if p.is_empty() { "0".to_string() } else { p })
985                    .collect();
986                format!("pl.date({})", parts.join(", "))
987            }
988            Node::Timestamp {
989                iso,
990                format,
991                unit,
992                zone,
993            } => format!(
994                "pl.lit({}).str.to_datetime({}, time_unit={}{})",
995                py_str(iso),
996                py_str(format),
997                py_str(time_unit_name(*unit)),
998                zone.as_ref().map_or(String::new(), |zone| format!(
999                    ", time_zone={}, ambiguous=\"earliest\"",
1000                    py_str(zone)
1001                ))
1002            ),
1003            Node::Bin(BinOp::Or | BinOp::Greater, left, right) => {
1004                format!("pl.max_horizontal({}, {})", left.python(), right.python())
1005            }
1006            Node::Bin(BinOp::And, left, right) => format!(
1007                "({}.cast(pl.UInt8) * {}.cast(pl.UInt8)).cast(pl.Boolean)",
1008                left.python_operand(),
1009                right.python_operand()
1010            ),
1011            Node::Bin(BinOp::Lesser, left, right) => {
1012                let (l, r) = (left.python_operand(), right.python_operand());
1013                format!(
1014                    "pl.when({l}.is_null() | {r}.is_null()).then(None).otherwise(pl.min_horizontal({}, {}))",
1015                    left.python(),
1016                    right.python()
1017                )
1018            }
1019            Node::Bin(op, left, right) => {
1020                // A literal on the right stays bare (`pl.col("a") > 1.0`); Python's
1021                // operators turn it into one.
1022                let right = match right.python_literal() {
1023                    Some(literal) => literal,
1024                    None => right.python_operand(),
1025                };
1026                format!("{} {} {}", left.python_operand(), op.python(), right)
1027            }
1028            Node::Coalesce(left, right) => {
1029                format!("pl.coalesce({}, {})", left.python(), right.python())
1030            }
1031            Node::Filter(values, predicate) => {
1032                format!("{}.filter({})", values.python_operand(), predicate.python())
1033            }
1034            Node::When(condition, then, otherwise) => format!(
1035                "pl.when({}).then({}).otherwise({})",
1036                condition.python(),
1037                then.python(),
1038                otherwise.python()
1039            ),
1040            Node::Op(inner, op) => format!("{}{}", inner.python_operand(), op.python()),
1041            Node::Alias(inner, name) => {
1042                // Only the outer name survives an alias of an alias.
1043                let mut inner = inner.as_ref();
1044                while let Node::Alias(deeper, _) = inner {
1045                    inner = deeper;
1046                }
1047                format!("{}.alias({})", inner.python_operand(), py_str(name))
1048            }
1049            _ => unreachable!("literals return above"),
1050        }
1051    }
1052
1053    /// As [`Self::python`], parenthesized where an operator or a method after it
1054    /// would otherwise bind to part of it.
1055    fn python_operand(&self) -> String {
1056        match self {
1057            // Calls, which bind as a whole.
1058            Node::Bin(BinOp::And | BinOp::Or | BinOp::Lesser | BinOp::Greater, ..) => self.python(),
1059            Node::Bin(..) => format!("({})", self.python()),
1060            _ => self.python(),
1061        }
1062    }
1063}
1064
1065/// What `&` and `|` make of an operand's type.
1066#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1067enum Kind {
1068    Bool,
1069    Int,
1070    Float,
1071    Temporal,
1072    Text,
1073    /// Unknown, or a type q has no rule here for (a null literal, a list).
1074    Other,
1075}
1076
1077impl Kind {
1078    fn of(node: &Node, schema: &Schema) -> Kind {
1079        if node.is_boolean() {
1080            return Kind::Bool;
1081        }
1082        match node.dtype(schema) {
1083            Some(DataType::Boolean) => Kind::Bool,
1084            Some(t) if t.is_integer() => Kind::Int,
1085            Some(t) if t.is_float() => Kind::Float,
1086            Some(t) if t.is_temporal() => Kind::Temporal,
1087            Some(DataType::String | DataType::Categorical(..) | DataType::Enum(..)) => Kind::Text,
1088            _ => Kind::Other,
1089        }
1090    }
1091
1092    fn name(self) -> &'static str {
1093        match self {
1094            Kind::Bool => "a boolean",
1095            Kind::Int | Kind::Float => "a number",
1096            Kind::Temporal => "a date or time",
1097            Kind::Text => "a string",
1098            Kind::Other => "a value",
1099        }
1100    }
1101}
1102
1103/// A whole number typed on one side of an integer becomes an integer literal.
1104fn whole_next_to_integer(left: &mut Node, right: &mut Node, schema: &Schema) {
1105    let whole = |node: &Node| match node {
1106        Node::Num(n) if n.fract() == 0.0 && n.abs() < i64::MAX as f64 => Some(*n as i64),
1107        _ => None,
1108    };
1109    let integer = |node: &Node| matches!(Kind::of(node, schema), Kind::Int | Kind::Bool);
1110    if let Some(n) = whole(left)
1111        && integer(right)
1112    {
1113        *left = Node::Int(n);
1114    } else if let Some(n) = whole(right)
1115        && integer(left)
1116    {
1117        *right = Node::Int(n);
1118    }
1119}
1120
1121/// Settle `op`, q's `&` or `|` on anything but two booleans known by form, by the
1122/// operands' types; see [`Node::resolve_types`].
1123fn resolve_min_max(
1124    op: &mut BinOp,
1125    left: &mut Node,
1126    right: &mut Node,
1127    schema: &Schema,
1128) -> Result<(), String> {
1129    let symbol = if *op == BinOp::Lesser { "&" } else { "|" };
1130    whole_next_to_integer(left, right, schema);
1131    let (l, r) = (Kind::of(left, schema), Kind::of(right, schema));
1132    let number = |k: Kind| matches!(k, Kind::Int | Kind::Float);
1133    let refuse = |a: Kind, b: Kind| {
1134        Err(format!(
1135            "`{symbol}` cannot combine {} with {}",
1136            a.name(),
1137            b.name()
1138        ))
1139    };
1140    match (l, r) {
1141        (Kind::Bool, Kind::Bool) => {
1142            *op = if *op == BinOp::Lesser {
1143                BinOp::And
1144            } else {
1145                BinOp::Or
1146            };
1147        }
1148        (Kind::Text, Kind::Bool | Kind::Int | Kind::Float | Kind::Temporal)
1149        | (Kind::Bool | Kind::Int | Kind::Float | Kind::Temporal, Kind::Text) => {
1150            return refuse(l, r);
1151        }
1152        // q's result would be the date or time, counted from 2000.01.01; Polars
1153        // counts from 1970, so there is no faithful cast.
1154        (Kind::Temporal, Kind::Bool | Kind::Int | Kind::Float)
1155        | (Kind::Bool | Kind::Int | Kind::Float, Kind::Temporal) => return refuse(l, r),
1156        // A boolean is 0/1 in the number's type.
1157        (Kind::Bool, k) if number(k) => bool_as(left, right, schema),
1158        (k, Kind::Bool) if number(k) => bool_as(right, left, schema),
1159        // Polars' common type of Int64 and UInt64 is Float64; q keeps a long. A
1160        // UInt64 past Int64's range becomes null.
1161        (Kind::Int, Kind::Int) => {
1162            let unsigned = |n: &Node| n.dtype(schema) == Some(DataType::UInt64);
1163            let signed = |n: &Node| n.dtype(schema).is_some_and(|t| t.is_signed_integer());
1164            if unsigned(left) && signed(right) {
1165                cast_in_place(left, CastTo::Int64);
1166            } else if unsigned(right) && signed(left) {
1167                cast_in_place(right, CastTo::Int64);
1168            }
1169        }
1170        _ => {}
1171    }
1172    Ok(())
1173}
1174
1175/// Cast the boolean `flag` to the type of the number beside it: a whole number as
1176/// typed is a long, as in q (Polars would read it as Int32).
1177fn bool_as(flag: &mut Node, number: &Node, schema: &Schema) {
1178    let to = match number {
1179        Node::Int(_) => Some(CastTo::Int64),
1180        _ => number.dtype(schema).as_ref().and_then(CastTo::number),
1181    };
1182    if let Some(to) = to {
1183        cast_in_place(flag, to);
1184    }
1185}
1186
1187/// Null where either side is null.
1188fn null_if_either(left: &Expr, right: &Expr) -> Then {
1189    when(left.clone().is_null().or(right.clone().is_null())).then(lit(NULL))
1190}
1191
1192/// The smaller or the larger of two, row by row; a null is skipped.
1193fn min_max(smaller: bool, left: Expr, right: Expr) -> Expr {
1194    let both = [left, right];
1195    let result = if smaller {
1196        polars::lazy::dsl::min_horizontal(both)
1197    } else {
1198        polars::lazy::dsl::max_horizontal(both)
1199    };
1200    // Two inputs, so never the empty-input error.
1201    result.expect("min/max of two expressions")
1202}
1203
1204/// `node` replaced by itself cast to `to`.
1205fn cast_in_place(node: &mut Node, to: CastTo) {
1206    let inner = std::mem::replace(node, Node::Null);
1207    *node = inner.op(Op::Cast(to));
1208}
1209
1210fn time_unit_name(unit: TimeUnit) -> &'static str {
1211    match unit {
1212        TimeUnit::Milliseconds => "ms",
1213        TimeUnit::Microseconds => "us",
1214        TimeUnit::Nanoseconds => "ns",
1215    }
1216}
1217
1218impl BinOp {
1219    fn python(self) -> &'static str {
1220        match self {
1221            BinOp::Add => "+",
1222            BinOp::Sub => "-",
1223            BinOp::Mul => "*",
1224            BinOp::TrueDiv => "/",
1225            BinOp::FloorDiv => "//",
1226            BinOp::Rem => "%",
1227            BinOp::Eq => "==",
1228            BinOp::Neq => "!=",
1229            BinOp::Lt => "<",
1230            BinOp::Gt => ">",
1231            BinOp::LtEq => "<=",
1232            BinOp::GtEq => ">=",
1233            BinOp::Any => "|",
1234            // Written as calls; see `Node::python`.
1235            BinOp::And | BinOp::Lesser => "&",
1236            BinOp::Or | BinOp::Greater => "|",
1237        }
1238    }
1239}
1240
1241impl Op {
1242    /// The method call, from its dot: `.str.to_uppercase()`.
1243    fn python(&self) -> String {
1244        use crate::export::python_script::py_str;
1245        let fixed = match self {
1246            Op::Mean => ".mean()",
1247            Op::Min => ".min()",
1248            Op::Max => ".max()",
1249            Op::Count => ".count()",
1250            Op::Std => ".std()",
1251            Op::Var => ".var()",
1252            Op::Median => ".median()",
1253            Op::Sum => ".sum()",
1254            Op::First => ".first()",
1255            Op::Last => ".last()",
1256            Op::NUnique => ".n_unique()",
1257            Op::Not => ".not_()",
1258            Op::IsNull => ".is_null()",
1259            Op::IsNotNull => ".is_not_null()",
1260            Op::LenChars => ".str.len_chars()",
1261            Op::Upper => ".str.to_uppercase()",
1262            Op::Lower => ".str.to_lowercase()",
1263            Op::Abs => ".abs()",
1264            Op::Floor => ".floor()",
1265            Op::Ceil => ".ceil()",
1266            Op::Sqrt => ".sqrt()",
1267            Op::Ln => ".log()",
1268            Op::Exp => ".exp()",
1269            Op::Date => ".dt.date()",
1270            Op::Time => ".dt.time()",
1271            Op::Year => ".dt.year()",
1272            Op::Quarter => ".dt.quarter()",
1273            Op::Month => ".dt.month()",
1274            Op::Week => ".dt.week()",
1275            Op::Day => ".dt.day()",
1276            Op::OrdinalDay => ".dt.ordinal_day()",
1277            Op::Weekday => ".dt.weekday()",
1278            Op::Hour => ".dt.hour()",
1279            Op::Minute => ".dt.minute()",
1280            Op::Second => ".dt.second()",
1281            Op::MonthStart => ".dt.month_start()",
1282            Op::MonthEnd => ".dt.month_end()",
1283            Op::Strip => ".str.strip_chars()",
1284            Op::Cast(CastTo::String) => ".cast(pl.String)",
1285            _ => "",
1286        };
1287        if !fixed.is_empty() {
1288            return fixed.to_string();
1289        }
1290        let format_arg = |format: &Option<String>| match format {
1291            Some(f) => format!("{}, strict=False", py_str(f)),
1292            None => "strict=False".to_string(),
1293        };
1294        match self {
1295            Op::Cast(to) => format!(".cast({}, strict=False)", to.python()),
1296            Op::DtFormat(f) => format!(".dt.to_string({})", py_str(f)),
1297            Op::StartsWith(s) => format!(".str.starts_with({})", py_str(s)),
1298            Op::EndsWith(s) => format!(".str.ends_with({})", py_str(s)),
1299            Op::ContainsLiteral(s) => format!(".str.contains({}, literal=True)", py_str(s)),
1300            Op::ContainsRegex(r) => format!(".str.contains({})", py_str(r)),
1301            Op::Part(sep, i) => format!(
1302                ".str.split({}).list.get({i}, null_on_oob=True)",
1303                py_str(sep)
1304            ),
1305            Op::Slice(start, Some(len)) => format!(".str.slice({start}, {len})"),
1306            Op::Slice(start, None) => format!(".str.slice({start})"),
1307            Op::ReplaceAll(from, to) => format!(
1308                ".str.replace_all({}, {}, literal=True)",
1309                py_str(from),
1310                py_str(to)
1311            ),
1312            Op::ToDate(format) => format!(".str.to_date({})", format_arg(format)),
1313            Op::ToDatetime(format) => format!(".str.to_datetime({})", format_arg(format)),
1314            Op::Round(d) => format!(".round({d}, mode=\"half_away_from_zero\")"),
1315            _ => unreachable!("fixed calls return above"),
1316        }
1317    }
1318}
1319
1320fn apply_op_expr(expr: Expr, op: &Op) -> Expr {
1321    let strptime = |format: &Option<String>| StrptimeOptions {
1322        format: format.as_deref().map(Into::into),
1323        // A value that does not match becomes null, as a failed parse does in q,
1324        // rather than one stray row failing the whole query.
1325        strict: false,
1326        ..Default::default()
1327    };
1328    match op {
1329        Op::Mean => expr.mean(),
1330        Op::Min => expr.min(),
1331        Op::Max => expr.max(),
1332        Op::Count => expr.count(),
1333        // Sample statistics (n - 1), like `std`; q's own var and dev divide by n.
1334        Op::Std => expr.std(1),
1335        Op::Var => expr.var(1),
1336        Op::Median => expr.median(),
1337        Op::Sum => expr.sum(),
1338        Op::First => expr.first(),
1339        Op::Last => expr.last(),
1340        Op::NUnique => expr.n_unique(),
1341        Op::Not => expr.not(),
1342        Op::IsNull => expr.is_null(),
1343        Op::IsNotNull => expr.is_not_null(),
1344        Op::LenChars => expr.str().len_chars(),
1345        Op::Upper => expr.str().to_uppercase(),
1346        Op::Lower => expr.str().to_lowercase(),
1347        Op::Abs => expr.abs(),
1348        Op::Floor => expr.floor(),
1349        Op::Ceil => expr.ceil(),
1350        Op::Sqrt => expr.sqrt(),
1351        Op::Ln => expr.log(lit(std::f64::consts::E)),
1352        Op::Exp => expr.exp(),
1353        Op::Date => expr.dt().date(),
1354        Op::Time => expr.dt().time(),
1355        Op::Year => expr.dt().year(),
1356        Op::Quarter => expr.dt().quarter(),
1357        Op::Month => expr.dt().month(),
1358        Op::Week => expr.dt().week(),
1359        Op::Day => expr.dt().day(),
1360        Op::OrdinalDay => expr.dt().ordinal_day(),
1361        Op::Weekday => expr.dt().weekday(),
1362        Op::Hour => expr.dt().hour(),
1363        Op::Minute => expr.dt().minute(),
1364        Op::Second => expr.dt().second(),
1365        Op::MonthStart => expr.dt().month_start(),
1366        Op::MonthEnd => expr.dt().month_end(),
1367        Op::DtFormat(f) => expr.dt().to_string(f),
1368        Op::StartsWith(s) => expr.str().starts_with(lit(s.as_str())),
1369        Op::EndsWith(s) => expr.str().ends_with(lit(s.as_str())),
1370        Op::ContainsLiteral(s) => expr.str().contains_literal(lit(s.as_str())),
1371        Op::ContainsRegex(r) => expr.str().contains(lit(r.as_str()), true),
1372        // Past the last piece is null, not an error; a negative index counts from the end.
1373        Op::Part(sep, i) => expr
1374            .str()
1375            .split(lit(sep.as_str()))
1376            .list()
1377            .get(lit(*i), true),
1378        Op::Slice(start, length) => {
1379            // No length: to the end of the string.
1380            let length = length.map_or_else(|| lit(NULL), lit);
1381            expr.str().slice(lit(*start), length)
1382        }
1383        Op::ReplaceAll(from, to) => {
1384            expr.str()
1385                .replace_all(lit(from.as_str()), lit(to.as_str()), true)
1386        }
1387        Op::Strip => expr.str().strip_chars(lit(NULL)),
1388        Op::ToDate(format) => expr.str().to_date(strptime(format)),
1389        Op::ToDatetime(format) => {
1390            expr.str()
1391                .to_datetime(None, None, strptime(format), lit("raise"))
1392        }
1393        // Half away from zero, the rounding people expect from a calculator or SQL.
1394        Op::Round(decimals) => expr.round(*decimals, RoundMode::HalfAwayFromZero),
1395        // Non-strict casts: a value that does not convert becomes null.
1396        Op::Cast(to) => expr.cast(to.dtype()),
1397    }
1398}
1399
1400/// An operand of `mod` or `xbar`: whole numbers as integer literals, so integer columns
1401/// keep their type (`5 xbar passenger_count` stays Int64).
1402fn int_or_node(tokens: &[Token]) -> Result<Node, String> {
1403    let whole = |n: f64| n.fract() == 0.0 && n.abs() < i64::MAX as f64;
1404    match tokens {
1405        [Token::Number(n)] if whole(*n) => Ok(Node::Int(*n as i64)),
1406        [Token::Op(minus), Token::Number(n)] if minus == "-" && whole(*n) => {
1407            Ok(Node::Int(-(*n as i64)))
1408        }
1409        _ => parse_node(tokens),
1410    }
1411}
1412
1413/// True when no `]` in `tokens` closes a `[` from outside them.
1414fn brackets_balanced(tokens: &[Token]) -> bool {
1415    let mut depth = 0usize;
1416    tokens.iter().all(|t| match t {
1417        Token::LBracket => {
1418            depth += 1;
1419            true
1420        }
1421        Token::RBracket => depth.checked_sub(1).map(|d| depth = d).is_some(),
1422        _ => true,
1423    })
1424}
1425
1426/// The error for `column` of a temporal type compared with the quoted `text`, naming
1427/// the literal the type takes: the text itself when it is one once unquoted.
1428fn quoted_temporal(column: &Node, text: &Node, schema: &Schema) -> Option<String> {
1429    let (Node::Col(name), Node::Str(s)) = (column, text) else {
1430        return None;
1431    };
1432    let shown = q_name(name);
1433    let unquoted = tokenize(s).ok();
1434    let literal = |is_kind: fn(&Token) -> bool, example: &str| match unquoted.as_deref() {
1435        Some([token]) if is_kind(token) => s.trim().to_string(),
1436        _ => example.to_string(),
1437    };
1438    let (kind, remedy) = match schema.get(name)? {
1439        DataType::Date => (
1440            "date",
1441            format!(
1442                "A date is {}",
1443                literal(|t| matches!(t, Token::DateLiteral(_)), "2024.01.01")
1444            ),
1445        ),
1446        DataType::Datetime(..) => (
1447            "timestamp",
1448            format!(
1449                "A timestamp is {}",
1450                literal(
1451                    |t| matches!(t, Token::TimestampLiteral { .. }),
1452                    "2024.01.01T05:00:00"
1453                )
1454            ),
1455        ),
1456        DataType::Time => (
1457            "time",
1458            format!(
1459                "A time has no literal; compare {shown}.hour, {shown}.minute or {shown}.second with a number"
1460            ),
1461        ),
1462        DataType::Duration(_) => ("duration", "A duration has no literal".to_string()),
1463        _ => return None,
1464    };
1465    Some(format!(
1466        "{shown} is a {kind}; \"{s}\" is a string. {remedy}"
1467    ))
1468}
1469
1470/// How q spells a column: bare when it can be, else `col["first name"]`.
1471pub(crate) fn q_name(name: &str) -> String {
1472    if is_plain_name(name) {
1473        name.to_string()
1474    } else {
1475        format!("col[\"{name}\"]")
1476    }
1477}
1478
1479/// Whether `name` reads as a column when typed bare, rather than needing `col["…"]`.
1480fn is_plain_name(name: &str) -> bool {
1481    let mut chars = name.chars();
1482    chars.next().is_some_and(|c| c.is_alphabetic() || c == '_')
1483        && chars.all(|c| c.is_alphanumeric() || c == '_')
1484        && !matches!(name, "select" | "where" | "by")
1485}
1486
1487/// OR of the conditions as a balanced tree, so a long `in` list nests
1488/// logarithmically rather than one level per element.
1489fn any_of(mut conditions: Vec<Node>) -> Node {
1490    if conditions.len() <= 1 {
1491        return conditions.pop().unwrap_or(Node::Bool(false));
1492    }
1493    let right = conditions.split_off(conditions.len() / 2);
1494    any_of(conditions).bin(BinOp::Any, any_of(right))
1495}
1496
1497/// A `like` pattern as an anchored regex: `*` is any run of characters, `?` any
1498/// one character, everything else literal.
1499fn like_regex(pattern: &str) -> String {
1500    let mut re = String::from("(?s)^");
1501    for c in pattern.chars() {
1502        match c {
1503            '*' => re.push_str(".*"),
1504            '?' => re.push('.'),
1505            _ => re.push_str(&regex::escape(c.encode_utf8(&mut [0; 4]))),
1506        }
1507    }
1508    re.push('$');
1509    re
1510}
1511
1512/// Operators whose right side is not an ordinary expression, or whose operands
1513/// need their tokens (literal checks, names). Everything else goes to `apply_op`.
1514fn apply_infix(left_tokens: &[Token], op: &str, right_tokens: &[Token]) -> Result<Node, String> {
1515    match op {
1516        "in" => {
1517            let list = match right_tokens {
1518                // One list: the first `[` closes at the last token, so `[1] + [2]` is not.
1519                [Token::LBracket, inner @ .., Token::RBracket] if brackets_balanced(inner) => inner,
1520                _ => {
1521                    return Err(
1522                        "in takes a list on its right, e.g. name in [\"Emma\", \"Olivia\"]"
1523                            .to_string(),
1524                    );
1525                }
1526            };
1527            let items = split_tokens(list, &Token::Comma);
1528            if items.iter().any(|item| item.is_empty()) {
1529                return Err(
1530                    "in needs a list of values, e.g. name in [\"Emma\", \"Olivia\"]".to_string(),
1531                );
1532            }
1533            let left = parse_node(left_tokens)?;
1534            // A column or literal repeated grows only as the list typed does; a larger
1535            // left side repeated per item is what multiplies.
1536            if left.size() > 1 {
1537                check_copies(&left, items.len())?;
1538            }
1539            // One `=` per value, so each value compares exactly as `x = value` would,
1540            // with the same literal casting (numbers, dates, timestamps).
1541            let conditions = items
1542                .iter()
1543                .map(|item| Ok(left.clone().bin(BinOp::Eq, parse_node(item)?)))
1544                .collect::<Result<Vec<_>, String>>()?;
1545            Ok(any_of(conditions))
1546        }
1547        "like" => {
1548            let [Token::String(pattern)] = right_tokens else {
1549                return Err(
1550                    "like takes a quoted pattern on its right, e.g. item like \"*Chicken*\""
1551                        .to_string(),
1552                );
1553            };
1554            let left = parse_node(left_tokens)?;
1555            // Cast first so numeric codes (zip, station ids read as numbers) match too.
1556            Ok(left.cast_text().op(Op::ContainsRegex(like_regex(pattern))))
1557        }
1558        "xbar" => {
1559            if let [Token::Number(n)] = left_tokens
1560                && *n <= 0.0
1561            {
1562                return Err(
1563                    "xbar needs a positive bucket size, e.g. 5 xbar fare_amount".to_string()
1564                );
1565            }
1566            let right = parse_node(right_tokens)?;
1567            let size = int_or_node(left_tokens)?;
1568            check_copies(&size, 2)?;
1569            // floor_div floors toward negative infinity for both ints and floats,
1570            // which is what makes every value land in the bucket at or below it.
1571            Ok(right
1572                .bin(BinOp::FloorDiv, size.clone())
1573                .bin(BinOp::Mul, size))
1574        }
1575        "mod" => {
1576            let right = int_or_node(right_tokens)?;
1577            let left = int_or_node(left_tokens)?;
1578            Ok(left.bin(BinOp::Rem, right))
1579        }
1580        "wavg" => {
1581            let values = parse_node(right_tokens)?;
1582            let weights = parse_node(left_tokens)?;
1583            // Below, weights appear five times and values three.
1584            check_copies(&weights, 5)?;
1585            check_copies(&values, 3)?;
1586            let weighted = weights.clone().bin(BinOp::Mul, values);
1587            // Only pairs with both a weight and a value count toward the total weight;
1588            // a null value would otherwise still pull the average toward zero.
1589            let total = Node::Filter(
1590                Box::new(weights),
1591                Box::new(weighted.clone().op(Op::IsNotNull)),
1592            )
1593            .op(Op::Sum);
1594            // No weight at all (every pair null) is no average, not 0/0 = NaN.
1595            let total = Node::When(
1596                Box::new(total.clone().bin(BinOp::Neq, Node::Int(0))),
1597                Box::new(total),
1598                Box::new(Node::Null),
1599            );
1600            let node = weighted.op(Op::Sum).bin(BinOp::TrueDiv, total);
1601            Ok(match simple_column_name(right_tokens) {
1602                Some(column) => node.alias(format!("wavg_{}", column)),
1603                None => node,
1604            })
1605        }
1606        _ => {
1607            // Parse right side first (right-to-left evaluation): it holds any
1608            // remaining operators, so c>c%n becomes c > (c%n).
1609            let right = parse_node(right_tokens)?;
1610            let left = parse_node(left_tokens)?;
1611            apply_op(left, op, right)
1612        }
1613    }
1614}
1615
1616fn apply_op(left: Node, op: &str, right: Node) -> Result<Node, String> {
1617    let op = match op {
1618        "+" => BinOp::Add,
1619        "-" => BinOp::Sub,
1620        "*" => BinOp::Mul,
1621        // `%` divides (q heritage); `/` is the alias everyone expects. Always to a
1622        // float, as q's `%`, even on two integers (Polars' `/` would floor them).
1623        "%" | "/" => BinOp::TrueDiv,
1624        "^" => return Ok(Node::Coalesce(Box::new(left), Box::new(right))),
1625        "=" => BinOp::Eq,
1626        "<" => BinOp::Lt,
1627        ">" => BinOp::Gt,
1628        "<=" => BinOp::LtEq,
1629        ">=" => BinOp::GtEq,
1630        "<>" | "!=" => BinOp::Neq,
1631        // q's Lesser and Greater: and/or on two booleans, min/max otherwise. Two
1632        // tests or other boolean forms are known here; columns wait for the schema
1633        // (`Node::resolve_types`). `&` repeats its operands (null where either is).
1634        "&" | "and" if left.is_boolean() && right.is_boolean() => BinOp::And,
1635        "&" | "and" => {
1636            check_copies(&left, 2)?;
1637            check_copies(&right, 2)?;
1638            BinOp::Lesser
1639        }
1640        "|" | "or" if left.is_boolean() && right.is_boolean() => BinOp::Or,
1641        "|" | "or" => BinOp::Greater,
1642        _ => return Err(format!("Unknown operator: {}", op)),
1643    };
1644    Ok(left.bin(op, right))
1645}
1646
1647/// The column name when the tokens are a bare column reference: `salary`, or
1648/// `col["first name"]` / `col[name]`. Anything more (literals, operators) is None.
1649fn simple_column_name(tokens: &[Token]) -> Option<String> {
1650    match tokens {
1651        [Token::Identifier(name)] => Some(name.clone()),
1652        [
1653            Token::Identifier(c),
1654            Token::LBracket,
1655            Token::String(name) | Token::Identifier(name),
1656            Token::RBracket,
1657        ] if c == "col" => Some(name.clone()),
1658        _ => None,
1659    }
1660}
1661
1662const WAVG_USAGE: &str = "wavg goes between weights and values, e.g. passengers wavg fare";
1663
1664/// Aggregation function names, lowercase.
1665const AGG_FUNCTIONS: [&str; 16] = [
1666    "avg", "mean", "min", "max", "count", "std", "stddev", "dev", "var", "med", "median", "sum",
1667    "first", "last", "nunique", "wavg",
1668];
1669
1670/// Scalar function names, lowercase.
1671const SCALAR_FUNCTIONS: [&str; 13] = [
1672    "len", "length", "not", "null", "upper", "lower", "abs", "floor", "ceil", "ceiling", "sqrt",
1673    "log", "exp",
1674];
1675
1676fn is_agg_function(name: &str) -> bool {
1677    AGG_FUNCTIONS.contains(&name.to_lowercase().as_str())
1678}
1679
1680fn is_function_name(name: &str) -> bool {
1681    // `wavg` is infix (`w wavg x`), so it never opens an expression.
1682    let name = name.to_lowercase();
1683    name != "wavg" && (is_agg_function(&name) || SCALAR_FUNCTIONS.contains(&name.as_str()))
1684}
1685
1686/// A call, `fn[args]` or `fn args`. The name is checked first so arguments parse once
1687/// (trying aggregates then scalars doubled work per nesting level).
1688fn parse_call(name: &str, args: &[Token]) -> Result<Node, String> {
1689    if is_agg_function(name) {
1690        parse_agg_function(name, args)
1691    } else {
1692        parse_function(name, args)
1693    }
1694}
1695
1696/// An aggregate call: `avg[a]`, `min[b]`.
1697fn parse_agg_function(name: &str, args: &[Token]) -> Result<Node, String> {
1698    if args.is_empty() {
1699        return Err(format!(
1700            "Aggregation function {} requires an argument",
1701            name
1702        ));
1703    }
1704    let fn_name = name.to_lowercase();
1705    if fn_name == "wavg" {
1706        return Err(WAVG_USAGE.to_string());
1707    }
1708    let node = parse_node(args)?;
1709    let op = match fn_name.as_str() {
1710        "avg" | "mean" => Op::Mean,
1711        "min" => Op::Min,
1712        "max" => Op::Max,
1713        "count" => Op::Count,
1714        "std" | "stddev" | "dev" => Op::Std,
1715        "var" => Op::Var,
1716        "med" | "median" => Op::Median,
1717        "sum" => Op::Sum,
1718        "first" => Op::First,
1719        "last" => Op::Last,
1720        "nunique" => Op::NUnique,
1721        _ => return Err(format!("Unknown aggregation function: {}", name)),
1722    };
1723    let node = node.op(op);
1724    // A bare-column aggregate is named `{fn}_{column}` so two aggregates of one column do
1725    // not collide; an explicit alias overrides it later.
1726    match simple_column_name(args) {
1727        Some(column) => Ok(node.alias(format!("{}_{}", fn_name, column))),
1728        None => Ok(node),
1729    }
1730}
1731
1732/// A scalar call: `not[a=b]`, `null[col]`, `len[x]`, `upper[x]`.
1733fn parse_function(name: &str, args: &[Token]) -> Result<Node, String> {
1734    if args.is_empty() {
1735        return Err(format!("Function {} requires an argument", name));
1736    }
1737    let name_lower = name.to_lowercase();
1738    if !SCALAR_FUNCTIONS.contains(&name_lower.as_str()) {
1739        return Err(format!("Unknown function: {}", name));
1740    }
1741    let node = parse_node(args)?;
1742    let op = match name_lower.as_str() {
1743        "not" => Op::Not,
1744        "null" => Op::IsNull,
1745        "len" | "length" => Op::LenChars,
1746        "upper" => Op::Upper,
1747        "lower" => Op::Lower,
1748        "abs" => Op::Abs,
1749        "floor" => Op::Floor,
1750        "ceil" | "ceiling" => Op::Ceil,
1751        "sqrt" => Op::Sqrt,
1752        "log" => Op::Ln,
1753        "exp" => Op::Exp,
1754        _ => return Err(format!("Unknown function: {}", name)),
1755    };
1756    Ok(node.op(op))
1757}
1758
1759/// An accessor's bracketed argument: `.part["-", 0]` has a string and a number.
1760#[derive(Debug, Clone, PartialEq)]
1761enum AccessorArg {
1762    Str(String),
1763    Num(f64),
1764}
1765
1766impl AccessorArg {
1767    /// As it goes into the result's auto-alias: `part_-_0`, not `part_-_0.0`.
1768    fn alias_text(&self) -> String {
1769        match self {
1770            AccessorArg::Str(s) => s.clone(),
1771            AccessorArg::Num(n) => n.to_string(),
1772        }
1773    }
1774}
1775
1776/// Every accessor with the argument counts it takes and an example for errors.
1777const ACCESSORS: &[(&str, usize, usize, &str)] = &[
1778    // Date and time parts.
1779    ("date", 0, 0, ".date"),
1780    ("time", 0, 0, ".time"),
1781    ("year", 0, 0, ".year"),
1782    ("quarter", 0, 0, ".quarter"),
1783    ("month", 0, 0, ".month"),
1784    ("week", 0, 0, ".week"),
1785    ("day", 0, 0, ".day"),
1786    ("doy", 0, 0, ".doy"),
1787    ("dow", 0, 0, ".dow"),
1788    ("weekday", 0, 0, ".weekday"),
1789    ("hour", 0, 0, ".hour"),
1790    ("minute", 0, 0, ".minute"),
1791    ("second", 0, 0, ".second"),
1792    ("month_start", 0, 0, ".month_start"),
1793    ("month_end", 0, 0, ".month_end"),
1794    ("format", 1, 1, ".format[\"%Y-%m\"]"),
1795    // Strings.
1796    ("len", 0, 0, ".len"),
1797    ("length", 0, 0, ".length"),
1798    ("upper", 0, 0, ".upper"),
1799    ("lower", 0, 0, ".lower"),
1800    ("starts_with", 1, 1, ".starts_with[\"x\"]"),
1801    ("ends_with", 1, 1, ".ends_with[\"x\"]"),
1802    ("contains", 1, 1, ".contains[\"x\"]"),
1803    ("part", 2, 2, ".part[\"-\", 0]"),
1804    ("slice", 1, 2, ".slice[0, 4]"),
1805    ("replace", 2, 2, ".replace[\"(P)\", \"\"]"),
1806    ("strip", 0, 0, ".strip"),
1807    ("to_date", 0, 1, ".to_date[\"%Y%m%d\"]"),
1808    ("to_datetime", 0, 1, ".to_datetime[\"%Y-%m-%d %H:%M\"]"),
1809    // Numbers and casts.
1810    ("round", 0, 1, ".round[1]"),
1811    ("int", 0, 0, ".int"),
1812    ("float", 0, 0, ".float"),
1813    ("str", 0, 0, ".str"),
1814];
1815
1816/// Names for the unknown-accessor error, by kind.
1817const ACCESSOR_HELP: &str = "Valid date/time: date, time, year, quarter, month, week, day, doy, dow, hour, minute, second, month_start, month_end, format. \
1818     Valid string: len, upper, lower, starts_with, ends_with, contains, part, slice, replace, strip, to_date, to_datetime. \
1819     Valid number: round, int, float, str";
1820
1821fn arg_count_text(min: usize, max: usize) -> String {
1822    match (min, max) {
1823        (0, 0) => "no arguments".to_string(),
1824        (1, 1) => "1 argument".to_string(),
1825        (a, b) if a == b => format!("{} arguments", a),
1826        (a, b) => format!("{} to {} arguments", a, b),
1827    }
1828}
1829
1830/// Apply accessor `name` with its bracketed arguments.
1831fn apply_accessor(node: Node, accessor: &str, args: &[AccessorArg]) -> Result<Node, String> {
1832    let name = accessor.to_lowercase();
1833    let Some(&(_, min, max, usage)) = ACCESSORS.iter().find(|(n, ..)| *n == name) else {
1834        return Err(format!(
1835            "Unknown accessor: '{}'. {}",
1836            accessor, ACCESSOR_HELP
1837        ));
1838    };
1839    if args.len() < min || args.len() > max {
1840        return Err(format!(
1841            "{} takes {}, e.g. {}; got {}",
1842            name,
1843            arg_count_text(min, max),
1844            usage,
1845            args.len()
1846        ));
1847    }
1848    let text = |i: usize| match args.get(i) {
1849        Some(AccessorArg::Str(s)) => Ok(s.clone()),
1850        _ => Err(format!(
1851            "{}: argument {} must be quoted text, e.g. {}",
1852            name,
1853            i + 1,
1854            usage
1855        )),
1856    };
1857    let int = |i: usize| match args.get(i) {
1858        Some(AccessorArg::Num(n)) if n.fract() == 0.0 && n.abs() <= u32::MAX as f64 => {
1859            Ok(*n as i64)
1860        }
1861        _ => Err(format!(
1862            "{}: argument {} must be a whole number, e.g. {}",
1863            name,
1864            i + 1,
1865            usage
1866        )),
1867    };
1868    // The string pieces cast first, so they also work on numbers and dates read
1869    // as such (NOAA's DATE, an integer zip code); a cast from String is a no-op.
1870    let as_str = || node.clone().cast_text();
1871    Ok(match name.as_str() {
1872        "date" => node.op(Op::Date),
1873        "time" => node.op(Op::Time),
1874        "year" => node.op(Op::Year),
1875        "quarter" => node.op(Op::Quarter),
1876        "month" => node.op(Op::Month),
1877        "week" => node.op(Op::Week),
1878        "day" => node.op(Op::Day),
1879        "doy" => node.op(Op::OrdinalDay),
1880        "dow" | "weekday" => node.op(Op::Weekday),
1881        "hour" => node.op(Op::Hour),
1882        "minute" => node.op(Op::Minute),
1883        "second" => node.op(Op::Second),
1884        "month_start" => node.op(Op::MonthStart),
1885        "month_end" => node.op(Op::MonthEnd),
1886        "format" => node.op(Op::DtFormat(text(0)?)),
1887        "len" | "length" => node.op(Op::LenChars),
1888        "upper" => node.op(Op::Upper),
1889        "lower" => node.op(Op::Lower),
1890        "starts_with" => node.op(Op::StartsWith(text(0)?)),
1891        "ends_with" => node.op(Op::EndsWith(text(0)?)),
1892        "contains" => node.op(Op::ContainsLiteral(text(0)?)),
1893        "part" => as_str().op(Op::Part(text(0)?, int(1)?)),
1894        "slice" => {
1895            let start = int(0)?;
1896            let length = match args.len() {
1897                2 => {
1898                    let n = int(1)?;
1899                    if n < 0 {
1900                        return Err(format!(
1901                            "slice: the length cannot be negative, e.g. {}",
1902                            usage
1903                        ));
1904                    }
1905                    Some(n as u64)
1906                }
1907                // No length: to the end of the string.
1908                _ => None,
1909            };
1910            as_str().op(Op::Slice(start, length))
1911        }
1912        "replace" => as_str().op(Op::ReplaceAll(text(0)?, text(1)?)),
1913        "strip" => as_str().op(Op::Strip),
1914        "to_date" => as_str().op(Op::ToDate(args.first().map(|_| text(0)).transpose()?)),
1915        "to_datetime" => as_str().op(Op::ToDatetime(args.first().map(|_| text(0)).transpose()?)),
1916        "round" => {
1917            let decimals = if args.is_empty() { 0 } else { int(0)? };
1918            let decimals = u32::try_from(decimals)
1919                .map_err(|_| format!("round: decimals cannot be negative, e.g. {}", usage))?;
1920            node.op(Op::Round(decimals))
1921        }
1922        "int" => node.op(Op::Cast(CastTo::Int64)),
1923        "float" => node.op(Op::Cast(CastTo::Float64)),
1924        "str" => node.op(Op::Cast(CastTo::String)),
1925        _ => {
1926            return Err(format!(
1927                "Unknown accessor: '{}'. {}",
1928                accessor, ACCESSOR_HELP
1929            ));
1930        }
1931    })
1932}
1933
1934/// The arguments inside an accessor's brackets: literals separated by commas.
1935fn parse_accessor_args(accessor: &str, tokens: &[Token]) -> Result<Vec<AccessorArg>, String> {
1936    if tokens.is_empty() {
1937        return Ok(Vec::new());
1938    }
1939    split_tokens(tokens, &Token::Comma)
1940        .iter()
1941        .map(|arg| match arg.as_slice() {
1942            [Token::String(s)] | [Token::Identifier(s)] => Ok(AccessorArg::Str(s.clone())),
1943            [Token::Number(n)] => Ok(AccessorArg::Num(*n)),
1944            [Token::Op(minus), Token::Number(n)] if minus == "-" => Ok(AccessorArg::Num(-n)),
1945            _ => Err(format!(
1946                "{} takes literal arguments, quoted text or numbers, e.g. .part[\"-\", 0]",
1947                accessor
1948            )),
1949        })
1950        .collect()
1951}
1952
1953/// Parse optional dot accessors from the remaining tokens, returning
1954/// (expr_with_accessors, remaining). With `base_name`, results are aliased
1955/// `{base}_{accessor}` (chained: `{base}_{acc1}_{acc2}`) to avoid duplicates.
1956fn parse_accessors<'a>(
1957    mut expr: Node,
1958    mut tokens: &'a [Token],
1959    base_name: Option<&str>,
1960) -> Result<(Node, &'a [Token]), String> {
1961    let mut alias_suffix = String::new();
1962    while let [Token::Dot, Token::Identifier(accessor), rest @ ..] = tokens {
1963        let (args, consumed) = if rest.first() == Some(&Token::LBracket) {
1964            let mut depth = 0;
1965            let close = rest
1966                .iter()
1967                .position(|t| {
1968                    match t {
1969                        Token::LBracket => depth += 1,
1970                        Token::RBracket => depth -= 1,
1971                        _ => {}
1972                    }
1973                    depth == 0
1974                })
1975                .ok_or_else(|| format!("Unmatched bracket after .{}", accessor))?;
1976            (parse_accessor_args(accessor, &rest[1..close])?, close + 3)
1977        } else {
1978            (Vec::new(), 2)
1979        };
1980        expr = apply_accessor(expr, accessor, &args)?;
1981        if !alias_suffix.is_empty() {
1982            alias_suffix.push('_');
1983        }
1984        alias_suffix.push_str(accessor);
1985        for arg in &args {
1986            alias_suffix.push('_');
1987            alias_suffix.push_str(&arg.alias_text());
1988        }
1989        tokens = &tokens[consumed..];
1990    }
1991    if !alias_suffix.is_empty() {
1992        let alias = match base_name {
1993            Some(name) => format!("{}_{}", name, alias_suffix),
1994            None => alias_suffix,
1995        };
1996        expr = expr.alias(alias);
1997    }
1998    Ok((expr, tokens))
1999}
2000
2001fn parse_term(tokens: &[Token]) -> Result<(Node, &[Token]), String> {
2002    if tokens.is_empty() {
2003        return Err("Unexpected end of expression".to_string());
2004    }
2005    match &tokens[0] {
2006        Token::Identifier(name) => {
2007            // Check if it's col[...] syntax for column names with spaces
2008            if name == "col" && tokens.len() > 1 && tokens[1] == Token::LBracket {
2009                let mut depth = 1;
2010                let mut i = 2;
2011                while i < tokens.len() && depth > 0 {
2012                    match tokens[i] {
2013                        Token::LBracket => depth += 1,
2014                        Token::RBracket => depth -= 1,
2015                        _ => {}
2016                    }
2017                    i += 1;
2018                }
2019                if depth > 0 {
2020                    return Err("Unmatched bracket in col[]".to_string());
2021                }
2022                let col_name_tokens = &tokens[2..i - 1];
2023                if col_name_tokens.len() != 1 {
2024                    return Err("col[] must contain a single string or identifier".to_string());
2025                }
2026                let col_name = match &col_name_tokens[0] {
2027                    Token::String(s) => s.clone(),
2028                    Token::Identifier(id) => id.clone(),
2029                    _ => return Err("col[] must contain a string or identifier".to_string()),
2030                };
2031                let expr = Node::Col(col_name.clone());
2032                let (expr, remaining) = parse_accessors(expr, &tokens[i..], Some(&col_name))?;
2033                Ok((expr, remaining))
2034            }
2035            // Check if it's a function call (using square brackets)
2036            else if tokens.len() > 1 && tokens[1] == Token::LBracket {
2037                let mut depth = 1;
2038                let mut i = 2;
2039                while i < tokens.len() && depth > 0 {
2040                    match tokens[i] {
2041                        Token::LBracket => depth += 1,
2042                        Token::RBracket => depth -= 1,
2043                        _ => {}
2044                    }
2045                    i += 1;
2046                }
2047                if depth > 0 {
2048                    return Err("Unmatched bracket in function call".to_string());
2049                }
2050                let expr = parse_call(name, &tokens[2..i - 1])?;
2051                parse_accessors(expr, &tokens[i..], None)
2052            } else {
2053                // Regular column reference
2054                // (Function calls without brackets are handled in parse_node)
2055                let expr = Node::Col(name.clone());
2056                let (expr, remaining) = parse_accessors(expr, &tokens[1..], Some(name))?;
2057                Ok((expr, remaining))
2058            }
2059        }
2060        Token::Number(n) => Ok((Node::Num(*n), &tokens[1..])), // Numbers don't support accessors
2061        Token::String(s) => Ok((Node::Str(s.clone()), &tokens[1..])), // Strings don't support accessors
2062        Token::DateLiteral(iso) => Ok((Node::Date(iso.clone()), &tokens[1..])),
2063        Token::TimestampLiteral {
2064            iso,
2065            format_str,
2066            time_unit,
2067        } => Ok((
2068            Node::Timestamp {
2069                iso: iso.clone(),
2070                format: format_str.clone(),
2071                unit: *time_unit,
2072                zone: None,
2073            },
2074            &tokens[1..],
2075        )),
2076        Token::LParen => {
2077            let mut depth = 1;
2078            let mut i = 1;
2079            while i < tokens.len() && depth > 0 {
2080                match tokens[i] {
2081                    Token::LParen => depth += 1,
2082                    Token::RParen => depth -= 1,
2083                    _ => {}
2084                }
2085                i += 1;
2086            }
2087            if depth > 0 {
2088                return Err("Unmatched parenthesis".to_string());
2089            }
2090            let inner = &tokens[1..i - 1];
2091            if split_tokens(inner, &Token::Comma).len() > 1 {
2092                return Err(PAREN_COMMA.to_string());
2093            }
2094            let inner = parse_node(inner)?;
2095            let (expr, remaining) = parse_accessors(inner, &tokens[i..], None)?;
2096            Ok((expr, remaining))
2097        }
2098        // Square brackets are only for function calls, not grouping
2099        // Parentheses are used for grouping
2100        _ => Err(format!(
2101            "Unexpected '{}' where an expression was expected",
2102            token_text(&tokens[0])
2103        )),
2104    }
2105}
2106
2107/// A list in parentheses: q's join, or several where conditions grouped.
2108const PAREN_COMMA: &str = "`,` inside parentheses joins lists in q, which datui does not \
2109     support; combine conditions with `&` or `|`";
2110
2111/// Deepest nesting the recursive-descent parser follows, so a long chain (`------x`,
2112/// `((((x))))`) errors instead of overflowing the stack (found by the `parse_query`
2113/// fuzz target). Each level costs about 10 KiB of stack in a debug build, so 64 is safe
2114/// on a 2 MiB worker thread and far beyond handwritten nesting.
2115const MAX_EXPR_DEPTH: u32 = 64;
2116
2117thread_local! {
2118    static EXPR_DEPTH: std::cell::Cell<u32> = const { std::cell::Cell::new(0) };
2119}
2120
2121/// Holds the recursion counter up while alive: every recursive path passes through
2122/// `parse_node`, which returns from many places, so the decrement is scoped.
2123struct DepthGuard;
2124
2125impl DepthGuard {
2126    /// `None` once the limit is reached, leaving the counter untouched.
2127    fn enter() -> Option<Self> {
2128        EXPR_DEPTH.with(|depth| {
2129            let next = depth.get() + 1;
2130            if next > MAX_EXPR_DEPTH {
2131                return None;
2132            }
2133            depth.set(next);
2134            Some(DepthGuard)
2135        })
2136    }
2137}
2138
2139impl Drop for DepthGuard {
2140    fn drop(&mut self) {
2141        EXPR_DEPTH.with(|depth| depth.set(depth.get().saturating_sub(1)));
2142    }
2143}
2144
2145// Parse expression with right-to-left operator precedence
2146// This means operators are evaluated from right to left: a+b*c is parsed as a+(b*c)
2147fn parse_node(tokens: &[Token]) -> Result<Node, String> {
2148    let Some(_depth_guard) = DepthGuard::enter() else {
2149        return Err(
2150            "Expression is nested too deeply. Simplify it or split it into steps.".to_string(),
2151        );
2152    };
2153
2154    if tokens.is_empty() {
2155        return Err("Empty expression".to_string());
2156    }
2157
2158    // First check if this starts with a function call (without brackets)
2159    // This needs to be checked before operator parsing to ensure correct precedence
2160    if let Token::Identifier(name) = &tokens[0]
2161        && is_function_name(name)
2162        && tokens.len() > 1
2163        && tokens[1] != Token::LBracket
2164    {
2165        // A call without brackets: the rest is the argument, built as the bracketed form is.
2166        return parse_call(name, &tokens[1..]);
2167    }
2168
2169    let mut op_pos = None;
2170    let mut depth = 0;
2171    let mut bracket_depth = 0;
2172
2173    // Scan from left to right to find the leftmost operator
2174    for (i, token) in tokens.iter().enumerate() {
2175        match token {
2176            Token::LParen => depth += 1,
2177            Token::RParen => depth -= 1,
2178            Token::LBracket => bracket_depth += 1,
2179            Token::RBracket => bracket_depth -= 1,
2180            _ if depth == 0 && bracket_depth == 0 && infix_op_at(tokens, i).is_some() => {
2181                op_pos = Some(i);
2182                break;
2183            }
2184            _ => {}
2185        }
2186    }
2187
2188    if let Some(pos) = op_pos {
2189        let left_tokens = &tokens[..pos];
2190        let right_tokens = &tokens[pos + 1..];
2191
2192        if let Some(op) = infix_op_at(tokens, pos) {
2193            // Unary minus next to a literal with an operator on the other side: -0.1+discount → (-0.1)+discount
2194            if left_tokens.is_empty()
2195                && op == "-"
2196                && !right_tokens.is_empty()
2197                && matches!(right_tokens[0], Token::Number(_))
2198                && let Token::Number(n) = right_tokens[0]
2199            {
2200                if right_tokens.len() >= 3
2201                    && let Some(bin_op) = infix_op_at(right_tokens, 1)
2202                {
2203                    if WORD_OPS.contains(&bin_op) {
2204                        // The word operators read their operands as tokens, so hand
2205                        // them the negative number as one: -7 mod 3 is (-7) mod 3.
2206                        return apply_infix(&[Token::Number(-n)], bin_op, &right_tokens[2..]);
2207                    }
2208                    let right = parse_node(&right_tokens[2..])?;
2209                    return apply_op(Node::Int(0).bin(BinOp::Sub, Node::Num(n)), bin_op, right);
2210                }
2211                if right_tokens.len() == 1 {
2212                    return Ok(Node::Int(0).bin(BinOp::Sub, Node::Num(n)));
2213                }
2214            }
2215            // Unary plus/minus when there is no left operand (e.g. -x, +x, -(a+b))
2216            if left_tokens.is_empty() && (op == "+" || op == "-") {
2217                let inner = parse_node(right_tokens)?;
2218                return if op == "-" {
2219                    Ok(Node::Int(0).bin(BinOp::Sub, inner))
2220                } else {
2221                    Ok(inner)
2222                };
2223            }
2224            if left_tokens.is_empty() {
2225                return Err("Missing left operand".to_string());
2226            }
2227            apply_infix(left_tokens, op, right_tokens)
2228        } else {
2229            Err("Expected operator".to_string())
2230        }
2231    } else {
2232        // No operator: a term. Callers pass complete expressions, so leftover tokens are an
2233        // error (not silently dropped, as `where x > 1 by dept` once lost `by dept`).
2234        let (expr, remaining) = parse_term(tokens)?;
2235        if let Some(extra) = remaining.first() {
2236            if matches!(&tokens[0], Token::Identifier(w) if w == "wavg") {
2237                return Err(WAVG_USAGE.to_string());
2238            }
2239            return Err(format!(
2240                "Unexpected '{}' after the expression",
2241                token_text(extra)
2242            ));
2243        }
2244        Ok(expr)
2245    }
2246}
2247
2248/// A parsed q query, ready to apply to a LazyFrame.
2249#[derive(Debug, Default)]
2250pub struct ParsedQuery {
2251    /// The select list; empty means every column.
2252    pub cols: Vec<Expr>,
2253    /// The where clause's conditions, applied as successive filters: each one over
2254    /// the rows the ones before it kept, so an aggregate in a later one sees only those.
2255    pub filters: Vec<Expr>,
2256    /// The by expressions.
2257    pub group_by: Vec<Expr>,
2258    /// Names of the by columns that have one (a plain column or an alias).
2259    pub group_by_names: Vec<String>,
2260    /// `select distinct`: drop duplicate result rows.
2261    pub distinct: bool,
2262}
2263
2264impl ParsedQuery {
2265    /// The query with text casts and date parts guarded for dates past the calendar,
2266    /// where Polars panics ([`crate::past_calendar::guard_expr`]). With `schema`, only
2267    /// date and datetime operations change.
2268    pub fn past_calendar_safe(self, schema: Option<&Schema>) -> Self {
2269        let guard = |e: Expr| crate::past_calendar::guard_expr(e, schema);
2270        Self {
2271            cols: self.cols.into_iter().map(guard).collect(),
2272            filters: self.filters.into_iter().map(guard).collect(),
2273            group_by: self.group_by.into_iter().map(guard).collect(),
2274            ..self
2275        }
2276    }
2277}
2278
2279/// Convert Polars-specific error messages to user-friendly query errors.
2280pub fn sanitize_query_error(msg: &str) -> String {
2281    let msg_lower = msg.to_lowercase();
2282    if msg_lower.contains("duplicate")
2283        && (msg_lower.contains("output name") || msg_lower.contains("projection"))
2284    {
2285        let name = msg
2286            .split('\'')
2287            .nth(1)
2288            .map(|s| s.to_string())
2289            .unwrap_or_else(|| "column".to_string());
2290        return format!(
2291            "Duplicate column name '{}' in result. Use aliases to rename columns, e.g. `select my_date: timestamp.date`",
2292            name
2293        );
2294    }
2295    if msg_lower.contains(".alias(") || msg_lower.contains("try renaming") {
2296        return "Duplicate column names in result. Use aliases to rename columns, e.g. `select my_date: timestamp.date`"
2297            .to_string();
2298    }
2299    msg.to_string()
2300}
2301
2302/// A q query as parsed, before it becomes Polars expressions: what
2303/// [`parse_query`] runs and "Copy as Python" writes out.
2304#[derive(Debug, Default)]
2305pub(crate) struct QueryNodes {
2306    pub cols: Vec<Node>,
2307    /// The where clause's conditions, in order; see [`ParsedQuery::filters`].
2308    pub filters: Vec<Node>,
2309    pub group_by: Vec<Node>,
2310    pub group_by_names: Vec<String>,
2311    pub distinct: bool,
2312}
2313
2314impl QueryNodes {
2315    fn into_parsed(self) -> ParsedQuery {
2316        let lower = |nodes: Vec<Node>| nodes.iter().map(Node::to_expr).collect();
2317        ParsedQuery {
2318            cols: lower(self.cols),
2319            filters: self.filters.iter().map(Node::to_predicate).collect(),
2320            group_by: lower(self.group_by),
2321            group_by_names: self.group_by_names,
2322            distinct: self.distinct,
2323        }
2324    }
2325
2326    /// Each timestamp literal read in the zone of the column it meets in `schema`; see
2327    /// [`Node::resolve_time_zones`].
2328    pub(crate) fn resolve_time_zones(&mut self, schema: &Schema) {
2329        let nodes = self
2330            .cols
2331            .iter_mut()
2332            .chain(self.filters.iter_mut())
2333            .chain(self.group_by.iter_mut());
2334        for node in nodes {
2335            node.resolve_time_zones(schema);
2336        }
2337    }
2338
2339    /// What the parser could not tell by form, settled by the types in `schema`; see
2340    /// [`Node::resolve_types`].
2341    pub(crate) fn resolve_types(&mut self, schema: &Schema) -> Result<(), String> {
2342        self.cols
2343            .iter_mut()
2344            .chain(self.filters.iter_mut())
2345            .chain(self.group_by.iter_mut())
2346            .try_for_each(|node| node.resolve_types(schema))
2347    }
2348
2349    /// Fails on a temporal column compared with quoted text; see
2350    /// [`Node::check_quoted_temporal`].
2351    fn check_quoted_temporal(&self, schema: &Schema) -> Result<(), String> {
2352        self.cols
2353            .iter()
2354            .chain(self.filters.iter())
2355            .chain(self.group_by.iter())
2356            .try_for_each(|node| node.check_quoted_temporal(schema))
2357    }
2358
2359    /// The where clause as Python: one `.filter(...)` per condition, in order, as
2360    /// datui runs them.
2361    pub(crate) fn python_filters(&self) -> Vec<String> {
2362        self.filters
2363            .iter()
2364            .map(|f| format!(".filter({})", f.python_predicate()))
2365            .collect()
2366    }
2367
2368    /// Python calls doing what `DataTableState::query` does: the where clause, then the
2369    /// grouping (ordered by keys named `key_names`) or the select list, then `distinct`.
2370    pub(crate) fn python_steps(&self, key_names: &[String]) -> Vec<String> {
2371        let mut steps = self.python_filters();
2372        if !self.group_by.is_empty() {
2373            let keys = python_list(&self.group_by);
2374            let aggs = if !self.cols.is_empty() {
2375                python_list(&self.cols)
2376            } else if self.group_by_names.is_empty() {
2377                "pl.all()".to_string()
2378            } else {
2379                let names: Vec<String> = self
2380                    .group_by_names
2381                    .iter()
2382                    .map(|n| crate::export::python_script::py_str(n))
2383                    .collect();
2384                format!("pl.all().exclude({})", names.join(", "))
2385            };
2386            steps.push(format!(".group_by({keys})"));
2387            steps.push(format!(".agg({aggs})"));
2388            steps.push(crate::export::python_script::sort_call(
2389                key_names,
2390                &vec![false; key_names.len()],
2391            ));
2392        } else if !self.cols.is_empty() {
2393            steps.push(format!(".select({})", python_list(&self.cols)));
2394        }
2395        if self.distinct {
2396            steps.push(".unique(keep=\"first\", maintain_order=True)".to_string());
2397        }
2398        steps
2399    }
2400}
2401
2402/// Expressions as Python arguments: a plain column by its name, as Polars reads a
2403/// string there, anything else as an expression.
2404fn python_list(nodes: &[Node]) -> String {
2405    nodes
2406        .iter()
2407        .map(|n| match n {
2408            Node::Col(name) => crate::export::python_script::py_str(name),
2409            n => n.python(),
2410        })
2411        .collect::<Vec<_>>()
2412        .join(", ")
2413}
2414
2415pub fn parse_query(query: &str) -> Result<ParsedQuery, String> {
2416    parse_nodes(query).map(QueryNodes::into_parsed)
2417}
2418
2419/// [`parse_query`] for data of `schema`: timestamp literals against zoned columns read
2420/// in that zone, and temporal columns compared with quoted text are errors.
2421pub fn parse_query_over(query: &str, schema: Option<&Schema>) -> Result<ParsedQuery, String> {
2422    let mut nodes = parse_nodes(query)?;
2423    if let Some(schema) = schema {
2424        nodes.resolve_time_zones(schema);
2425        nodes.check_quoted_temporal(schema)?;
2426        nodes.resolve_types(schema)?;
2427    }
2428    Ok(nodes.into_parsed())
2429}
2430
2431/// Parse a q query into nodes. An empty query selects every column.
2432pub(crate) fn parse_nodes(query: &str) -> Result<QueryNodes, String> {
2433    // Empty query is equivalent to "select" - return all columns with no filter or grouping
2434    let trimmed = query.trim();
2435    if trimmed.is_empty() {
2436        return Ok(QueryNodes::default());
2437    }
2438
2439    let tokens = tokenize(query)?;
2440    if tokens.is_empty() || tokens[0] != Token::Select {
2441        return Err("Query must start with 'select'".to_string());
2442    }
2443    // `distinct` after `select` is the keyword unless it is a column or alias
2444    // (`distinct: x`, `distinct, a`, `distinct + 1`); `col["distinct"]` always names it.
2445    let distinct = tokens.get(1) == Some(&Token::Identifier("distinct".to_string()))
2446        && !matches!(
2447            tokens.get(2),
2448            Some(Token::Colon | Token::Comma | Token::Dot | Token::Op(_))
2449        );
2450    let body = strip_from(&tokens[if distinct { 2 } else { 1 }..])?;
2451    let body = &body[..];
2452
2453    let mut parts = split_tokens(body, &Token::Where);
2454    let select_by_tokens = parts.remove(0);
2455    let where_tokens = if !parts.is_empty() {
2456        Some(parts.remove(0))
2457    } else {
2458        None
2459    };
2460    if !parts.is_empty() {
2461        return Err(
2462            "Unexpected second 'where': list the conditions after one where, separated by ','"
2463                .to_string(),
2464        );
2465    }
2466
2467    // `by` after `where` reads naturally but grouping comes first; catch it here,
2468    // outside parentheses and brackets, so the error can name the clause order.
2469    if let Some(ref wt) = where_tokens {
2470        let mut depth = 0;
2471        let mut bracket_depth = 0;
2472        for token in wt {
2473            match token {
2474                Token::LParen => depth += 1,
2475                Token::RParen => depth -= 1,
2476                Token::LBracket => bracket_depth += 1,
2477                Token::RBracket => bracket_depth -= 1,
2478                Token::By if depth == 0 && bracket_depth == 0 => {
2479                    return Err(format!(
2480                        "Unexpected 'by' after the where clause: {}",
2481                        CLAUSE_ORDER
2482                    ));
2483                }
2484                _ => {}
2485            }
2486        }
2487    }
2488
2489    let mut select_by_parts = split_tokens(&select_by_tokens, &Token::By);
2490    let cols_tokens = select_by_parts.remove(0);
2491    let by_tokens = if !select_by_parts.is_empty() {
2492        Some(select_by_parts.remove(0))
2493    } else {
2494        None
2495    };
2496    if !select_by_parts.is_empty() {
2497        return Err(format!("Unexpected second 'by': {}", CLAUSE_ORDER));
2498    }
2499
2500    let mut cols = Vec::new();
2501    if !cols_tokens.is_empty() {
2502        for chunk in split_tokens(&cols_tokens, &Token::Comma) {
2503            if chunk.is_empty() {
2504                continue;
2505            }
2506            // Find colon position (if any) - need to account for col[...] syntax
2507            let mut colon_pos = None;
2508            let mut depth = 0;
2509            for (i, token) in chunk.iter().enumerate() {
2510                match token {
2511                    Token::LBracket => depth += 1,
2512                    Token::RBracket => depth -= 1,
2513                    Token::Colon if depth == 0 => {
2514                        colon_pos = Some(i);
2515                        break;
2516                    }
2517                    _ => {}
2518                }
2519            }
2520            if let Some(pos) = colon_pos {
2521                // Has alias: parse left side for alias name, right side for expression
2522                let alias_tokens = &chunk[..pos];
2523                let expr_tokens = &chunk[pos + 1..];
2524
2525                // Parse alias - could be simple identifier or col[...]
2526                let alias_name = if alias_tokens.len() == 1 {
2527                    if let Token::Identifier(name) = &alias_tokens[0] {
2528                        name.clone()
2529                    } else {
2530                        return Err("Expected identifier or col[] for alias".to_string());
2531                    }
2532                } else if alias_tokens.len() == 4
2533                    && alias_tokens[0] == Token::Identifier("col".to_string())
2534                    && alias_tokens[1] == Token::LBracket
2535                    && alias_tokens[3] == Token::RBracket
2536                {
2537                    // col[...] syntax for alias
2538                    match &alias_tokens[2] {
2539                        Token::String(name) | Token::Identifier(name) => name.clone(),
2540                        _ => {
2541                            return Err(
2542                                "Expected string or identifier in col[] for alias".to_string()
2543                            );
2544                        }
2545                    }
2546                } else {
2547                    // Try to parse as expression and extract name (for simple cases)
2548                    // For now, require explicit identifier or col[]
2549                    return Err("Alias must be an identifier or col[]".to_string());
2550                };
2551
2552                let expr = parse_node(expr_tokens)?;
2553                cols.push(expr.alias(alias_name));
2554            } else {
2555                cols.push(parse_node(&chunk)?);
2556            }
2557        }
2558    }
2559
2560    let mut group_by_cols = Vec::new();
2561    let mut group_by_col_names = Vec::new();
2562    if let Some(bt) = by_tokens {
2563        for chunk in split_tokens(&bt, &Token::Comma) {
2564            if chunk.is_empty() {
2565                continue;
2566            }
2567            // Support column assignment in by clause (like select)
2568            // Find colon position (if any) - need to account for col[...] syntax
2569            let mut colon_pos = None;
2570            let mut depth = 0;
2571            for (i, token) in chunk.iter().enumerate() {
2572                match token {
2573                    Token::LBracket => depth += 1,
2574                    Token::RBracket => depth -= 1,
2575                    Token::Colon if depth == 0 => {
2576                        colon_pos = Some(i);
2577                        break;
2578                    }
2579                    _ => {}
2580                }
2581            }
2582            if let Some(pos) = colon_pos {
2583                // Has alias: parse left side for alias name, right side for expression
2584                let alias_tokens = &chunk[..pos];
2585                let expr_tokens = &chunk[pos + 1..];
2586
2587                // Parse alias - could be simple identifier or col[...]
2588                let alias_name = if alias_tokens.len() == 1 {
2589                    if let Token::Identifier(name) = &alias_tokens[0] {
2590                        name.clone()
2591                    } else {
2592                        return Err(
2593                            "Expected identifier or col[] for alias in by clause".to_string()
2594                        );
2595                    }
2596                } else if alias_tokens.len() == 4
2597                    && alias_tokens[0] == Token::Identifier("col".to_string())
2598                    && alias_tokens[1] == Token::LBracket
2599                    && alias_tokens[3] == Token::RBracket
2600                {
2601                    // col[...] syntax for alias
2602                    match &alias_tokens[2] {
2603                        Token::String(name) | Token::Identifier(name) => name.clone(),
2604                        _ => {
2605                            return Err(
2606                                "Expected string or identifier in col[] for alias in by clause"
2607                                    .to_string(),
2608                            );
2609                        }
2610                    }
2611                } else {
2612                    return Err("Alias must be an identifier or col[] in by clause".to_string());
2613                };
2614
2615                let expr = parse_node(expr_tokens)?;
2616                group_by_cols.push(expr.alias(alias_name.clone()));
2617                group_by_col_names.push(alias_name); // Use alias name
2618            } else {
2619                let expr = parse_node(&chunk)?;
2620                group_by_cols.push(expr.clone());
2621                // The column name of a simple expression: `name`, or `col[name]`.
2622                if chunk.len() == 1 {
2623                    if let Token::Identifier(name) = &chunk[0] {
2624                        group_by_col_names.push(name.clone());
2625                    }
2626                } else if chunk.len() == 4
2627                    && chunk[0] == Token::Identifier("col".to_string())
2628                    && chunk[1] == Token::LBracket
2629                    && chunk[3] == Token::RBracket
2630                {
2631                    // col[...] syntax
2632                    match &chunk[2] {
2633                        Token::String(name) | Token::Identifier(name) => {
2634                            group_by_col_names.push(name.clone());
2635                        }
2636                        _ => {}
2637                    }
2638                } else {
2639                    // Complex expressions have no simple name; sorting uses the expression itself.
2640                }
2641            }
2642        }
2643    }
2644
2645    // q splits where only on its top-level commas, into conditions applied in turn;
2646    // each is an ordinary expression (`&` and `|` are operators like any other).
2647    let mut filters = Vec::new();
2648    if let Some(wt) = where_tokens.filter(|wt| !wt.is_empty()) {
2649        for chunk in split_tokens(&wt, &Token::Comma) {
2650            if chunk.is_empty() {
2651                return Err(
2652                    "Empty condition in where: a ',' with nothing before or after it".to_string(),
2653                );
2654            }
2655            filters.push(parse_node(&chunk)?);
2656        }
2657    }
2658
2659    Ok(QueryNodes {
2660        cols,
2661        filters,
2662        group_by: group_by_cols,
2663        group_by_names: group_by_col_names,
2664        distinct,
2665    })
2666}
2667
2668/// One expression as Polars runs it.
2669#[cfg(test)]
2670fn parse_expr(tokens: &[Token]) -> Result<Expr, String> {
2671    parse_node(tokens).map(|n| n.to_expr())
2672}
2673
2674#[cfg(test)]
2675mod tests;