polars_sql/
sql_expr.rs

Help
1//! Expressions that are supported by the Polars SQL interface.
2//!
3//! This is useful for syntax highlighting
4//!
5//! This module defines:
6//! - all Polars SQL keywords [`all_keywords`]
7//! - all of polars SQL functions [`all_functions`]
8
9use std::fmt::Display;
10use std::ops::Div;
11
12use polars_core::prelude::*;
13use polars_lazy::prelude::*;
14use polars_plan::plans::DynLiteralValue;
15use polars_plan::prelude::typed_lit;
16use polars_time::Duration;
17use rand::distributions::Alphanumeric;
18use rand::{Rng, thread_rng};
19#[cfg(feature = "serde")]
20use serde::{Deserialize, Serialize};
21use sqlparser::ast::{
22    BinaryOperator as SQLBinaryOperator, CastFormat, CastKind, DataType as SQLDataType,
23    DateTimeField, Expr as SQLExpr, Function as SQLFunction, Ident, Interval, Query as Subquery,
24    SelectItem, Subscript, TimezoneInfo, TrimWhereField, UnaryOperator, Value as SQLValue,
25};
26use sqlparser::dialect::GenericDialect;
27use sqlparser::parser::{Parser, ParserOptions};
28
29use crate::SQLContext;
30use crate::functions::SQLFunctionVisitor;
31use crate::types::{
32    bitstring_to_bytes_literal, is_iso_date, is_iso_datetime, is_iso_time, map_sql_dtype_to_polars,
33};
34
35#[inline]
36#[cold]
37#[must_use]
38/// Convert a Display-able error to PolarsError::SQLInterface
39pub fn to_sql_interface_err(err: impl Display) -> PolarsError {
40    PolarsError::SQLInterface(err.to_string().into())
41}
42
43#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
44#[derive(Clone, Copy, PartialEq, Debug, Eq, Hash)]
45/// Categorises the type of (allowed) subquery constraint
46pub enum SubqueryRestriction {
47    /// Subquery must return a single column
48    SingleColumn,
49    // SingleRow,
50    // SingleValue,
51    // Any
52}
53
54/// Recursively walks a SQL Expr to create a polars Expr
55pub(crate) struct SQLExprVisitor<'a> {
56    ctx: &'a mut SQLContext,
57    active_schema: Option<&'a Schema>,
58}
59
60impl SQLExprVisitor<'_> {
61    fn array_expr_to_series(&mut self, elements: &[SQLExpr]) -> PolarsResult<Series> {
62        let array_elements = elements
63            .iter()
64            .map(|e| match e {
65                SQLExpr::Value(v) => self.visit_any_value(v, None),
66                SQLExpr::UnaryOp { op, expr } => match expr.as_ref() {
67                    SQLExpr::Value(v) => self.visit_any_value(v, Some(op)),
68                    _ => Err(polars_err!(SQLInterface: "expression {:?} is not currently supported", e)),
69                },
70                SQLExpr::Array(_) => {
71                    // TODO: nested arrays (handle FnMut issues)
72                    // let srs = self.array_expr_to_series(&[e.clone()])?;
73                    // Ok(AnyValue::List(srs))
74                    Err(polars_err!(SQLInterface: "nested array literals are not currently supported:\n{:?}", e))
75                },
76                _ => Err(polars_err!(SQLInterface: "expression {:?} is not currently supported", e)),
77            })
78            .collect::<PolarsResult<Vec<_>>>()?;
79
80        Series::from_any_values(PlSmallStr::EMPTY, &array_elements, true)
81    }
82
83    fn visit_expr(&mut self, expr: &SQLExpr) -> PolarsResult<Expr> {
84        match expr {
85            SQLExpr::AllOp {
86                left,
87                compare_op,
88                right,
89            } => self.visit_all(left, compare_op, right),
90            SQLExpr::AnyOp {
91                left,
92                compare_op,
93                right,
94                is_some: _,
95            } => self.visit_any(left, compare_op, right),
96            SQLExpr::Array(arr) => self.visit_array_expr(&arr.elem, true, None),
97            SQLExpr::Between {
98                expr,
99                negated,
100                low,
101                high,
102            } => self.visit_between(expr, *negated, low, high),
103            SQLExpr::BinaryOp { left, op, right } => self.visit_binary_op(left, op, right),
104            SQLExpr::Cast {
105                kind,
106                expr,
107                data_type,
108                format,
109            } => self.visit_cast(expr, data_type, format, kind),
110            SQLExpr::Ceil { expr, .. } => Ok(self.visit_expr(expr)?.ceil()),
111            SQLExpr::CompoundIdentifier(idents) => self.visit_compound_identifier(idents),
112            SQLExpr::Extract {
113                field,
114                syntax: _,
115                expr,
116            } => parse_extract_date_part(self.visit_expr(expr)?, field),
117            SQLExpr::Floor { expr, .. } => Ok(self.visit_expr(expr)?.floor()),
118            SQLExpr::Function(function) => self.visit_function(function),
119            SQLExpr::Identifier(ident) => self.visit_identifier(ident),
120            SQLExpr::InList {
121                expr,
122                list,
123                negated,
124            } => {
125                let expr = self.visit_expr(expr)?;
126                let elems = self.visit_array_expr(list, false, Some(&expr))?;
127                let is_in = expr.is_in(elems, false);
128                Ok(if *negated { is_in.not() } else { is_in })
129            },
130            SQLExpr::InSubquery {
131                expr,
132                subquery,
133                negated,
134            } => self.visit_in_subquery(expr, subquery, *negated),
135            SQLExpr::Interval(interval) => Ok(lit(interval_to_duration(interval, true)?)),
136            SQLExpr::IsDistinctFrom(e1, e2) => {
137                Ok(self.visit_expr(e1)?.neq_missing(self.visit_expr(e2)?))
138            },
139            SQLExpr::IsFalse(expr) => Ok(self.visit_expr(expr)?.eq(lit(false))),
140            SQLExpr::IsNotDistinctFrom(e1, e2) => {
141                Ok(self.visit_expr(e1)?.eq_missing(self.visit_expr(e2)?))
142            },
143            SQLExpr::IsNotFalse(expr) => Ok(self.visit_expr(expr)?.eq(lit(false)).not()),
144            SQLExpr::IsNotNull(expr) => Ok(self.visit_expr(expr)?.is_not_null()),
145            SQLExpr::IsNotTrue(expr) => Ok(self.visit_expr(expr)?.eq(lit(true)).not()),
146            SQLExpr::IsNull(expr) => Ok(self.visit_expr(expr)?.is_null()),
147            SQLExpr::IsTrue(expr) => Ok(self.visit_expr(expr)?.eq(lit(true))),
148            SQLExpr::Like {
149                negated,
150                any,
151                expr,
152                pattern,
153                escape_char,
154            } => {
155                if *any {
156                    polars_bail!(SQLSyntax: "LIKE ANY is not a supported syntax")
157                }
158                self.visit_like(*negated, expr, pattern, escape_char, false)
159            },
160            SQLExpr::ILike {
161                negated,
162                any,
163                expr,
164                pattern,
165                escape_char,
166            } => {
167                if *any {
168                    polars_bail!(SQLSyntax: "ILIKE ANY is not a supported syntax")
169                }
170                self.visit_like(*negated, expr, pattern, escape_char, true)
171            },
172            SQLExpr::Nested(expr) => self.visit_expr(expr),
173            SQLExpr::Position { expr, r#in } => Ok(
174                // note: SQL is 1-indexed
175                (self
176                    .visit_expr(r#in)?
177                    .str()
178                    .find(self.visit_expr(expr)?, true)
179                    + typed_lit(1u32))
180                .fill_null(typed_lit(0u32)),
181            ),
182            SQLExpr::RLike {
183                // note: parses both RLIKE and REGEXP
184                negated,
185                expr,
186                pattern,
187                regexp: _,
188            } => {
189                let matches = self
190                    .visit_expr(expr)?
191                    .str()
192                    .contains(self.visit_expr(pattern)?, true);
193                Ok(if *negated { matches.not() } else { matches })
194            },
195            SQLExpr::Subscript { expr, subscript } => self.visit_subscript(expr, subscript),
196            SQLExpr::Subquery(_) => polars_bail!(SQLInterface: "unexpected subquery"),
197            SQLExpr::Trim {
198                expr,
199                trim_where,
200                trim_what,
201                trim_characters,
202            } => self.visit_trim(expr, trim_where, trim_what, trim_characters),
203            SQLExpr::TypedString { data_type, value } => match data_type {
204                SQLDataType::Date => {
205                    if is_iso_date(value) {
206                        Ok(lit(value.as_str()).cast(DataType::Date))
207                    } else {
208                        polars_bail!(SQLSyntax: "invalid DATE literal '{}'", value)
209                    }
210                },
211                SQLDataType::Time(None, TimezoneInfo::None) => {
212                    if is_iso_time(value) {
213                        Ok(lit(value.as_str()).str().to_time(StrptimeOptions {
214                            strict: true,
215                            ..Default::default()
216                        }))
217                    } else {
218                        polars_bail!(SQLSyntax: "invalid TIME literal '{}'", value)
219                    }
220                },
221                SQLDataType::Timestamp(None, TimezoneInfo::None) | SQLDataType::Datetime(None) => {
222                    if is_iso_datetime(value) {
223                        Ok(lit(value.as_str()).str().to_datetime(
224                            None,
225                            None,
226                            StrptimeOptions {
227                                strict: true,
228                                ..Default::default()
229                            },
230                            lit("latest"),
231                        ))
232                    } else {
233                        let fn_name = match data_type {
234                            SQLDataType::Timestamp(_, _) => "TIMESTAMP",
235                            SQLDataType::Datetime(_) => "DATETIME",
236                            _ => unreachable!(),
237                        };
238                        polars_bail!(SQLSyntax: "invalid {} literal '{}'", fn_name, value)
239                    }
240                },
241                _ => {
242                    polars_bail!(SQLInterface: "typed literal should be one of DATE, DATETIME, TIME, or TIMESTAMP (found {})", data_type)
243                },
244            },
245            SQLExpr::UnaryOp { op, expr } => self.visit_unary_op(op, expr),
246            SQLExpr::Value(value) => self.visit_literal(value),
247            SQLExpr::Wildcard(_) => Ok(Expr::Wildcard),
248            e @ SQLExpr::Case { .. } => self.visit_case_when_then(e),
249            other => {
250                polars_bail!(SQLInterface: "expression {:?} is not currently supported", other)
251            },
252        }
253    }
254
255    fn visit_subquery(
256        &mut self,
257        subquery: &Subquery,
258        restriction: SubqueryRestriction,
259    ) -> PolarsResult<Expr> {
260        if subquery.with.is_some() {
261            polars_bail!(SQLSyntax: "SQL subquery cannot be a CTE 'WITH' clause");
262        }
263        let mut lf = self.ctx.execute_query_no_ctes(subquery)?;
264        let schema = self.ctx.get_frame_schema(&mut lf)?;
265
266        if restriction == SubqueryRestriction::SingleColumn {
267            if schema.len() != 1 {
268                polars_bail!(SQLSyntax: "SQL subquery returns more than one column");
269            }
270            let rand_string: String = thread_rng()
271                .sample_iter(&Alphanumeric)
272                .take(16)
273                .map(char::from)
274                .collect();
275
276            let schema_entry = schema.get_at_index(0);
277            if let Some((old_name, _)) = schema_entry {
278                let new_name = String::from(old_name.as_str()) + rand_string.as_str();
279                lf = lf.rename([old_name.to_string()], [new_name.clone()], true);
280                return Ok(Expr::SubPlan(
281                    SpecialEq::new(Arc::new(lf.logical_plan)),
282                    vec![new_name],
283                ));
284            }
285        };
286        polars_bail!(SQLInterface: "subquery type not supported");
287    }
288
289    /// Visit a single SQL identifier.
290    ///
291    /// e.g. column
292    fn visit_identifier(&self, ident: &Ident) -> PolarsResult<Expr> {
293        Ok(col(ident.value.as_str()))
294    }
295
296    /// Visit a compound SQL identifier
297    ///
298    /// e.g. tbl.column, struct.field, tbl.struct.field (inc. nested struct fields)
299    fn visit_compound_identifier(&mut self, idents: &[Ident]) -> PolarsResult<Expr> {
300        Ok(resolve_compound_identifier(self.ctx, idents, self.active_schema)?[0].clone())
301    }
302
303    fn visit_like(
304        &mut self,
305        negated: bool,
306        expr: &SQLExpr,
307        pattern: &SQLExpr,
308        escape_char: &Option<String>,
309        case_insensitive: bool,
310    ) -> PolarsResult<Expr> {
311        if escape_char.is_some() {
312            polars_bail!(SQLInterface: "ESCAPE char for LIKE/ILIKE is not currently supported; found '{}'", escape_char.clone().unwrap());
313        }
314        let pat = match self.visit_expr(pattern) {
315            Ok(Expr::Literal(lv)) if lv.extract_str().is_some() => {
316                PlSmallStr::from_str(lv.extract_str().unwrap())
317            },
318            _ => {
319                polars_bail!(SQLSyntax: "LIKE/ILIKE pattern must be a string literal; found {}", pattern)
320            },
321        };
322        if pat.is_empty() || (!case_insensitive && pat.chars().all(|c| !matches!(c, '%' | '_'))) {
323            // empty string or other exact literal match (eg: no wildcard chars)
324            let op = if negated {
325                SQLBinaryOperator::NotEq
326            } else {
327                SQLBinaryOperator::Eq
328            };
329            self.visit_binary_op(expr, &op, pattern)
330        } else {
331            // create regex from pattern containing SQL wildcard chars ('%' => '.*', '_' => '.')
332            let mut rx = regex::escape(pat.as_str())
333                .replace('%', ".*")
334                .replace('_', ".");
335
336            rx = format!(
337                "^{}{}$",
338                if case_insensitive { "(?is)" } else { "(?s)" },
339                rx
340            );
341
342            let expr = self.visit_expr(expr)?;
343            let matches = expr.str().contains(lit(rx), true);
344            Ok(if negated { matches.not() } else { matches })
345        }
346    }
347
348    fn visit_subscript(&mut self, expr: &SQLExpr, subscript: &Subscript) -> PolarsResult<Expr> {
349        let expr = self.visit_expr(expr)?;
350        Ok(match subscript {
351            Subscript::Index { index } => {
352                let idx = adjust_one_indexed_param(self.visit_expr(index)?, true);
353                expr.list().get(idx, true)
354            },
355            Subscript::Slice { .. } => {
356                polars_bail!(SQLSyntax: "array slice syntax is not currently supported")
357            },
358        })
359    }
360
361    /// Handle implicit temporal string comparisons.
362    ///
363    /// eg: "dt >= '2024-04-30'", or "dtm::date = '2077-10-10'"
364    fn convert_temporal_strings(&mut self, left: &Expr, right: &Expr) -> Expr {
365        if let (Some(name), Some(s), expr_dtype) = match (left, right) {
366            // identify "col <op> string" expressions
367            (Expr::Column(name), Expr::Literal(lv)) if lv.extract_str().is_some() => {
368                (Some(name.clone()), Some(lv.extract_str().unwrap()), None)
369            },
370            // identify "CAST(expr AS type) <op> string" and/or "expr::type <op> string" expressions
371            (Expr::Cast { expr, dtype, .. }, Expr::Literal(lv)) if lv.extract_str().is_some() => {
372                let s = lv.extract_str().unwrap();
373                match &**expr {
374                    Expr::Column(name) => (Some(name.clone()), Some(s), Some(dtype)),
375                    _ => (None, Some(s), Some(dtype)),
376                }
377            },
378            _ => (None, None, None),
379        } {
380            if expr_dtype.is_none() && self.active_schema.is_none() {
381                right.clone()
382            } else {
383                let left_dtype = expr_dtype.or_else(|| {
384                    self.active_schema
385                        .as_ref()
386                        .and_then(|schema| schema.get(&name))
387                });
388                match left_dtype {
389                    Some(DataType::Time) if is_iso_time(s) => {
390                        right.clone().str().to_time(StrptimeOptions {
391                            strict: true,
392                            ..Default::default()
393                        })
394                    },
395                    Some(DataType::Date) if is_iso_date(s) => {
396                        right.clone().str().to_date(StrptimeOptions {
397                            strict: true,
398                            ..Default::default()
399                        })
400                    },
401                    Some(DataType::Datetime(tu, tz)) if is_iso_datetime(s) || is_iso_date(s) => {
402                        if s.len() == 10 {
403                            // handle upcast from ISO date string (10 chars) to datetime
404                            lit(format!("{}T00:00:00", s))
405                        } else {
406                            lit(s.replacen(' ', "T", 1))
407                        }
408                        .str()
409                        .to_datetime(
410                            Some(*tu),
411                            tz.clone(),
412                            StrptimeOptions {
413                                strict: true,
414                                ..Default::default()
415                            },
416                            lit("latest"),
417                        )
418                    },
419                    _ => right.clone(),
420                }
421            }
422        } else {
423            right.clone()
424        }
425    }
426
427    fn struct_field_access_expr(
428        &mut self,
429        expr: &Expr,
430        path: &str,
431        infer_index: bool,
432    ) -> PolarsResult<Expr> {
433        let path_elems = if path.starts_with('{') && path.ends_with('}') {
434            path.trim_matches(|c| c == '{' || c == '}')
435        } else {
436            path
437        }
438        .split(',');
439
440        let mut expr = expr.clone();
441        for p in path_elems {
442            let p = p.trim();
443            expr = if infer_index {
444                match p.parse::<i64>() {
445                    Ok(idx) => expr.list().get(lit(idx), true),
446                    Err(_) => expr.struct_().field_by_name(p),
447                }
448            } else {
449                expr.struct_().field_by_name(p)
450            }
451        }
452        Ok(expr)
453    }
454
455    /// Visit a SQL binary operator.
456    ///
457    /// e.g. "column + 1", "column1 <= column2"
458    fn visit_binary_op(
459        &mut self,
460        left: &SQLExpr,
461        op: &SQLBinaryOperator,
462        right: &SQLExpr,
463    ) -> PolarsResult<Expr> {
464        // need special handling for interval offsets and comparisons
465        let (lhs, mut rhs) = match (left, op, right) {
466            (_, SQLBinaryOperator::Minus, SQLExpr::Interval(v)) => {
467                let duration = interval_to_duration(v, false)?;
468                return Ok(self
469                    .visit_expr(left)?
470                    .dt()
471                    .offset_by(lit(format!("-{}", duration))));
472            },
473            (_, SQLBinaryOperator::Plus, SQLExpr::Interval(v)) => {
474                let duration = interval_to_duration(v, false)?;
475                return Ok(self
476                    .visit_expr(left)?
477                    .dt()
478                    .offset_by(lit(format!("{}", duration))));
479            },
480            (SQLExpr::Interval(v1), _, SQLExpr::Interval(v2)) => {
481                // shortcut interval comparison evaluation (-> bool)
482                let d1 = interval_to_duration(v1, false)?;
483                let d2 = interval_to_duration(v2, false)?;
484                let res = match op {
485                    SQLBinaryOperator::Gt => Ok(lit(d1 > d2)),
486                    SQLBinaryOperator::Lt => Ok(lit(d1 < d2)),
487                    SQLBinaryOperator::GtEq => Ok(lit(d1 >= d2)),
488                    SQLBinaryOperator::LtEq => Ok(lit(d1 <= d2)),
489                    SQLBinaryOperator::NotEq => Ok(lit(d1 != d2)),
490                    SQLBinaryOperator::Eq | SQLBinaryOperator::Spaceship => Ok(lit(d1 == d2)),
491                    _ => polars_bail!(SQLInterface: "invalid interval comparison operator"),
492                };
493                if res.is_ok() {
494                    return res;
495                }
496                (self.visit_expr(left)?, self.visit_expr(right)?)
497            },
498            _ => (self.visit_expr(left)?, self.visit_expr(right)?),
499        };
500        rhs = self.convert_temporal_strings(&lhs, &rhs);
501
502        Ok(match op {
503            // ----
504            // Bitwise operators
505            // ----
506            SQLBinaryOperator::BitwiseAnd => lhs.and(rhs),  // "x & y"
507            SQLBinaryOperator::BitwiseOr => lhs.or(rhs),  // "x | y"
508            SQLBinaryOperator::Xor => lhs.xor(rhs),  // "x XOR y"
509
510            // ----
511            // General operators
512            // ----
513            SQLBinaryOperator::And => lhs.and(rhs),  // "x AND y"
514            SQLBinaryOperator::Divide => lhs / rhs,  // "x / y"
515            SQLBinaryOperator::DuckIntegerDivide => lhs.floor_div(rhs).cast(DataType::Int64),  // "x // y"
516            SQLBinaryOperator::Eq => lhs.eq(rhs),  // "x = y"
517            SQLBinaryOperator::Gt => lhs.gt(rhs),  // "x > y"
518            SQLBinaryOperator::GtEq => lhs.gt_eq(rhs),  // "x >= y"
519            SQLBinaryOperator::Lt => lhs.lt(rhs),  // "x < y"
520            SQLBinaryOperator::LtEq => lhs.lt_eq(rhs),  // "x <= y"
521            SQLBinaryOperator::Minus => lhs - rhs,  // "x - y"
522            SQLBinaryOperator::Modulo => lhs % rhs,  // "x % y"
523            SQLBinaryOperator::Multiply => lhs * rhs,  // "x * y"
524            SQLBinaryOperator::NotEq => lhs.eq(rhs).not(),  // "x != y"
525            SQLBinaryOperator::Or => lhs.or(rhs),  // "x OR y"
526            SQLBinaryOperator::Plus => lhs + rhs,  // "x + y"
527            SQLBinaryOperator::Spaceship => lhs.eq_missing(rhs),  // "x <=> y"
528            SQLBinaryOperator::StringConcat => {  // "x || y"
529                lhs.cast(DataType::String) + rhs.cast(DataType::String)
530            },
531            SQLBinaryOperator::PGStartsWith => lhs.str().starts_with(rhs),  // "x ^@ y"
532            // ----
533            // Regular expression operators
534            // ----
535            SQLBinaryOperator::PGRegexMatch => match rhs {  // "x ~ y"
536                Expr::Literal(ref lv) if lv.extract_str().is_some() => lhs.str().contains(rhs, true),
537                _ => polars_bail!(SQLSyntax: "invalid pattern for '~' operator: {:?}", rhs),
538            },
539            SQLBinaryOperator::PGRegexNotMatch => match rhs {  // "x !~ y"
540                Expr::Literal(ref lv) if lv.extract_str().is_some() => lhs.str().contains(rhs, true).not(),
541                _ => polars_bail!(SQLSyntax: "invalid pattern for '!~' operator: {:?}", rhs),
542            },
543            SQLBinaryOperator::PGRegexIMatch => match rhs {  // "x ~* y"
544                Expr::Literal(ref lv) if lv.extract_str().is_some() => {
545                    let pat = lv.extract_str().unwrap();
546                    lhs.str().contains(lit(format!("(?i){}", pat)), true)
547                },
548                _ => polars_bail!(SQLSyntax: "invalid pattern for '~*' operator: {:?}", rhs),
549            },
550            SQLBinaryOperator::PGRegexNotIMatch => match rhs {  // "x !~* y"
551                Expr::Literal(ref lv) if lv.extract_str().is_some() => {
552                    let pat = lv.extract_str().unwrap();
553                    lhs.str().contains(lit(format!("(?i){}", pat)), true).not()
554                },
555                _ => {
556                    polars_bail!(SQLSyntax: "invalid pattern for '!~*' operator: {:?}", rhs)
557                },
558            },
559            // ----
560            // LIKE/ILIKE operators
561            // ----
562            SQLBinaryOperator::PGLikeMatch  // "x ~~ y"
563            | SQLBinaryOperator::PGNotLikeMatch  // "x !~~ y"
564            | SQLBinaryOperator::PGILikeMatch  // "x ~~* y"
565            | SQLBinaryOperator::PGNotILikeMatch => {  // "x !~~* y"
566                let expr = if matches!(
567                    op,
568                    SQLBinaryOperator::PGLikeMatch | SQLBinaryOperator::PGNotLikeMatch
569                ) {
570                    SQLExpr::Like {
571                        negated: matches!(op, SQLBinaryOperator::PGNotLikeMatch),
572                        any: false,
573                        expr: Box::new(left.clone()),
574                        pattern: Box::new(right.clone()),
575                        escape_char: None,
576                    }
577                } else {
578                    SQLExpr::ILike {
579                        negated: matches!(op, SQLBinaryOperator::PGNotILikeMatch),
580                        any: false,
581                        expr: Box::new(left.clone()),
582                        pattern: Box::new(right.clone()),
583                        escape_char: None,
584                    }
585                };
586                self.visit_expr(&expr)?
587            },
588            // ----
589            // JSON/Struct field access operators
590            // ----
591            SQLBinaryOperator::Arrow | SQLBinaryOperator::LongArrow => match rhs {  // "x -> y", "x ->> y"
592                Expr::Literal(lv) if lv.extract_str().is_some() => {
593                    let path = lv.extract_str().unwrap();
594                    let mut expr = self.struct_field_access_expr(&lhs, path, false)?;
595                    if let SQLBinaryOperator::LongArrow = op {
596                        expr = expr.cast(DataType::String);
597                    }
598                    expr
599                },
600                Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Int(idx))) => {
601                    let mut expr = self.struct_field_access_expr(&lhs, &idx.to_string(), true)?;
602                    if let SQLBinaryOperator::LongArrow = op {
603                        expr = expr.cast(DataType::String);
604                    }
605                    expr
606                },
607                _ => {
608                    polars_bail!(SQLSyntax: "invalid json/struct path-extract definition: {:?}", right)
609                },
610            },
611            SQLBinaryOperator::HashArrow | SQLBinaryOperator::HashLongArrow => {  // "x #> y", "x #>> y"
612                match rhs {
613                    Expr::Literal(lv) if lv.extract_str().is_some() => {
614                        let path = lv.extract_str().unwrap();
615                        let mut expr = self.struct_field_access_expr(&lhs, path, true)?;
616                        if let SQLBinaryOperator::HashLongArrow = op {
617                            expr = expr.cast(DataType::String);
618                        }
619                        expr
620                    },
621                    _ => {
622                        polars_bail!(SQLSyntax: "invalid json/struct path-extract definition: {:?}", rhs)
623                    }
624                }
625            },
626            other => {
627                polars_bail!(SQLInterface: "operator {:?} is not currently supported", other)
628            },
629        })
630    }
631
632    /// Visit a SQL unary operator.
633    ///
634    /// e.g. +column or -column
635    fn visit_unary_op(&mut self, op: &UnaryOperator, expr: &SQLExpr) -> PolarsResult<Expr> {
636        let expr = self.visit_expr(expr)?;
637        Ok(match (op, expr.clone()) {
638            // simplify the parse tree by special-casing common unary +/- ops
639            (UnaryOperator::Plus, Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Int(n)))) => {
640                lit(n)
641            },
642            (UnaryOperator::Plus, Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Float(n)))) => {
643                lit(n)
644            },
645            (UnaryOperator::Minus, Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Int(n)))) => {
646                lit(-n)
647            },
648            (UnaryOperator::Minus, Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Float(n)))) => {
649                lit(-n)
650            },
651            // general case
652            (UnaryOperator::Plus, _) => lit(0) + expr,
653            (UnaryOperator::Minus, _) => lit(0) - expr,
654            (UnaryOperator::Not, _) => expr.not(),
655            other => polars_bail!(SQLInterface: "unary operator {:?} is not supported", other),
656        })
657    }
658
659    /// Visit a SQL function.
660    ///
661    /// e.g. SUM(column) or COUNT(*)
662    ///
663    /// See [SQLFunctionVisitor] for more details
664    fn visit_function(&mut self, function: &SQLFunction) -> PolarsResult<Expr> {
665        let mut visitor = SQLFunctionVisitor {
666            func: function,
667            ctx: self.ctx,
668            active_schema: self.active_schema,
669        };
670        visitor.visit_function()
671    }
672
673    /// Visit a SQL `ALL` expression.
674    ///
675    /// e.g. `a > ALL(y)`
676    fn visit_all(
677        &mut self,
678        left: &SQLExpr,
679        compare_op: &SQLBinaryOperator,
680        right: &SQLExpr,
681    ) -> PolarsResult<Expr> {
682        let left = self.visit_expr(left)?;
683        let right = self.visit_expr(right)?;
684
685        match compare_op {
686            SQLBinaryOperator::Gt => Ok(left.gt(right.max())),
687            SQLBinaryOperator::Lt => Ok(left.lt(right.min())),
688            SQLBinaryOperator::GtEq => Ok(left.gt_eq(right.max())),
689            SQLBinaryOperator::LtEq => Ok(left.lt_eq(right.min())),
690            SQLBinaryOperator::Eq => polars_bail!(SQLSyntax: "ALL cannot be used with ="),
691            SQLBinaryOperator::NotEq => polars_bail!(SQLSyntax: "ALL cannot be used with !="),
692            _ => polars_bail!(SQLInterface: "invalid comparison operator"),
693        }
694    }
695
696    /// Visit a SQL `ANY` expression.
697    ///
698    /// e.g. `a != ANY(y)`
699    fn visit_any(
700        &mut self,
701        left: &SQLExpr,
702        compare_op: &SQLBinaryOperator,
703        right: &SQLExpr,
704    ) -> PolarsResult<Expr> {
705        let left = self.visit_expr(left)?;
706        let right = self.visit_expr(right)?;
707
708        match compare_op {
709            SQLBinaryOperator::Gt => Ok(left.gt(right.min())),
710            SQLBinaryOperator::Lt => Ok(left.lt(right.max())),
711            SQLBinaryOperator::GtEq => Ok(left.gt_eq(right.min())),
712            SQLBinaryOperator::LtEq => Ok(left.lt_eq(right.max())),
713            SQLBinaryOperator::Eq => Ok(left.is_in(right, false)),
714            SQLBinaryOperator::NotEq => Ok(left.is_in(right, false).not()),
715            _ => polars_bail!(SQLInterface: "invalid comparison operator"),
716        }
717    }
718
719    /// Visit a SQL `ARRAY` list (including `IN` values).
720    fn visit_array_expr(
721        &mut self,
722        elements: &[SQLExpr],
723        result_as_element: bool,
724        dtype_expr_match: Option<&Expr>,
725    ) -> PolarsResult<Expr> {
726        let mut elems = self.array_expr_to_series(elements)?;
727
728        // handle implicit temporal strings, eg: "dt IN ('2024-04-30','2024-05-01')".
729        // (not yet as versatile as the temporal string conversions in visit_binary_op)
730        if let (Some(Expr::Column(name)), Some(schema)) =
731            (dtype_expr_match, self.active_schema.as_ref())
732        {
733            if elems.dtype() == &DataType::String {
734                if let Some(dtype) = schema.get(name) {
735                    if matches!(
736                        dtype,
737                        DataType::Date | DataType::Time | DataType::Datetime(_, _)
738                    ) {
739                        elems = elems.strict_cast(dtype)?;
740                    }
741                }
742            }
743        }
744
745        // if we are parsing the list as an element in a series, implode.
746        // otherwise, return the series as-is.
747        let res = if result_as_element {
748            elems.implode()?.into_series()
749        } else {
750            elems
751        };
752        Ok(lit(res))
753    }
754
755    /// Visit a SQL `CAST` or `TRY_CAST` expression.
756    ///
757    /// e.g. `CAST(col AS INT)`, `col::int4`, or `TRY_CAST(col AS VARCHAR)`,
758    fn visit_cast(
759        &mut self,
760        expr: &SQLExpr,
761        dtype: &SQLDataType,
762        format: &Option<CastFormat>,
763        cast_kind: &CastKind,
764    ) -> PolarsResult<Expr> {
765        if format.is_some() {
766            return Err(
767                polars_err!(SQLInterface: "use of FORMAT is not currently supported in CAST"),
768            );
769        }
770        let expr = self.visit_expr(expr)?;
771
772        #[cfg(feature = "json")]
773        if dtype == &SQLDataType::JSON {
774            return Ok(expr.str().json_decode(None, None));
775        }
776        let polars_type = map_sql_dtype_to_polars(dtype)?;
777        Ok(match cast_kind {
778            CastKind::Cast | CastKind::DoubleColon => expr.strict_cast(polars_type),
779            CastKind::TryCast | CastKind::SafeCast => expr.cast(polars_type),
780        })
781    }
782
783    /// Visit a SQL literal.
784    ///
785    /// e.g. 1, 'foo', 1.0, NULL
786    ///
787    /// See [SQLValue] and [LiteralValue] for more details
788    fn visit_literal(&self, value: &SQLValue) -> PolarsResult<Expr> {
789        // note: double-quoted strings will be parsed as identifiers, not literals
790        Ok(match value {
791            SQLValue::Boolean(b) => lit(*b),
792            SQLValue::DollarQuotedString(s) => lit(s.value.clone()),
793            #[cfg(feature = "binary_encoding")]
794            SQLValue::HexStringLiteral(x) => {
795                if x.len() % 2 != 0 {
796                    polars_bail!(SQLSyntax: "hex string literal must have an even number of digits; found '{}'", x)
797                };
798                lit(hex::decode(x.clone()).unwrap())
799            },
800            SQLValue::Null => Expr::Literal(LiteralValue::untyped_null()),
801            SQLValue::Number(s, _) => {
802                // Check for existence of decimal separator dot
803                if s.contains('.') {
804                    s.parse::<f64>().map(lit).map_err(|_| ())
805                } else {
806                    s.parse::<i64>().map(lit).map_err(|_| ())
807                }
808                .map_err(|_| polars_err!(SQLInterface: "cannot parse literal: {:?}", s))?
809            },
810            SQLValue::SingleQuotedByteStringLiteral(b) => {
811                // note: for PostgreSQL this represents a BIT string literal (eg: b'10101') not a BYTE string
812                // literal (see https://www.postgresql.org/docs/current/datatype-bit.html), but sqlparser-rs
813                // patterned the token name after BigQuery (where b'str' really IS a byte string)
814                bitstring_to_bytes_literal(b)?
815            },
816            SQLValue::SingleQuotedString(s) => lit(s.clone()),
817            other => {
818                polars_bail!(SQLInterface: "value {:?} is not a supported literal type", other)
819            },
820        })
821    }
822
823    /// Visit a SQL literal (like [visit_literal]), but return AnyValue instead of Expr.
824    fn visit_any_value(
825        &self,
826        value: &SQLValue,
827        op: Option<&UnaryOperator>,
828    ) -> PolarsResult<AnyValue> {
829        Ok(match value {
830            SQLValue::Boolean(b) => AnyValue::Boolean(*b),
831            SQLValue::DollarQuotedString(s) => AnyValue::StringOwned(s.clone().value.into()),
832            #[cfg(feature = "binary_encoding")]
833            SQLValue::HexStringLiteral(x) => {
834                if x.len() % 2 != 0 {
835                    polars_bail!(SQLSyntax: "hex string literal must have an even number of digits; found '{}'", x)
836                };
837                AnyValue::BinaryOwned(hex::decode(x.clone()).unwrap())
838            },
839            SQLValue::Null => AnyValue::Null,
840            SQLValue::Number(s, _) => {
841                let negate = match op {
842                    Some(UnaryOperator::Minus) => true,
843                    // no op should be taken as plus.
844                    Some(UnaryOperator::Plus) | None => false,
845                    Some(op) => {
846                        polars_bail!(SQLInterface: "unary op {:?} not supported for numeric SQL value", op)
847                    },
848                };
849                // Check for existence of decimal separator dot
850                if s.contains('.') {
851                    s.parse::<f64>()
852                        .map(|n: f64| AnyValue::Float64(if negate { -n } else { n }))
853                        .map_err(|_| ())
854                } else {
855                    s.parse::<i64>()
856                        .map(|n: i64| AnyValue::Int64(if negate { -n } else { n }))
857                        .map_err(|_| ())
858                }
859                .map_err(|_| polars_err!(SQLInterface: "cannot parse literal: {:?}", s))?
860            },
861            SQLValue::SingleQuotedByteStringLiteral(b) => {
862                // note: for PostgreSQL this represents a BIT literal (eg: b'10101') not BYTE
863                let bytes_literal = bitstring_to_bytes_literal(b)?;
864                match bytes_literal {
865                    Expr::Literal(lv) if lv.extract_binary().is_some() => {
866                        AnyValue::BinaryOwned(lv.extract_binary().unwrap().to_vec())
867                    },
868                    _ => {
869                        polars_bail!(SQLInterface: "failed to parse bitstring literal: {:?}", b)
870                    },
871                }
872            },
873            SQLValue::SingleQuotedString(s) => AnyValue::StringOwned(s.as_str().into()),
874            other => polars_bail!(SQLInterface: "value {:?} is not currently supported", other),
875        })
876    }
877
878    /// Visit a SQL `BETWEEN` expression.
879    /// See [sqlparser::ast::Expr::Between] for more details
880    fn visit_between(
881        &mut self,
882        expr: &SQLExpr,
883        negated: bool,
884        low: &SQLExpr,
885        high: &SQLExpr,
886    ) -> PolarsResult<Expr> {
887        let expr = self.visit_expr(expr)?;
888        let low = self.visit_expr(low)?;
889        let high = self.visit_expr(high)?;
890
891        let low = self.convert_temporal_strings(&expr, &low);
892        let high = self.convert_temporal_strings(&expr, &high);
893        Ok(if negated {
894            expr.clone().lt(low).or(expr.gt(high))
895        } else {
896            expr.clone().gt_eq(low).and(expr.lt_eq(high))
897        })
898    }
899
900    /// Visit a SQL `TRIM` function.
901    /// See [sqlparser::ast::Expr::Trim] for more details
902    fn visit_trim(
903        &mut self,
904        expr: &SQLExpr,
905        trim_where: &Option<TrimWhereField>,
906        trim_what: &Option<Box<SQLExpr>>,
907        trim_characters: &Option<Vec<SQLExpr>>,
908    ) -> PolarsResult<Expr> {
909        if trim_characters.is_some() {
910            // TODO: allow compact snowflake/bigquery syntax?
911            return Err(polars_err!(SQLSyntax: "unsupported TRIM syntax (custom chars)"));
912        };
913        let expr = self.visit_expr(expr)?;
914        let trim_what = trim_what.as_ref().map(|e| self.visit_expr(e)).transpose()?;
915        let trim_what = match trim_what {
916            Some(Expr::Literal(lv)) if lv.extract_str().is_some() => {
917                Some(PlSmallStr::from_str(lv.extract_str().unwrap()))
918            },
919            None => None,
920            _ => return self.err(&expr),
921        };
922        Ok(match (trim_where, trim_what) {
923            (None | Some(TrimWhereField::Both), None) => {
924                expr.str().strip_chars(lit(LiteralValue::untyped_null()))
925            },
926            (None | Some(TrimWhereField::Both), Some(val)) => expr.str().strip_chars(lit(val)),
927            (Some(TrimWhereField::Leading), None) => expr
928                .str()
929                .strip_chars_start(lit(LiteralValue::untyped_null())),
930            (Some(TrimWhereField::Leading), Some(val)) => expr.str().strip_chars_start(lit(val)),
931            (Some(TrimWhereField::Trailing), None) => expr
932                .str()
933                .strip_chars_end(lit(LiteralValue::untyped_null())),
934            (Some(TrimWhereField::Trailing), Some(val)) => expr.str().strip_chars_end(lit(val)),
935        })
936    }
937
938    /// Visit a SQL subquery inside and `IN` expression.
939    fn visit_in_subquery(
940        &mut self,
941        expr: &SQLExpr,
942        subquery: &Subquery,
943        negated: bool,
944    ) -> PolarsResult<Expr> {
945        let subquery_result = self.visit_subquery(subquery, SubqueryRestriction::SingleColumn)?;
946        let expr = self.visit_expr(expr)?;
947        Ok(if negated {
948            expr.is_in(subquery_result, false).not()
949        } else {
950            expr.is_in(subquery_result, false)
951        })
952    }
953
954    /// Visit `CASE` control flow expression.
955    fn visit_case_when_then(&mut self, expr: &SQLExpr) -> PolarsResult<Expr> {
956        if let SQLExpr::Case {
957            operand,
958            conditions,
959            results,
960            else_result,
961        } = expr
962        {
963            polars_ensure!(
964                conditions.len() == results.len(),
965                SQLSyntax: "WHEN and THEN expressions must have the same length"
966            );
967            polars_ensure!(
968                !conditions.is_empty(),
969                SQLSyntax: "WHEN and THEN expressions must have at least one element"
970            );
971
972            let mut when_thens = conditions.iter().zip(results.iter());
973            let first = when_thens.next();
974            if first.is_none() {
975                polars_bail!(SQLSyntax: "WHEN and THEN expressions must have at least one element");
976            }
977            let else_res = match else_result {
978                Some(else_res) => self.visit_expr(else_res)?,
979                None => lit(LiteralValue::untyped_null()), // ELSE clause is optional; when omitted, it is implicitly NULL
980            };
981            if let Some(operand_expr) = operand {
982                let first_operand_expr = self.visit_expr(operand_expr)?;
983
984                let first = first.unwrap();
985                let first_cond = first_operand_expr.eq(self.visit_expr(first.0)?);
986                let first_then = self.visit_expr(first.1)?;
987                let expr = when(first_cond).then(first_then);
988                let next = when_thens.next();
989
990                let mut when_then = if let Some((cond, res)) = next {
991                    let second_operand_expr = self.visit_expr(operand_expr)?;
992                    let cond = second_operand_expr.eq(self.visit_expr(cond)?);
993                    let res = self.visit_expr(res)?;
994                    expr.when(cond).then(res)
995                } else {
996                    return Ok(expr.otherwise(else_res));
997                };
998                for (cond, res) in when_thens {
999                    let new_operand_expr = self.visit_expr(operand_expr)?;
1000                    let cond = new_operand_expr.eq(self.visit_expr(cond)?);
1001                    let res = self.visit_expr(res)?;
1002                    when_then = when_then.when(cond).then(res);
1003                }
1004                return Ok(when_then.otherwise(else_res));
1005            }
1006
1007            let first = first.unwrap();
1008            let first_cond = self.visit_expr(first.0)?;
1009            let first_then = self.visit_expr(first.1)?;
1010            let expr = when(first_cond).then(first_then);
1011            let next = when_thens.next();
1012
1013            let mut when_then = if let Some((cond, res)) = next {
1014                let cond = self.visit_expr(cond)?;
1015                let res = self.visit_expr(res)?;
1016                expr.when(cond).then(res)
1017            } else {
1018                return Ok(expr.otherwise(else_res));
1019            };
1020            for (cond, res) in when_thens {
1021                let cond = self.visit_expr(cond)?;
1022                let res = self.visit_expr(res)?;
1023                when_then = when_then.when(cond).then(res);
1024            }
1025            Ok(when_then.otherwise(else_res))
1026        } else {
1027            unreachable!()
1028        }
1029    }
1030
1031    fn err(&self, expr: &Expr) -> PolarsResult<Expr> {
1032        polars_bail!(SQLInterface: "expression {:?} is not currently supported", expr);
1033    }
1034}
1035
1036/// parse a SQL expression to a polars expression
1037/// # Example
1038/// ```rust
1039/// # use polars_sql::{SQLContext, sql_expr};
1040/// # use polars_core::prelude::*;
1041/// # use polars_lazy::prelude::*;
1042/// # fn main() {
1043///
1044/// let mut ctx = SQLContext::new();
1045/// let df = df! {
1046///    "a" =>  [1, 2, 3],
1047/// }
1048/// .unwrap();
1049/// let expr = sql_expr("MAX(a)").unwrap();
1050/// df.lazy().select(vec![expr]).collect().unwrap();
1051/// # }
1052/// ```
1053pub fn sql_expr<S: AsRef<str>>(s: S) -> PolarsResult<Expr> {
1054    let mut ctx = SQLContext::new();
1055
1056    let mut parser = Parser::new(&GenericDialect);
1057    parser = parser.with_options(ParserOptions {
1058        trailing_commas: true,
1059        ..Default::default()
1060    });
1061
1062    let mut ast = parser
1063        .try_with_sql(s.as_ref())
1064        .map_err(to_sql_interface_err)?;
1065    let expr = ast.parse_select_item().map_err(to_sql_interface_err)?;
1066
1067    Ok(match &expr {
1068        SelectItem::ExprWithAlias { expr, alias } => {
1069            let expr = parse_sql_expr(expr, &mut ctx, None)?;
1070            expr.alias(alias.value.as_str())
1071        },
1072        SelectItem::UnnamedExpr(expr) => parse_sql_expr(expr, &mut ctx, None)?,
1073        _ => polars_bail!(SQLInterface: "unable to parse '{}' as Expr", s.as_ref()),
1074    })
1075}
1076
1077pub(crate) fn interval_to_duration(interval: &Interval, fixed: bool) -> PolarsResult<Duration> {
1078    if interval.last_field.is_some()
1079        || interval.leading_field.is_some()
1080        || interval.leading_precision.is_some()
1081        || interval.fractional_seconds_precision.is_some()
1082    {
1083        polars_bail!(SQLSyntax: "unsupported interval syntax ('{}')", interval)
1084    }
1085    let s = match &*interval.value {
1086        SQLExpr::UnaryOp { .. } => {
1087            polars_bail!(SQLSyntax: "unary ops are not valid on interval strings; found {}", interval.value)
1088        },
1089        SQLExpr::Value(SQLValue::SingleQuotedString(s)) => Some(s),
1090        _ => None,
1091    };
1092    match s {
1093        Some(s) if s.contains('-') => {
1094            polars_bail!(SQLInterface: "minus signs are not yet supported in interval strings; found '{}'", s)
1095        },
1096        Some(s) => {
1097            // years, quarters, and months do not have a fixed duration; these
1098            // interval parts can only be used with respect to a reference point
1099            let duration = Duration::parse_interval(s);
1100            if fixed && duration.months() != 0 {
1101                polars_bail!(SQLSyntax: "fixed-duration interval cannot contain years, quarters, or months; found {}", s)
1102            };
1103            Ok(duration)
1104        },
1105        None => polars_bail!(SQLSyntax: "invalid interval {:?}", interval),
1106    }
1107}
1108
1109pub(crate) fn parse_sql_expr(
1110    expr: &SQLExpr,
1111    ctx: &mut SQLContext,
1112    active_schema: Option<&Schema>,
1113) -> PolarsResult<Expr> {
1114    let mut visitor = SQLExprVisitor { ctx, active_schema };
1115    visitor.visit_expr(expr)
1116}
1117
1118pub(crate) fn parse_sql_array(expr: &SQLExpr, ctx: &mut SQLContext) -> PolarsResult<Series> {
1119    match expr {
1120        SQLExpr::Array(arr) => {
1121            let mut visitor = SQLExprVisitor {
1122                ctx,
1123                active_schema: None,
1124            };
1125            visitor.array_expr_to_series(arr.elem.as_slice())
1126        },
1127        _ => polars_bail!(SQLSyntax: "Expected array expression, found {:?}", expr),
1128    }
1129}
1130
1131pub(crate) fn parse_extract_date_part(expr: Expr, field: &DateTimeField) -> PolarsResult<Expr> {
1132    let field = match field {
1133        // handle 'DATE_PART' and all valid abbreviations/alternates
1134        DateTimeField::Custom(Ident { value, .. }) => {
1135            let value = value.to_ascii_lowercase();
1136            match value.as_str() {
1137                "millennium" | "millennia" => &DateTimeField::Millennium,
1138                "century" | "centuries" => &DateTimeField::Century,
1139                "decade" | "decades" => &DateTimeField::Decade,
1140                "isoyear" => &DateTimeField::Isoyear,
1141                "year" | "years" | "y" => &DateTimeField::Year,
1142                "quarter" | "quarters" => &DateTimeField::Quarter,
1143                "month" | "months" | "mon" | "mons" => &DateTimeField::Month,
1144                "dayofyear" | "doy" => &DateTimeField::DayOfYear,
1145                "dayofweek" | "dow" => &DateTimeField::DayOfWeek,
1146                "isoweek" | "week" | "weeks" => &DateTimeField::IsoWeek,
1147                "isodow" => &DateTimeField::Isodow,
1148                "day" | "days" | "d" => &DateTimeField::Day,
1149                "hour" | "hours" | "h" => &DateTimeField::Hour,
1150                "minute" | "minutes" | "mins" | "min" | "m" => &DateTimeField::Minute,
1151                "second" | "seconds" | "sec" | "secs" | "s" => &DateTimeField::Second,
1152                "millisecond" | "milliseconds" | "ms" => &DateTimeField::Millisecond,
1153                "microsecond" | "microseconds" | "us" => &DateTimeField::Microsecond,
1154                "nanosecond" | "nanoseconds" | "ns" => &DateTimeField::Nanosecond,
1155                #[cfg(feature = "timezones")]
1156                "timezone" => &DateTimeField::Timezone,
1157                "time" => &DateTimeField::Time,
1158                "epoch" => &DateTimeField::Epoch,
1159                _ => {
1160                    polars_bail!(SQLSyntax: "EXTRACT/DATE_PART does not support '{}' part", value)
1161                },
1162            }
1163        },
1164        _ => field,
1165    };
1166    Ok(match field {
1167        DateTimeField::Millennium => expr.dt().millennium(),
1168        DateTimeField::Century => expr.dt().century(),
1169        DateTimeField::Decade => expr.dt().year() / typed_lit(10i32),
1170        DateTimeField::Isoyear => expr.dt().iso_year(),
1171        DateTimeField::Year => expr.dt().year(),
1172        DateTimeField::Quarter => expr.dt().quarter(),
1173        DateTimeField::Month => expr.dt().month(),
1174        DateTimeField::Week(weekday) => {
1175            if weekday.is_some() {
1176                polars_bail!(SQLSyntax: "EXTRACT/DATE_PART does not support '{}' part", field)
1177            }
1178            expr.dt().week()
1179        },
1180        DateTimeField::IsoWeek => expr.dt().week(),
1181        DateTimeField::DayOfYear | DateTimeField::Doy => expr.dt().ordinal_day(),
1182        DateTimeField::DayOfWeek | DateTimeField::Dow => {
1183            let w = expr.dt().weekday();
1184            when(w.clone().eq(typed_lit(7i8)))
1185                .then(typed_lit(0i8))
1186                .otherwise(w)
1187        },
1188        DateTimeField::Isodow => expr.dt().weekday(),
1189        DateTimeField::Day => expr.dt().day(),
1190        DateTimeField::Hour => expr.dt().hour(),
1191        DateTimeField::Minute => expr.dt().minute(),
1192        DateTimeField::Second => expr.dt().second(),
1193        DateTimeField::Millisecond | DateTimeField::Milliseconds => {
1194            (expr.clone().dt().second() * typed_lit(1_000f64))
1195                + expr.dt().nanosecond().div(typed_lit(1_000_000f64))
1196        },
1197        DateTimeField::Microsecond | DateTimeField::Microseconds => {
1198            (expr.clone().dt().second() * typed_lit(1_000_000f64))
1199                + expr.dt().nanosecond().div(typed_lit(1_000f64))
1200        },
1201        DateTimeField::Nanosecond | DateTimeField::Nanoseconds => {
1202            (expr.clone().dt().second() * typed_lit(1_000_000_000f64)) + expr.dt().nanosecond()
1203        },
1204        DateTimeField::Time => expr.dt().time(),
1205        #[cfg(feature = "timezones")]
1206        DateTimeField::Timezone => expr.dt().base_utc_offset().dt().total_seconds(),
1207        DateTimeField::Epoch => {
1208            expr.clone()
1209                .dt()
1210                .timestamp(TimeUnit::Nanoseconds)
1211                .div(typed_lit(1_000_000_000i64))
1212                + expr.dt().nanosecond().div(typed_lit(1_000_000_000f64))
1213        },
1214        _ => {
1215            polars_bail!(SQLSyntax: "EXTRACT/DATE_PART does not support '{}' part", field)
1216        },
1217    })
1218}
1219
1220/// Allow an expression that represents a 1-indexed parameter to
1221/// be adjusted from 1-indexed (SQL) to 0-indexed (Rust/Polars)
1222pub(crate) fn adjust_one_indexed_param(idx: Expr, null_if_zero: bool) -> Expr {
1223    match idx {
1224        Expr::Literal(sc) if sc.is_null() => lit(LiteralValue::untyped_null()),
1225        Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Int(0))) => {
1226            if null_if_zero {
1227                lit(LiteralValue::untyped_null())
1228            } else {
1229                idx
1230            }
1231        },
1232        Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Int(n))) if n < 0 => idx,
1233        Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Int(n))) => lit(n - 1),
1234        // TODO: when 'saturating_sub' is available, should be able
1235        //  to streamline the when/then/otherwise block below -
1236        _ => when(idx.clone().gt(lit(0)))
1237            .then(idx.clone() - lit(1))
1238            .otherwise(if null_if_zero {
1239                when(idx.clone().eq(lit(0)))
1240                    .then(lit(LiteralValue::untyped_null()))
1241                    .otherwise(idx.clone())
1242            } else {
1243                idx.clone()
1244            }),
1245    }
1246}
1247
1248fn resolve_column<'a>(
1249    ctx: &'a mut SQLContext,
1250    ident_root: &'a Ident,
1251    name: &'a str,
1252    dtype: &'a DataType,
1253) -> PolarsResult<(Expr, Option<&'a DataType>)> {
1254    let resolved = ctx.resolve_name(&ident_root.value, name);
1255    let resolved = resolved.as_str();
1256    Ok((
1257        if name != resolved {
1258            col(resolved).alias(name)
1259        } else {
1260            col(name)
1261        },
1262        Some(dtype),
1263    ))
1264}
1265
1266pub(crate) fn resolve_compound_identifier(
1267    ctx: &mut SQLContext,
1268    idents: &[Ident],
1269    active_schema: Option<&Schema>,
1270) -> PolarsResult<Vec<Expr>> {
1271    // inference priority: table > struct > column
1272    let ident_root = &idents[0];
1273    let mut remaining_idents = idents.iter().skip(1);
1274    let mut lf = ctx.get_table_from_current_scope(&ident_root.value);
1275
1276    let schema = if let Some(ref mut lf) = lf {
1277        lf.schema_with_arenas(&mut ctx.lp_arena, &mut ctx.expr_arena)
1278    } else {
1279        Ok(Arc::new(if let Some(active_schema) = active_schema {
1280            active_schema.clone()
1281        } else {
1282            Schema::default()
1283        }))
1284    }?;
1285
1286    let col_dtype: PolarsResult<(Expr, Option<&DataType>)> = if lf.is_none() && schema.is_empty() {
1287        Ok((col(ident_root.value.as_str()), None))
1288    } else {
1289        let name = &remaining_idents.next().unwrap().value;
1290        if lf.is_some() && name == "*" {
1291            return Ok(schema
1292                .iter_names_and_dtypes()
1293                .map(|(name, dtype)| resolve_column(ctx, ident_root, name, dtype).unwrap().0)
1294                .collect::<Vec<_>>());
1295        } else if let Some((_, name, dtype)) = schema.get_full(name) {
1296            resolve_column(ctx, ident_root, name, dtype)
1297        } else if lf.is_none() {
1298            remaining_idents = idents.iter().skip(1);
1299            Ok((
1300                col(ident_root.value.as_str()),
1301                schema.get(&ident_root.value),
1302            ))
1303        } else {
1304            polars_bail!(
1305                SQLInterface: "no column named '{}' found in table '{}'",
1306                name,
1307                ident_root
1308            )
1309        }
1310    };
1311
1312    // additional ident levels index into struct fields
1313    let (mut column, mut dtype) = col_dtype?;
1314    for ident in remaining_idents {
1315        let name = ident.value.as_str();
1316        match dtype {
1317            Some(DataType::Struct(fields)) if name == "*" => {
1318                return Ok(fields
1319                    .iter()
1320                    .map(|fld| column.clone().struct_().field_by_name(&fld.name))
1321                    .collect());
1322            },
1323            Some(DataType::Struct(fields)) => {
1324                dtype = fields
1325                    .iter()
1326                    .find(|fld| fld.name == name)
1327                    .map(|fld| &fld.dtype);
1328            },
1329            Some(dtype) if name == "*" => {
1330                polars_bail!(SQLSyntax: "cannot expand '*' on non-Struct dtype; found {:?}", dtype)
1331            },
1332            _ => {
1333                dtype = None;
1334            },
1335        }
1336        column = column.struct_().field_by_name(name);
1337    }
1338    Ok(vec![column])
1339}
polars_sql/sql_expr.rs

polars_sql/
sql_expr.rs
