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