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