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_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 DateLiteral(String),
20 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
39fn 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(); 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
90fn 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(); 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(); 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 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(); 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 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 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
328fn 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
350const CLAUSE_ORDER: &str = "clause order is select [by group] [where conditions]";
352
353const FROM_ORDER: &str = "clause order is select [by group] [from df] [where conditions]";
355
356const TABLE: &str = "df";
358
359fn 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
381fn 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
424const WORD_OPS: [&str; 7] = ["in", "like", "xbar", "mod", "wavg", "and", "or"];
428
429fn 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#[derive(Debug, Clone, PartialEq)]
446pub(crate) enum Node {
447 Col(String),
448 Num(f64),
450 Int(i64),
452 Str(String),
453 Bool(bool),
454 Null,
455 Date(String),
457 Timestamp {
459 iso: String,
460 format: String,
461 unit: TimeUnit,
462 zone: Option<String>,
465 },
466 Bin(BinOp, Box<Node>, Box<Node>),
467 Coalesce(Box<Node>, Box<Node>),
468 Filter(Box<Node>, Box<Node>),
470 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 TrueDiv,
483 FloorDiv,
485 Rem,
486 Eq,
487 Neq,
488 Lt,
489 Gt,
490 LtEq,
491 GtEq,
492 And,
496 Or,
498 Lesser,
501 Greater,
503 Any,
505}
506
507#[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 ContainsRegex(String),
553 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 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 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 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
632const MAX_EXPR_NODES: usize = 10_000;
635
636fn 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 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 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 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 let zone = TimeZone::opt_try_new(zone.as_deref()).ok().flatten();
712 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 fn python_operand(&self) -> String {
1054 match self {
1055 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1065enum Kind {
1066 Bool,
1067 Int,
1068 Float,
1069 Temporal,
1070 Text,
1071 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
1101fn 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
1119fn 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 (Kind::Temporal, Kind::Bool | Kind::Int | Kind::Float)
1153 | (Kind::Bool | Kind::Int | Kind::Float, Kind::Temporal) => return refuse(l, r),
1154 (Kind::Bool, k) if number(k) => bool_as(left, right, schema),
1156 (k, Kind::Bool) if number(k) => bool_as(right, left, schema),
1157 (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
1173fn 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
1185fn null_if_either(left: &Expr, right: &Expr) -> Then {
1187 when(left.clone().is_null().or(right.clone().is_null())).then(lit(NULL))
1188}
1189
1190fn 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 result.expect("min/max of two expressions")
1200}
1201
1202fn 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 BinOp::And | BinOp::Lesser => "&",
1234 BinOp::Or | BinOp::Greater => "|",
1235 }
1236 }
1237}
1238
1239impl Op {
1240 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 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 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 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 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 Op::Round(decimals) => expr.round(*decimals, RoundMode::HalfAwayFromZero),
1393 Op::Cast(to) => expr.cast(to.dtype()),
1395 }
1396}
1397
1398fn 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
1411fn 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
1424fn 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
1468pub(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
1477fn 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
1485fn 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
1495fn 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(®ex::escape(c.encode_utf8(&mut [0; 4]))),
1504 }
1505 }
1506 re.push('$');
1507 re
1508}
1509
1510fn 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 [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 if left.size() > 1 {
1535 check_copies(&left, items.len())?;
1536 }
1537 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 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 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 check_copies(&weights, 5)?;
1583 check_copies(&values, 3)?;
1584 let weighted = weights.clone().bin(BinOp::Mul, values);
1585 let total = Node::Filter(
1588 Box::new(weights),
1589 Box::new(weighted.clone().op(Op::IsNotNull)),
1590 )
1591 .op(Op::Sum);
1592 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 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 "%" | "/" => 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 "&" | "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
1645fn 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
1662const AGG_FUNCTIONS: [&str; 16] = [
1664 "avg", "mean", "min", "max", "count", "std", "stddev", "dev", "var", "med", "median", "sum",
1665 "first", "last", "nunique", "wavg",
1666];
1667
1668const 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 let name = name.to_lowercase();
1681 name != "wavg" && (is_agg_function(&name) || SCALAR_FUNCTIONS.contains(&name.as_str()))
1682}
1683
1684fn 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
1694fn 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 match simple_column_name(args) {
1725 Some(column) => Ok(node.alias(format!("{}_{}", fn_name, column))),
1726 None => Ok(node),
1727 }
1728}
1729
1730fn 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#[derive(Debug, Clone, PartialEq)]
1759enum AccessorArg {
1760 Str(String),
1761 Num(f64),
1762}
1763
1764impl AccessorArg {
1765 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
1774const ACCESSORS: &[(&str, usize, usize, &str)] = &[
1776 ("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 ("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 ("round", 0, 1, ".round[1]"),
1809 ("int", 0, 0, ".int"),
1810 ("float", 0, 0, ".float"),
1811 ("str", 0, 0, ".str"),
1812];
1813
1814const 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
1828fn 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 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 _ => 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
1932fn 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
1951fn 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 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 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 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..])), Token::String(s) => Ok((Node::Str(s.clone()), &tokens[1..])), 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 _ => Err(format!(
2099 "Unexpected '{}' where an expression was expected",
2100 token_text(&tokens[0])
2101 )),
2102 }
2103}
2104
2105const PAREN_COMMA: &str = "`,` inside parentheses joins lists in q, which datui does not \
2107 support; combine conditions with `&` or `|`";
2108
2109const MAX_EXPR_DEPTH: u32 = 64;
2114
2115thread_local! {
2116 static EXPR_DEPTH: std::cell::Cell<u32> = const { std::cell::Cell::new(0) };
2117}
2118
2119struct DepthGuard;
2122
2123impl DepthGuard {
2124 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
2143fn 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 if let Token::Identifier(name) = &tokens[0]
2159 && is_function_name(name)
2160 && tokens.len() > 1
2161 && tokens[1] != Token::LBracket
2162 {
2163 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 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 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 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 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 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#[derive(Debug, Default)]
2248pub struct ParsedQuery {
2249 pub cols: Vec<Expr>,
2251 pub filters: Vec<Expr>,
2254 pub group_by: Vec<Expr>,
2256 pub group_by_names: Vec<String>,
2258 pub distinct: bool,
2260}
2261
2262impl ParsedQuery {
2263 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
2277pub 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#[derive(Debug, Default)]
2303pub(crate) struct QueryNodes {
2304 pub cols: Vec<Node>,
2305 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 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 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 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 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 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
2400fn 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
2417pub 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
2429pub(crate) fn parse_nodes(query: &str) -> Result<QueryNodes, String> {
2431 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 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 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 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 let alias_tokens = &chunk[..pos];
2521 let expr_tokens = &chunk[pos + 1..];
2522
2523 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 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 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 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 let alias_tokens = &chunk[..pos];
2583 let expr_tokens = &chunk[pos + 1..];
2584
2585 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 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); } else {
2617 let expr = parse_node(&chunk)?;
2618 group_by_cols.push(expr.clone());
2619 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 match &chunk[2] {
2631 Token::String(name) | Token::Identifier(name) => {
2632 group_by_col_names.push(name.clone());
2633 }
2634 _ => {}
2635 }
2636 } else {
2637 }
2639 }
2640 }
2641 }
2642
2643 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#[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;