1use std::ops::Deref;
2use std::sync::RwLock;
3
4use polars_core::frame::row::Row;
5use polars_core::prelude::*;
6use polars_lazy::prelude::*;
7use polars_ops::frame::{JoinCoalesce, MaintainOrderJoin};
8use polars_plan::dsl::function_expr::StructFunction;
9use polars_plan::prelude::*;
10use polars_utils::aliases::{PlHashSet, PlIndexSet};
11use polars_utils::format_pl_smallstr;
12use sqlparser::ast::{
13 BinaryOperator as SQLBinaryOperator, CreateTable, CreateTableLikeKind, CreateTableOptions,
14 Delete, Distinct, ExcludeSelectItem, Expr as SQLExpr, Fetch, FromTable, FunctionArg,
15 GroupByExpr, HiveDistributionStyle, HiveFormat, Ident, JoinConstraint, JoinOperator,
16 LimitClause, NamedWindowDefinition, NamedWindowExpr, ObjectName, ObjectType, OrderBy,
17 OrderByKind, Query, RenameSelectItem, Select, SelectFlavor, SelectItem,
18 SelectItemQualifiedWildcardKind, SetExpr, SetOperator, SetQuantifier, Statement, TableAlias,
19 TableFactor, TableWithJoins, Truncate, UnaryOperator as SQLUnaryOperator, Value as SQLValue,
20 ValueWithSpan, Values, Visit, WildcardAdditionalOptions, WindowSpec,
21};
22use sqlparser::dialect::GenericDialect;
23use sqlparser::parser::{Parser, ParserOptions};
24
25use crate::function_registry::{DefaultFunctionRegistry, FunctionRegistry};
26use crate::sql_expr::{
27 parse_sql_array, parse_sql_expr, resolve_compound_identifier, to_sql_interface_err,
28};
29use crate::sql_visitors::{
30 QualifyExpression, TableIdentifierCollector, check_for_ambiguous_column_refs,
31 expr_has_window_functions, expr_refers_to_table,
32};
33use crate::table_functions::PolarsTableFunctions;
34use crate::types::map_sql_dtype_to_polars;
35
36#[derive(Clone)]
37pub struct TableInfo {
38 pub(crate) frame: LazyFrame,
39 pub(crate) name: PlSmallStr,
40 pub(crate) schema: Arc<Schema>,
41}
42
43struct SelectModifiers {
44 exclude: PlHashSet<String>, ilike: Option<regex::Regex>, rename: PlHashMap<PlSmallStr, PlSmallStr>, replace: Vec<Expr>, }
49impl SelectModifiers {
50 fn matches_ilike(&self, s: &str) -> bool {
51 match &self.ilike {
52 Some(rx) => rx.is_match(s),
53 None => true,
54 }
55 }
56 fn renamed_cols(&self) -> Vec<Expr> {
57 self.rename
58 .iter()
59 .map(|(before, after)| col(before.clone()).alias(after.clone()))
60 .collect()
61 }
62}
63
64enum ProjectionItem {
66 QualifiedExprs(PlSmallStr, Vec<Expr>),
67 Exprs(Vec<Expr>),
68}
69
70fn expr_output_name(expr: &Expr) -> Option<&PlSmallStr> {
72 match expr {
73 Expr::Column(name) | Expr::Alias(_, name) => Some(name),
74 _ => None,
75 }
76}
77
78fn disambiguate_projection_cols(
80 items: Vec<ProjectionItem>,
81 schema: &Schema,
82) -> PolarsResult<Vec<Expr>> {
83 let mut qualified_wildcard_names: PlHashMap<PlSmallStr, usize> = PlHashMap::new();
85 let mut other_names: PlHashSet<PlSmallStr> = PlHashSet::new();
86 for item in &items {
87 match item {
88 ProjectionItem::QualifiedExprs(_, exprs) => {
89 for expr in exprs {
90 if let Some(name) = expr_output_name(expr) {
91 *qualified_wildcard_names.entry(name.clone()).or_insert(0) += 1;
92 }
93 }
94 },
95 ProjectionItem::Exprs(exprs) => {
96 for expr in exprs {
97 if let Some(name) = expr_output_name(expr) {
98 other_names.insert(name.clone());
99 }
100 }
101 },
102 }
103 }
104
105 let needs_suffix: PlHashSet<PlSmallStr> = qualified_wildcard_names
107 .into_iter()
108 .filter(|(name, count)| *count > 1 || other_names.contains(name))
109 .map(|(name, _)| name)
110 .collect();
111
112 let mut result: Vec<Expr> = Vec::new();
114 for item in items {
115 match item {
116 ProjectionItem::QualifiedExprs(tbl_name, exprs) if !needs_suffix.is_empty() => {
117 for expr in exprs {
118 if let Some(name) = expr_output_name(&expr) {
119 if needs_suffix.contains(name) {
120 let suffixed = format_pl_smallstr!("{}:{}", name, tbl_name);
121 if schema.contains(suffixed.as_str()) {
122 result.push(col(suffixed));
123 continue;
124 }
125 if other_names.contains(name) {
126 polars_bail!(
127 SQLInterface:
128 "column '{}' is duplicated in the SELECT (explicitly, and via the `*` wildcard)", name
129 );
130 }
131 }
132 }
133 result.push(expr);
134 }
135 },
136 ProjectionItem::QualifiedExprs(_, exprs) | ProjectionItem::Exprs(exprs) => {
137 result.extend(exprs);
138 },
139 }
140 }
141 Ok(result)
142}
143
144#[derive(Clone, Copy, PartialEq, Eq)]
146pub(crate) enum FilterMode {
147 KeepTrue,
150 RemoveTrue,
153}
154
155#[derive(Clone)]
157pub struct SQLContext {
158 pub(crate) table_map: Arc<RwLock<PlHashMap<String, LazyFrame>>>,
159 pub(crate) function_registry: Arc<dyn FunctionRegistry>,
160 pub(crate) lp_arena: Arena<IR>,
161 pub(crate) expr_arena: Arena<AExpr>,
162
163 cte_map: PlHashMap<String, LazyFrame>,
164 table_aliases: PlHashMap<String, String>,
165 joined_aliases: PlHashMap<String, PlHashMap<String, String>>,
166 pub(crate) named_windows: PlHashMap<String, WindowSpec>,
167}
168
169impl Default for SQLContext {
170 fn default() -> Self {
171 Self {
172 function_registry: Arc::new(DefaultFunctionRegistry {}),
173 table_map: Default::default(),
174 cte_map: Default::default(),
175 table_aliases: Default::default(),
176 joined_aliases: Default::default(),
177 named_windows: Default::default(),
178 lp_arena: Default::default(),
179 expr_arena: Default::default(),
180 }
181 }
182}
183
184impl SQLContext {
185 pub fn new() -> Self {
193 Self::default()
194 }
195
196 pub fn get_tables(&self) -> Vec<String> {
198 let mut tables = Vec::from_iter(self.table_map.read().unwrap().keys().cloned());
199 tables.sort_unstable();
200 tables
201 }
202
203 pub fn register(&self, name: &str, lf: LazyFrame) {
219 self.table_map.write().unwrap().insert(name.to_owned(), lf);
220 }
221
222 pub fn unregister(&self, name: &str) {
224 self.table_map.write().unwrap().remove(&name.to_owned());
225 }
226
227 pub fn execute(&mut self, query: &str) -> PolarsResult<LazyFrame> {
246 let mut parser = Parser::new(&GenericDialect);
247 parser = parser.with_options(ParserOptions {
248 trailing_commas: true,
249 ..Default::default()
250 });
251
252 let ast = parser
253 .try_with_sql(query)
254 .map_err(to_sql_interface_err)?
255 .parse_statements()
256 .map_err(to_sql_interface_err)?;
257
258 polars_ensure!(ast.len() == 1, SQLInterface: "one (and only one) statement can be parsed at a time");
259 let res = self.execute_statement(ast.first().unwrap())?;
260
261 let lp_arena = std::mem::take(&mut self.lp_arena);
264 let expr_arena = std::mem::take(&mut self.expr_arena);
265 res.set_cached_arena(lp_arena, expr_arena);
266
267 self.cte_map.clear();
269 self.table_aliases.clear();
270 self.joined_aliases.clear();
271 self.named_windows.clear();
272
273 Ok(res)
274 }
275
276 pub fn with_function_registry(mut self, function_registry: Arc<dyn FunctionRegistry>) -> Self {
279 self.function_registry = function_registry;
280 self
281 }
282
283 pub fn registry(&self) -> &Arc<dyn FunctionRegistry> {
285 &self.function_registry
286 }
287
288 pub fn registry_mut(&mut self) -> &mut dyn FunctionRegistry {
290 Arc::get_mut(&mut self.function_registry).unwrap()
291 }
292}
293
294impl SQLContext {
295 pub(crate) fn isolated(&self) -> Self {
296 Self {
297 table_map: Arc::new(RwLock::new(self.table_map.read().unwrap().clone())),
299 named_windows: self.named_windows.clone(),
300 cte_map: self.cte_map.clone(),
301
302 ..Default::default()
303 }
304 }
305
306 pub(crate) fn execute_statement(&mut self, stmt: &Statement) -> PolarsResult<LazyFrame> {
307 let ast = stmt;
308 Ok(match ast {
309 Statement::Query(query) => self.execute_query(query)?,
310 stmt @ Statement::ShowTables { .. } => self.execute_show_tables(stmt)?,
311 stmt @ Statement::CreateTable { .. } => self.execute_create_table(stmt)?,
312 stmt @ Statement::Drop {
313 object_type: ObjectType::Table,
314 ..
315 } => self.execute_drop_table(stmt)?,
316 stmt @ Statement::Explain { .. } => self.execute_explain(stmt)?,
317 stmt @ Statement::Truncate { .. } => self.execute_truncate_table(stmt)?,
318 stmt @ Statement::Delete { .. } => self.execute_delete_from_table(stmt)?,
319 _ => polars_bail!(
320 SQLInterface: "statement type is not supported:\n{:?}", ast,
321 ),
322 })
323 }
324
325 pub(crate) fn execute_query(&mut self, query: &Query) -> PolarsResult<LazyFrame> {
326 self.register_ctes(query)?;
327 self.execute_query_no_ctes(query)
328 }
329
330 pub(crate) fn execute_query_no_ctes(&mut self, query: &Query) -> PolarsResult<LazyFrame> {
331 self.validate_query(query)?;
332
333 let lf = self.process_query(&query.body, query)?;
334 self.process_limit_offset(lf, &query.limit_clause, &query.fetch)
335 }
336
337 pub(crate) fn get_frame_schema(&mut self, frame: &mut LazyFrame) -> PolarsResult<SchemaRef> {
338 frame.schema_with_arenas(&mut self.lp_arena, &mut self.expr_arena)
339 }
340
341 pub(super) fn get_table_from_current_scope(&self, name: &str) -> Option<LazyFrame> {
342 self.table_map
346 .read()
347 .unwrap()
348 .get(name)
349 .cloned()
350 .or_else(|| self.cte_map.get(name).cloned())
351 .or_else(|| {
352 self.table_aliases.get(name).and_then(|alias| {
353 self.table_map
354 .read()
355 .unwrap()
356 .get(alias.as_str())
357 .or_else(|| self.cte_map.get(alias.as_str()))
358 .cloned()
359 })
360 })
361 }
362
363 pub(crate) fn execute_isolated<F>(&mut self, query: F) -> PolarsResult<LazyFrame>
367 where
368 F: FnOnce(&mut Self) -> PolarsResult<LazyFrame>,
369 {
370 let mut ctx = self.isolated();
371
372 let lf = query(&mut ctx)?;
374
375 lf.set_cached_arena(ctx.lp_arena, ctx.expr_arena);
377
378 Ok(lf)
379 }
380
381 fn expr_or_ordinal(
382 &mut self,
383 e: &SQLExpr,
384 exprs: &[Expr],
385 selected: Option<&[Expr]>,
386 schema: Option<&Schema>,
387 clause: &str,
388 ) -> PolarsResult<Expr> {
389 match e {
390 SQLExpr::UnaryOp {
391 op: SQLUnaryOperator::Minus,
392 expr,
393 } if matches!(
394 **expr,
395 SQLExpr::Value(ValueWithSpan {
396 value: SQLValue::Number(_, _),
397 ..
398 })
399 ) =>
400 {
401 if let SQLExpr::Value(ValueWithSpan {
402 value: SQLValue::Number(ref idx, _),
403 ..
404 }) = **expr
405 {
406 Err(polars_err!(
407 SQLSyntax:
408 "negative ordinal values are invalid for {}; found -{}",
409 clause,
410 idx
411 ))
412 } else {
413 unreachable!()
414 }
415 },
416 SQLExpr::Value(ValueWithSpan {
417 value: SQLValue::Number(idx, _),
418 ..
419 }) => {
420 let idx = idx.parse::<usize>().map_err(|_| {
422 polars_err!(
423 SQLSyntax:
424 "negative ordinal values are invalid for {}; found {}",
425 clause,
426 idx
427 )
428 })?;
429 let cols = if let Some(cols) = selected {
432 cols
433 } else {
434 exprs
435 };
436 Ok(cols
437 .get(idx - 1)
438 .ok_or_else(|| {
439 polars_err!(
440 SQLInterface:
441 "{} ordinal value must refer to a valid column; found {}",
442 clause,
443 idx
444 )
445 })?
446 .clone())
447 },
448 SQLExpr::Value(v) => Err(polars_err!(
449 SQLSyntax:
450 "{} requires a valid expression or positive ordinal; found {}", clause, v,
451 )),
452 _ => {
453 let mut expr = parse_sql_expr(e, self, schema)?;
456 if matches!(e, SQLExpr::CompoundIdentifier(_)) {
457 if let Some(schema) = schema {
458 expr = expr.map_expr(|ex| match &ex {
459 Expr::Column(name) => {
460 let prefixed = format!("__POLARS_ORIG_{}", name.as_str());
461 if schema.contains(prefixed.as_str()) {
462 col(prefixed)
463 } else {
464 ex
465 }
466 },
467 _ => ex,
468 });
469 }
470 }
471 Ok(expr)
472 },
473 }
474 }
475
476 pub(super) fn resolve_name(&self, tbl_name: &str, column_name: &str) -> String {
477 if let Some(aliases) = self.joined_aliases.get(tbl_name) {
478 if let Some(name) = aliases.get(column_name) {
479 return name.to_string();
480 }
481 }
482 column_name.to_string()
483 }
484
485 fn process_query(&mut self, expr: &SetExpr, query: &Query) -> PolarsResult<LazyFrame> {
486 match expr {
487 SetExpr::Select(select_stmt) => self.execute_select(select_stmt, query),
488 SetExpr::Query(nested_query) => {
489 let lf = self.execute_query_no_ctes(nested_query)?;
490 self.process_order_by(lf, &query.order_by, None)
491 },
492 SetExpr::SetOperation {
493 op: SetOperator::Union,
494 set_quantifier,
495 left,
496 right,
497 } => self.process_union(left, right, set_quantifier, query),
498
499 #[cfg(feature = "semi_anti_join")]
500 SetExpr::SetOperation {
501 op: SetOperator::Intersect | SetOperator::Except,
502 set_quantifier,
503 left,
504 right,
505 } => self.process_except_intersect(left, right, set_quantifier, query),
506
507 SetExpr::Values(Values {
508 explicit_row: _,
509 rows,
510 value_keyword: _,
511 }) => self.process_values(rows.iter().map(|p| &p.content)),
512
513 SetExpr::Table(tbl) => {
514 if let Some(table_name) = tbl.table_name.as_ref() {
515 self.get_table_from_current_scope(table_name)
516 .ok_or_else(|| {
517 polars_err!(
518 SQLInterface: "no table or alias named '{}' found",
519 tbl
520 )
521 })
522 } else {
523 polars_bail!(SQLInterface: "'TABLE' requires valid table name")
524 }
525 },
526 op => {
527 polars_bail!(SQLInterface: "'{}' operation is currently unsupported", op)
528 },
529 }
530 }
531
532 #[cfg(feature = "semi_anti_join")]
533 fn process_except_intersect(
534 &mut self,
535 left: &SetExpr,
536 right: &SetExpr,
537 quantifier: &SetQuantifier,
538 query: &Query,
539 ) -> PolarsResult<LazyFrame> {
540 let (join_type, op_name) = match *query.body {
541 SetExpr::SetOperation {
542 op: SetOperator::Except,
543 ..
544 } => (JoinType::Anti, "EXCEPT"),
545 _ => (JoinType::Semi, "INTERSECT"),
546 };
547
548 let mut lf = self.execute_isolated(|ctx| ctx.process_query(left, query))?;
551 let mut rf = self.execute_isolated(|ctx| ctx.process_query(right, query))?;
552 let lf_schema = self.get_frame_schema(&mut lf)?;
553
554 let lf_cols: Vec<_> = lf_schema.iter_names_cloned().map(col).collect();
555 let rf_cols = match quantifier {
556 SetQuantifier::ByName => None,
557 SetQuantifier::Distinct | SetQuantifier::None => {
558 let rf_schema = self.get_frame_schema(&mut rf)?;
559 let rf_cols: Vec<_> = rf_schema.iter_names_cloned().map(col).collect();
560 if lf_cols.len() != rf_cols.len() {
561 polars_bail!(SQLInterface: "{} requires equal number of columns in each table (use '{} BY NAME' to combine mismatched tables)", op_name, op_name)
562 }
563 Some(rf_cols)
564 },
565 _ => {
566 polars_bail!(SQLInterface: "'{} {}' is not supported", op_name, quantifier.to_string())
567 },
568 };
569 let join = lf.join_builder().with(rf).how(join_type).join_nulls(true);
570 let joined_tbl = match rf_cols {
571 Some(rf_cols) => join.left_on(lf_cols).right_on(rf_cols).finish(),
572 None => join.on(lf_cols).finish(),
573 };
574 let lf = joined_tbl.unique(None, UniqueKeepStrategy::Any);
575 self.process_order_by(lf, &query.order_by, None)
576 }
577
578 fn process_union(
579 &mut self,
580 left: &SetExpr,
581 right: &SetExpr,
582 quantifier: &SetQuantifier,
583 query: &Query,
584 ) -> PolarsResult<LazyFrame> {
585 let quantifier = *quantifier;
586
587 let mut lf = self.execute_isolated(|ctx| ctx.process_query(left, query))?;
590 let mut rf = self.execute_isolated(|ctx| ctx.process_query(right, query))?;
591
592 let opts = UnionArgs {
593 parallel: true,
594 to_supertypes: true,
595 maintain_order: false,
596 ..Default::default()
597 };
598 let lf = match quantifier {
599 SetQuantifier::All | SetQuantifier::Distinct | SetQuantifier::None => {
601 let lf_schema = self.get_frame_schema(&mut lf)?;
602 let rf_schema = self.get_frame_schema(&mut rf)?;
603 if lf_schema.len() != rf_schema.len() {
604 polars_bail!(SQLInterface: "UNION requires equal number of columns in each table (use 'UNION BY NAME' to combine mismatched tables)")
605 }
606 if lf_schema.iter_names().ne(rf_schema.iter_names()) {
609 rf = rf.rename(rf_schema.iter_names(), lf_schema.iter_names(), true);
610 }
611 let concatenated = concat(vec![lf, rf], opts);
612 match quantifier {
613 SetQuantifier::Distinct | SetQuantifier::None => {
614 concatenated.map(|lf| lf.unique(None, UniqueKeepStrategy::Any))
615 },
616 _ => concatenated,
617 }
618 },
619 #[cfg(feature = "diagonal_concat")]
621 SetQuantifier::AllByName => concat_lf_diagonal(vec![lf, rf], opts),
622 #[cfg(feature = "diagonal_concat")]
624 SetQuantifier::ByName | SetQuantifier::DistinctByName => {
625 let concatenated = concat_lf_diagonal(vec![lf, rf], opts);
626 concatenated.map(|lf| lf.unique(None, UniqueKeepStrategy::Any))
627 },
628 #[allow(unreachable_patterns)]
629 _ => {
630 polars_bail!(SQLInterface: "'UNION {}' is not currently supported", quantifier)
631 },
632 }?;
633
634 self.process_order_by(lf, &query.order_by, None)
635 }
636
637 fn process_unnest_lateral(
640 &self,
641 lf: LazyFrame,
642 alias: &Option<TableAlias>,
643 array_exprs: &[SQLExpr],
644 with_offset: bool,
645 ) -> PolarsResult<LazyFrame> {
646 let alias = alias
647 .as_ref()
648 .ok_or_else(|| polars_err!(SQLSyntax: "UNNEST table must have an alias"))?;
649 polars_ensure!(!with_offset, SQLInterface: "UNNEST tables do not (yet) support WITH ORDINALITY|OFFSET");
650
651 let (mut explode_cols, mut rename_from, mut rename_to) = (
652 Vec::with_capacity(array_exprs.len()),
653 Vec::with_capacity(array_exprs.len()),
654 Vec::with_capacity(array_exprs.len()),
655 );
656 let is_single_col = array_exprs.len() == 1;
657
658 for (i, arr_expr) in array_exprs.iter().enumerate() {
659 let col_name = match arr_expr {
660 SQLExpr::Identifier(ident) => PlSmallStr::from_str(&ident.value),
661 SQLExpr::CompoundIdentifier(parts) => {
662 PlSmallStr::from_str(&parts.last().unwrap().value)
663 },
664 SQLExpr::Array(_) => polars_bail!(
665 SQLInterface: "CROSS JOIN UNNEST with both literal arrays and column references is not supported"
666 ),
667 other => polars_bail!(
668 SQLSyntax: "UNNEST expects column references or array literals, found {:?}", other
669 ),
670 };
671 if let Some(name) = alias
673 .columns
674 .get(i)
675 .map(|c| c.name.value.as_str())
676 .or_else(|| is_single_col.then_some(alias.name.value.as_str()))
677 .filter(|name| !name.is_empty() && *name != col_name.as_str())
678 {
679 rename_from.push(col_name.clone());
680 rename_to.push(PlSmallStr::from_str(name));
681 }
682 explode_cols.push(col_name);
683 }
684
685 let mut lf = lf.explode(
686 Selector::ByName {
687 names: Arc::from(explode_cols),
688 strict: true,
689 },
690 ExplodeOptions {
691 empty_as_null: true,
692 keep_nulls: true,
693 },
694 );
695 if !rename_from.is_empty() {
696 lf = lf.rename(rename_from, rename_to, true);
697 }
698 Ok(lf)
699 }
700
701 fn process_values<'a>(
702 &mut self,
703 values: impl Iterator<Item = &'a Vec<SQLExpr>>,
704 ) -> PolarsResult<LazyFrame> {
705 let frame_rows: Vec<Row> = values.map(|row| {
706 let row_data: Result<Vec<_>, _> = row.iter().map(|expr| {
707 let expr = parse_sql_expr(expr, self, None)?;
708 match expr {
709 Expr::Literal(value) => {
710 value.to_any_value()
711 .ok_or_else(|| polars_err!(SQLInterface: "invalid literal value: {:?}", value))
712 .map(|av| av.into_static())
713 },
714 _ => polars_bail!(SQLInterface: "VALUES clause expects literals; found {}", expr),
715 }
716 }).collect();
717 row_data.map(Row::new)
718 }).collect::<Result<_, _>>()?;
719
720 Ok(DataFrame::from_rows(frame_rows.as_ref())?.lazy())
721 }
722
723 fn execute_explain(&mut self, stmt: &Statement) -> PolarsResult<LazyFrame> {
725 match stmt {
726 Statement::Explain { statement, .. } => {
727 let lf = self.execute_statement(statement)?;
728 let plan = lf.describe_optimized_plan()?;
729 let plan = plan
730 .split('\n')
731 .collect::<Series>()
732 .with_name(PlSmallStr::from_static("Logical Plan"))
733 .into_column();
734 let df = DataFrame::new_infer_height(vec![plan])?;
735 Ok(df.lazy())
736 },
737 _ => polars_bail!(SQLInterface: "unexpected statement type; expected EXPLAIN"),
738 }
739 }
740
741 fn execute_show_tables(&mut self, _: &Statement) -> PolarsResult<LazyFrame> {
743 let tables = Column::new("name".into(), self.get_tables());
744 let df = DataFrame::new_infer_height(vec![tables])?;
745 Ok(df.lazy())
746 }
747
748 fn execute_drop_table(&mut self, stmt: &Statement) -> PolarsResult<LazyFrame> {
750 if let Statement::Drop {
752 object_type: _,
753 names,
754 if_exists,
755
756 cascade,
758 restrict,
759 purge,
760 temporary,
761 table,
762 } = stmt
763 {
764 polars_ensure!(!cascade, SQLInterface: "`DROP ... CASCADE` is not supported");
765 polars_ensure!(!purge, SQLInterface: "`DROP ... PURGE` is not supported");
766 polars_ensure!(!restrict, SQLInterface: "`DROP ... RESTRICT` is not supported");
767 polars_ensure!(!temporary, SQLInterface: "`DROP TEMPORARY` is not supported");
768 polars_ensure!(table.is_none(), SQLInterface: "`DROP ... ON <table>` is not supported");
769
770 for name in names {
771 let tbl = name.to_string();
772 if self.table_map.write().unwrap().remove(&tbl).is_none() && !if_exists {
775 polars_bail!(SQLInterface: "table '{}' does not exist", tbl);
776 }
777 }
778 Ok(DataFrame::empty().lazy())
779 } else {
780 polars_bail!(SQLInterface: "unexpected statement type; expected DROP")
781 }
782 }
783
784 fn execute_delete_from_table(&mut self, stmt: &Statement) -> PolarsResult<LazyFrame> {
786 if let Statement::Delete(Delete {
787 tables,
788 from,
789 using,
790 selection,
791 returning,
792 order_by,
793 limit,
794 delete_token: _,
795 optimizer_hints: _,
796 output: _,
797 }) = stmt
798 {
799 let error_message: Option<&'static str> = if !tables.is_empty() {
800 Some("DELETE expects exactly one table name")
801 } else if using.is_some() {
802 Some("DELETE does not support the USING clause")
803 } else if returning.is_some() {
804 Some("DELETE does not support the RETURNING clause")
805 } else if limit.is_some() {
806 Some("DELETE does not support the LIMIT clause")
807 } else if !order_by.is_empty() {
808 Some("DELETE does not support the ORDER BY clause")
809 } else {
810 None
811 };
812
813 if let Some(msg) = error_message {
814 polars_bail!(SQLInterface: msg);
815 }
816
817 let from_tables = match &from {
818 FromTable::WithFromKeyword(from) => from,
819 FromTable::WithoutKeyword(from) => from,
820 };
821 if from_tables.len() > 1 {
822 polars_bail!(SQLInterface: "cannot have multiple tables in DELETE FROM (found {})", from_tables.len())
823 }
824 let tbl_expr = from_tables.first().unwrap();
825 if !tbl_expr.joins.is_empty() {
826 polars_bail!(SQLInterface: "DELETE does not support table JOINs")
827 }
828 let (_, lf) = self.get_table(&tbl_expr.relation)?;
829 if selection.is_none() {
830 Ok(lf.clear())
832 } else {
833 Ok(self.process_where(lf.clone(), selection, FilterMode::RemoveTrue, None)?)
835 }
836 } else {
837 polars_bail!(SQLInterface: "unexpected statement type; expected DELETE")
838 }
839 }
840
841 fn execute_truncate_table(&mut self, stmt: &Statement) -> PolarsResult<LazyFrame> {
843 if let Statement::Truncate(Truncate {
844 table_names,
845 partitions,
846 table: _, if_exists,
848 identity,
849 cascade,
850 on_cluster,
851 }) = stmt
852 {
853 polars_ensure!(identity.is_none(), SQLInterface: "`TRUNCATE ... RESTART/CONTINUE IDENTITY` is not supported");
854 polars_ensure!(cascade.is_none(), SQLInterface: "`TRUNCATE ... CASCADE/RESTRICT` is not supported");
855 polars_ensure!(on_cluster.is_none(), SQLInterface: "`TRUNCATE ... ON CLUSTER` is not supported");
856
857 match partitions {
858 None => {
859 if table_names.len() != 1 {
860 polars_bail!(SQLInterface: "TRUNCATE expects exactly one table name; found {}", table_names.len())
861 }
862 let tbl = table_names[0].name.to_string();
863 if let Some(lf) = self.table_map.write().unwrap().get_mut(&tbl) {
864 *lf = lf.clone().clear();
865 Ok(lf.clone())
866 } else if *if_exists {
867 Ok(DataFrame::empty().lazy())
869 } else {
870 polars_bail!(SQLInterface: "table '{}' does not exist", tbl);
871 }
872 },
873 _ => {
874 polars_bail!(SQLInterface: "TRUNCATE does not support use of 'partitions'")
875 },
876 }
877 } else {
878 polars_bail!(SQLInterface: "unexpected statement type; expected TRUNCATE")
879 }
880 }
881
882 fn register_cte(&mut self, name: &str, lf: LazyFrame) {
883 self.cte_map.insert(name.to_owned(), lf);
884 }
885
886 fn register_ctes(&mut self, query: &Query) -> PolarsResult<()> {
887 if let Some(with) = &query.with {
888 if with.recursive {
889 polars_bail!(SQLInterface: "recursive CTEs are not supported")
890 }
891 for cte in &with.cte_tables {
892 let cte_name = cte.alias.name.value.clone();
894 let mut lf = self.execute_isolated(|ctx| ctx.execute_query(&cte.query))?;
895 lf = self.rename_columns_from_table_alias(lf, &cte.alias)?;
896 self.register_cte(&cte_name, lf);
897 }
898 }
899 Ok(())
900 }
901
902 fn register_named_windows(
903 &mut self,
904 named_windows: &[NamedWindowDefinition],
905 ) -> PolarsResult<()> {
906 for NamedWindowDefinition(name, expr) in named_windows {
907 let spec = match expr {
908 NamedWindowExpr::NamedWindow(ref_name) => self
909 .named_windows
910 .get(&ref_name.value)
911 .ok_or_else(|| {
912 polars_err!(
913 SQLInterface:
914 "named window '{}' references undefined window '{}'",
915 name.value, ref_name.value
916 )
917 })?
918 .clone(),
919 NamedWindowExpr::WindowSpec(spec) => spec.clone(),
920 };
921 self.named_windows.insert(name.value.clone(), spec);
922 }
923 Ok(())
924 }
925
926 pub(crate) fn execute_from_statement(
928 &mut self,
929 tbl_expr: &TableWithJoins,
930 ) -> PolarsResult<LazyFrame> {
931 let (l_name, mut lf) = self.get_table(&tbl_expr.relation)?;
932 if !tbl_expr.joins.is_empty() {
933 for join in &tbl_expr.joins {
934 if let (
936 JoinOperator::CrossJoin(JoinConstraint::None),
937 TableFactor::UNNEST {
938 alias,
939 array_exprs,
940 with_offset,
941 ..
942 },
943 ) = (&join.join_operator, &join.relation)
944 {
945 if array_exprs.iter().any(|e| !matches!(e, SQLExpr::Array(_))) {
946 lf = self.process_unnest_lateral(lf, alias, array_exprs, *with_offset)?;
947 continue;
948 }
949 }
950
951 let (r_name, mut rf) = self.get_table(&join.relation)?;
952 if r_name.is_empty() {
953 polars_bail!(
955 SQLInterface:
956 "cannot JOIN on unnamed relation; please provide an alias"
957 )
958 }
959 let left_schema = self.get_frame_schema(&mut lf)?;
960 let right_schema = self.get_frame_schema(&mut rf)?;
961
962 lf = match &join.join_operator {
963 op @ (JoinOperator::Join(constraint) | JoinOperator::FullOuter(constraint)
965 | JoinOperator::Left(constraint)
966 | JoinOperator::LeftOuter(constraint)
967 | JoinOperator::Right(constraint)
968 | JoinOperator::RightOuter(constraint)
969 | JoinOperator::Inner(constraint)
970 | JoinOperator::Anti(constraint)
971 | JoinOperator::Semi(constraint)
972 | JoinOperator::LeftAnti(constraint)
973 | JoinOperator::LeftSemi(constraint)
974 | JoinOperator::RightAnti(constraint)
975 | JoinOperator::RightSemi(constraint)) => {
976 let (lf, rf) = match op {
977 JoinOperator::RightAnti(_) | JoinOperator::RightSemi(_) => (rf, lf),
978 _ => (lf, rf),
979 };
980 self.process_join(
981 &TableInfo {
982 frame: lf,
983 name: (&l_name).into(),
984 schema: left_schema.clone(),
985 },
986 &TableInfo {
987 frame: rf,
988 name: (&r_name).into(),
989 schema: right_schema.clone(),
990 },
991 constraint,
992 match op {
993 JoinOperator::Join(_) | JoinOperator::Inner(_) => JoinType::Inner,
994 JoinOperator::Left(_) | JoinOperator::LeftOuter(_) => {
995 JoinType::Left
996 },
997 JoinOperator::Right(_) | JoinOperator::RightOuter(_) => {
998 JoinType::Right
999 },
1000 JoinOperator::FullOuter(_) => JoinType::Full,
1001 #[cfg(feature = "semi_anti_join")]
1002 JoinOperator::Anti(_)
1003 | JoinOperator::LeftAnti(_)
1004 | JoinOperator::RightAnti(_) => JoinType::Anti,
1005 #[cfg(feature = "semi_anti_join")]
1006 JoinOperator::Semi(_)
1007 | JoinOperator::LeftSemi(_)
1008 | JoinOperator::RightSemi(_) => JoinType::Semi,
1009 join_type => polars_bail!(
1010 SQLInterface:
1011 "join type '{:?}' not currently supported",
1012 join_type
1013 ),
1014 },
1015 )?
1016 },
1017 JoinOperator::CrossJoin(JoinConstraint::None) => {
1018 lf.cross_join(rf, Some(format_pl_smallstr!(":{}", r_name)))
1019 },
1020 JoinOperator::CrossJoin(constraint) => {
1021 polars_bail!(
1022 SQLInterface:
1023 "CROSS JOIN does not support {:?} constraint; consider INNER JOIN instead",
1024 constraint
1025 )
1026 },
1027 join_type => {
1028 polars_bail!(SQLInterface: "join type '{:?}' not currently supported", join_type)
1029 },
1030 };
1031
1032 let joined_schema = self.get_frame_schema(&mut lf)?;
1034
1035 self.joined_aliases.insert(
1036 r_name.clone(),
1037 right_schema
1038 .iter_names()
1039 .filter_map(|name| {
1040 let aliased_name = format!("{name}:{r_name}");
1042 if left_schema.contains(name)
1043 && joined_schema.contains(aliased_name.as_str())
1044 {
1045 Some((name.to_string(), aliased_name))
1046 } else {
1047 None
1048 }
1049 })
1050 .collect::<PlHashMap<String, String>>(),
1051 );
1052 }
1053 };
1054 Ok(lf)
1055 }
1056
1057 fn validate_select(&self, select_stmt: &Select) -> PolarsResult<()> {
1059 let Select {
1062 distinct: _,
1064 from: _,
1065 group_by: _,
1066 having: _,
1067 named_window: _,
1068 projection: _,
1069 qualify: _,
1070 selection: _,
1071
1072 flavor: _,
1074 select_token: _,
1075 top_before_distinct: _,
1076 window_before_qualify: _,
1077
1078 ref cluster_by,
1080 ref connect_by,
1081 ref distribute_by,
1082 ref exclude,
1083 ref into,
1084 ref lateral_views,
1085 ref optimizer_hints,
1086 ref prewhere,
1087 ref select_modifiers,
1088 ref sort_by,
1089 ref top,
1090 ref value_table_mode,
1091 } = *select_stmt;
1092
1093 polars_ensure!(cluster_by.is_empty(), SQLInterface: "`CLUSTER BY` clause is not supported");
1095 polars_ensure!(connect_by.is_empty(), SQLInterface: "`CONNECT BY` clause is not supported");
1096 polars_ensure!(distribute_by.is_empty(), SQLInterface: "`DISTRIBUTE BY` clause is not supported");
1097 polars_ensure!(exclude.is_none(), SQLInterface: "`EXCLUDE` clause is not supported");
1098 polars_ensure!(into.is_none(), SQLInterface: "`SELECT INTO` clause is not supported");
1099 polars_ensure!(lateral_views.is_empty(), SQLInterface: "`LATERAL VIEW` clause is not supported");
1100 polars_ensure!(optimizer_hints.is_empty(), SQLInterface: "optimizer hints are not supported");
1101 polars_ensure!(prewhere.is_none(), SQLInterface: "`PREWHERE` clause is not supported");
1102 polars_ensure!(select_modifiers.is_none(), SQLInterface: "`SELECT` modifiers are not supported");
1103 polars_ensure!(sort_by.is_empty(), SQLInterface: "`SORT BY` clause is not supported; use `ORDER BY` instead");
1104 polars_ensure!(top.is_none(), SQLInterface: "`TOP` clause is not supported; use `LIMIT` instead");
1105 polars_ensure!(value_table_mode.is_none(), SQLInterface: "`SELECT AS VALUE/STRUCT` is not supported");
1106
1107 Ok(())
1108 }
1109
1110 fn validate_table_factor(&self, factor: &TableFactor) -> PolarsResult<()> {
1112 let TableFactor::Table {
1115 name: _,
1117 alias: _,
1118 args: _,
1119
1120 ref index_hints,
1122 ref json_path,
1123 ref partitions,
1124 ref sample,
1125 ref version,
1126 ref with_hints,
1127 with_ordinality,
1128 } = *factor
1129 else {
1130 return Ok(());
1131 };
1132
1133 polars_ensure!(!with_ordinality, SQLInterface: "`WITH ORDINALITY` is not supported");
1134 polars_ensure!(index_hints.is_empty(), SQLInterface: "table index hints are not supported");
1135 polars_ensure!(json_path.is_none(), SQLInterface: "table JSON path access is not supported");
1136 polars_ensure!(partitions.is_empty(), SQLInterface: "table `PARTITION` selection is not supported");
1137 polars_ensure!(sample.is_none(), SQLInterface: "table `SAMPLE` clause is not supported");
1138 polars_ensure!(version.is_none(), SQLInterface: "table version (time-travel) qualifiers are not supported");
1139 polars_ensure!(with_hints.is_empty(), SQLInterface: "table `WITH (...)` hints are not supported");
1140
1141 Ok(())
1142 }
1143
1144 fn validate_create_table(&self, create_table: &CreateTable) -> PolarsResult<()> {
1146 let CreateTable {
1149 name: _,
1151 columns: _,
1152 query: _,
1153 like: _,
1154 if_not_exists: _, ref copy_grants,
1158 ref dynamic,
1159 ref external,
1160 ref global,
1161 ref iceberg,
1162 ref or_replace,
1163 ref require_user,
1164 ref snapshot,
1165 ref strict,
1166 ref temporary,
1167 ref transient,
1168 ref volatile,
1169 ref without_rowid,
1170
1171 ref backup,
1173 ref base_location,
1174 ref catalog,
1175 ref catalog_sync,
1176 ref change_tracking,
1177 ref clone,
1178 ref cluster_by,
1179 ref clustered_by,
1180 ref comment,
1181 ref constraints,
1182 ref data_retention_time_in_days,
1183 ref default_ddl_collation,
1184 ref distkey,
1185 ref diststyle,
1186 ref enable_schema_evolution,
1187 ref external_volume,
1188 ref file_format,
1189 ref for_values,
1190 ref hive_distribution,
1191 ref hive_formats,
1192 ref inherits,
1193 ref initialize,
1194 ref location,
1195 ref max_data_extension_time_in_days,
1196 ref on_cluster,
1197 ref on_commit,
1198 ref order_by,
1199 ref partition_by,
1200 ref partition_of,
1201 ref primary_key,
1202 ref refresh_mode,
1203 ref sortkey,
1204 ref storage_serialization_policy,
1205 ref table_options,
1206 ref target_lag,
1207 ref version,
1208 ref warehouse,
1209 ref with_aggregation_policy,
1210 ref with_row_access_policy,
1211 ref with_storage_lifecycle_policy,
1212 ref with_tags,
1213 } = *create_table;
1214
1215 polars_ensure!(!copy_grants, SQLInterface: "`COPY GRANTS` is not supported");
1216 polars_ensure!(!dynamic, SQLInterface: "`CREATE DYNAMIC TABLE` is not supported");
1217 polars_ensure!(!external, SQLInterface: "`CREATE EXTERNAL TABLE` is not supported");
1218 polars_ensure!(!iceberg, SQLInterface: "`CREATE ICEBERG TABLE` is not supported");
1219 polars_ensure!(!or_replace, SQLInterface: "`CREATE OR REPLACE TABLE` is not supported");
1220 polars_ensure!(!require_user, SQLInterface: "`REQUIRE USER` is not supported");
1221 polars_ensure!(!snapshot, SQLInterface: "`CREATE SNAPSHOT TABLE` is not supported");
1222 polars_ensure!(!strict, SQLInterface: "`STRICT` tables are not supported");
1223 polars_ensure!(!temporary, SQLInterface: "`CREATE TEMPORARY TABLE` is not supported");
1224 polars_ensure!(!transient, SQLInterface: "`CREATE TRANSIENT TABLE` is not supported");
1225 polars_ensure!(!volatile, SQLInterface: "`CREATE VOLATILE TABLE` is not supported");
1226 polars_ensure!(!without_rowid, SQLInterface: "`WITHOUT ROWID` is not supported");
1227 polars_ensure!(backup.is_none(), SQLInterface: "`BACKUP` is not supported");
1228 polars_ensure!(base_location.is_none(), SQLInterface: "`BASE_LOCATION` is not supported");
1229 polars_ensure!(catalog.is_none(), SQLInterface: "`CATALOG` is not supported");
1230 polars_ensure!(catalog_sync.is_none(), SQLInterface: "`CATALOG_SYNC` is not supported");
1231 polars_ensure!(change_tracking.is_none(), SQLInterface: "`CHANGE_TRACKING` is not supported");
1232 polars_ensure!(clone.is_none(), SQLInterface: "`CREATE TABLE ... CLONE` is not supported");
1233 polars_ensure!(cluster_by.is_none(), SQLInterface: "table `CLUSTER BY` clauses are not supported");
1234 polars_ensure!(clustered_by.is_none(), SQLInterface: "table `CLUSTERED BY` clauses are not supported");
1235 polars_ensure!(comment.is_none(), SQLInterface: "table `COMMENT` clauses are not supported");
1236 polars_ensure!(constraints.is_empty(), SQLInterface: "table constraints are not supported");
1237 polars_ensure!(data_retention_time_in_days.is_none(), SQLInterface: "`DATA_RETENTION_TIME_IN_DAYS` is not supported");
1238 polars_ensure!(default_ddl_collation.is_none(), SQLInterface: "`DEFAULT_DDL_COLLATION` is not supported");
1239 polars_ensure!(distkey.is_none(), SQLInterface: "`DISTKEY` is not supported");
1240 polars_ensure!(diststyle.is_none(), SQLInterface: "`DISTSTYLE` is not supported");
1241 polars_ensure!(enable_schema_evolution.is_none(), SQLInterface: "`ENABLE_SCHEMA_EVOLUTION` is not supported");
1242 polars_ensure!(external_volume.is_none(), SQLInterface: "`EXTERNAL_VOLUME` is not supported");
1243 polars_ensure!(file_format.is_none(), SQLInterface: "table `STORED AS` clauses are not supported");
1244 polars_ensure!(for_values.is_none(), SQLInterface: "`FOR VALUES` clauses are not supported");
1245 polars_ensure!(global.is_none(), SQLInterface: "`CREATE GLOBAL/LOCAL TABLE` is not supported");
1246 polars_ensure!(hive_formats.as_ref().is_none_or(|f| *f == HiveFormat::default()), SQLInterface: "Hive table format clauses are not supported");
1247 polars_ensure!(inherits.is_none(), SQLInterface: "table `INHERITS` clauses are not supported");
1248 polars_ensure!(initialize.is_none(), SQLInterface: "`INITIALIZE` is not supported");
1249 polars_ensure!(location.is_none(), SQLInterface: "table `LOCATION` clauses are not supported");
1250 polars_ensure!(matches!(hive_distribution, HiveDistributionStyle::NONE), SQLInterface: "Hive table distribution clauses are not supported");
1251 polars_ensure!(matches!(table_options, CreateTableOptions::None), SQLInterface: "table `WITH`/`OPTIONS` clauses are not supported");
1252 polars_ensure!(max_data_extension_time_in_days.is_none(), SQLInterface: "`MAX_DATA_EXTENSION_TIME_IN_DAYS` is not supported");
1253 polars_ensure!(on_cluster.is_none(), SQLInterface: "`ON CLUSTER` clauses are not supported");
1254 polars_ensure!(on_commit.is_none(), SQLInterface: "`ON COMMIT` clauses are not supported");
1255 polars_ensure!(order_by.is_none(), SQLInterface: "table `ORDER BY` clauses are not supported");
1256 polars_ensure!(partition_by.is_none(), SQLInterface: "table `PARTITION BY` clauses are not supported");
1257 polars_ensure!(partition_of.is_none(), SQLInterface: "`PARTITION OF` clauses are not supported");
1258 polars_ensure!(primary_key.is_none(), SQLInterface: "inline `PRIMARY KEY` clauses are not supported");
1259 polars_ensure!(refresh_mode.is_none(), SQLInterface: "`REFRESH_MODE` is not supported");
1260 polars_ensure!(sortkey.is_none(), SQLInterface: "`SORTKEY` is not supported");
1261 polars_ensure!(storage_serialization_policy.is_none(), SQLInterface: "`STORAGE_SERIALIZATION_POLICY` is not supported");
1262 polars_ensure!(target_lag.is_none(), SQLInterface: "`TARGET_LAG` is not supported");
1263 polars_ensure!(version.is_none(), SQLInterface: "table version (time-travel) qualifiers are not supported");
1264 polars_ensure!(warehouse.is_none(), SQLInterface: "`WAREHOUSE` is not supported");
1265 polars_ensure!(with_aggregation_policy.is_none(), SQLInterface: "`WITH AGGREGATION POLICY` is not supported");
1266 polars_ensure!(with_row_access_policy.is_none(), SQLInterface: "`WITH ROW ACCESS POLICY` is not supported");
1267 polars_ensure!(with_storage_lifecycle_policy.is_none(), SQLInterface: "`WITH STORAGE LIFECYCLE POLICY` is not supported");
1268 polars_ensure!(with_tags.is_none(), SQLInterface: "`WITH TAG` is not supported");
1269
1270 Ok(())
1271 }
1272
1273 fn validate_query(&self, query: &Query) -> PolarsResult<()> {
1275 let Query {
1277 with: _,
1279 body: _,
1280 order_by: _,
1281 limit_clause: _,
1282 fetch,
1283
1284 for_clause,
1286 format_clause,
1287 locks,
1288 pipe_operators,
1289 settings,
1290 } = query;
1291
1292 polars_ensure!(for_clause.is_none(), SQLInterface: "`FOR` clause is not supported");
1294 polars_ensure!(format_clause.is_none(), SQLInterface: "`FORMAT` clause is not supported");
1295 polars_ensure!(locks.is_empty(), SQLInterface: "`FOR UPDATE/SHARE` locking clause is not supported");
1296 polars_ensure!(pipe_operators.is_empty(), SQLInterface: "pipe operators are not supported");
1297 polars_ensure!(settings.is_none(), SQLInterface: "`SETTINGS` clause is not supported");
1298
1299 if let Some(Fetch {
1301 quantity: _, percent,
1303 with_ties,
1304 }) = fetch
1305 {
1306 polars_ensure!(!percent, SQLInterface: "`FETCH` with `PERCENT` is not supported");
1307 polars_ensure!(!with_ties, SQLInterface: "`FETCH` with `WITH TIES` is not supported");
1308 }
1309 Ok(())
1310 }
1311
1312 fn execute_select(&mut self, select_stmt: &Select, query: &Query) -> PolarsResult<LazyFrame> {
1314 self.validate_select(select_stmt)?;
1316
1317 self.register_named_windows(&select_stmt.named_window)?;
1319
1320 let mut implicit_join_filter: Option<SQLExpr> = None;
1322 let (mut lf, base_table_name) = if select_stmt.from.is_empty() {
1323 (DataFrame::empty().lazy(), None)
1324 } else {
1325 let from = &select_stmt.from;
1326 let first = from.first().unwrap();
1327 let mut lf = self.execute_from_statement(first)?;
1328 let base_name = get_table_name(&first.relation);
1329 if from.len() > 1 {
1330 implicit_join_filter =
1331 self.process_implicit_joins(&mut lf, from, &select_stmt.selection)?;
1332 }
1333 (lf, base_name)
1334 };
1335
1336 if let Some(ref base_name) = base_table_name {
1338 if !self.joined_aliases.is_empty() {
1339 let using_cols: PlHashSet<String> = select_stmt
1341 .from
1342 .first()
1343 .into_iter()
1344 .flat_map(|t| t.joins.iter())
1345 .filter_map(|join| get_using_cols(&join.join_operator))
1346 .flatten()
1347 .collect();
1348
1349 let check_expr = |e| {
1351 check_for_ambiguous_column_refs(e, &self.joined_aliases, base_name, &using_cols)
1352 };
1353 for item in &select_stmt.projection {
1354 match item {
1355 SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => {
1356 check_expr(e)?
1357 },
1358 _ => {},
1359 }
1360 }
1361 if let Some(ref where_expr) = select_stmt.selection {
1362 check_expr(where_expr)?;
1363 }
1364 }
1365 }
1366
1367 let effective_where = if implicit_join_filter.is_some() {
1369 &implicit_join_filter
1370 } else {
1371 &select_stmt.selection
1372 };
1373 let mut schema = self.get_frame_schema(&mut lf)?;
1374 lf = self.process_where(
1375 lf,
1376 effective_where,
1377 FilterMode::KeepTrue,
1378 Some(schema.clone()),
1379 )?;
1380
1381 let mut select_modifiers = SelectModifiers {
1383 ilike: None,
1384 exclude: PlHashSet::new(),
1385 rename: PlHashMap::new(),
1386 replace: vec![],
1387 };
1388
1389 let window_fn_columns = if select_stmt.qualify.is_some() {
1392 select_stmt
1393 .projection
1394 .iter()
1395 .filter_map(|item| match item {
1396 SelectItem::ExprWithAlias { expr, alias }
1397 if expr_has_window_functions(expr) =>
1398 {
1399 Some(alias.value.clone())
1400 },
1401 _ => None,
1402 })
1403 .collect::<PlHashSet<_>>()
1404 } else {
1405 PlHashSet::new()
1406 };
1407
1408 let mut projections =
1409 self.column_projections(select_stmt, &schema, &mut select_modifiers)?;
1410
1411 let mut explode_names = Vec::new();
1413 let mut explode_exprs = Vec::new();
1414 let mut explode_lookup = PlHashMap::new();
1415
1416 for expr in &projections {
1417 for e in expr {
1418 if let Expr::Explode { input, .. } = e {
1419 match input.as_ref() {
1420 Expr::Column(name) => explode_names.push(name.clone()),
1421 other_expr => {
1422 if !has_expr(other_expr, |e| matches!(e, Expr::Agg(_) | Expr::Len)) {
1424 let temp_name =
1425 format_pl_smallstr!("__POLARS_UNNEST_{}", explode_exprs.len());
1426 explode_exprs.push(other_expr.clone().alias(temp_name.as_str()));
1427 explode_lookup.insert(other_expr.clone(), temp_name.clone());
1428 explode_names.push(temp_name);
1429 }
1430 },
1431 }
1432 }
1433 }
1434 }
1435 if !explode_names.is_empty() {
1436 if !explode_exprs.is_empty() {
1437 lf = lf.with_columns(explode_exprs);
1438 }
1439 lf = lf.explode(
1440 Selector::ByName {
1441 names: Arc::from(explode_names),
1442 strict: true,
1443 },
1444 ExplodeOptions {
1445 empty_as_null: true,
1446 keep_nulls: true,
1447 },
1448 );
1449 projections = projections
1450 .into_iter()
1451 .map(|p| {
1452 p.map_expr(|e| match e {
1454 Expr::Explode { input, .. } => explode_lookup
1455 .get(input.as_ref())
1456 .map(|name| Expr::Column(name.clone()))
1457 .unwrap_or_else(|| input.as_ref().clone()),
1458 _ => e,
1459 })
1460 })
1461 .collect();
1462
1463 schema = self.get_frame_schema(&mut lf)?;
1464 }
1465
1466 let mut group_by_keys: Vec<Expr> = Vec::new();
1468 match &select_stmt.group_by {
1469 GroupByExpr::Expressions(group_by_exprs, modifiers) => {
1471 if !modifiers.is_empty() {
1472 polars_bail!(SQLInterface: "GROUP BY does not support CUBE, ROLLUP, or TOTALS modifiers")
1473 }
1474 group_by_keys = group_by_exprs
1476 .iter()
1477 .map(|e| match e {
1478 SQLExpr::Identifier(ident) => {
1479 resolve_select_alias(&ident.value, &projections, &schema).map_or_else(
1480 || {
1481 self.expr_or_ordinal(
1482 e,
1483 &projections,
1484 None,
1485 Some(&schema),
1486 "GROUP BY",
1487 )
1488 },
1489 Ok,
1490 )
1491 },
1492 _ => self.expr_or_ordinal(e, &projections, None, Some(&schema), "GROUP BY"),
1493 })
1494 .collect::<PolarsResult<_>>()?
1495 },
1496 GroupByExpr::All(modifiers) => {
1499 if !modifiers.is_empty() {
1500 polars_bail!(SQLInterface: "GROUP BY does not support CUBE, ROLLUP, or TOTALS modifiers")
1501 }
1502 projections.iter().for_each(|expr| match expr {
1503 Expr::Agg(_) | Expr::Len | Expr::Literal(_) => (),
1505 Expr::Column(_) => group_by_keys.push(expr.clone()),
1506 Expr::Alias(e, _)
1507 if matches!(&**e, Expr::Agg(_) | Expr::Len | Expr::Literal(_)) => {},
1508 Expr::Alias(e, _) if matches!(&**e, Expr::Column(_)) => {
1509 if let Expr::Column(name) = &**e {
1510 group_by_keys.push(col(name.clone()));
1511 }
1512 },
1513 _ => {
1514 if !has_expr(expr, |e| {
1516 matches!(e, Expr::Agg(_))
1517 || matches!(e, Expr::Len)
1518 || matches!(e, Expr::Over { .. })
1519 || {
1520 #[cfg(feature = "dynamic_group_by")]
1521 {
1522 matches!(e, Expr::Rolling { .. })
1523 }
1524 #[cfg(not(feature = "dynamic_group_by"))]
1525 {
1526 false
1527 }
1528 }
1529 }) {
1530 group_by_keys.push(expr.clone())
1531 }
1532 },
1533 });
1534 },
1535 };
1536
1537 lf = if group_by_keys.is_empty() {
1538 if select_stmt.having.is_some() {
1540 polars_bail!(SQLSyntax: "HAVING clause not valid outside of GROUP BY; found:\n{:?}", select_stmt.having);
1541 };
1542
1543 let mut retained_cols = Vec::with_capacity(projections.len());
1545 let mut retained_names = Vec::with_capacity(projections.len());
1546 let have_order_by = query.order_by.is_some();
1547
1548 let mut projection_heights = ExprSqlProjectionHeightBehavior::InheritsContext;
1550
1551 for p in projections.iter() {
1556 let name = p.to_field(schema.deref())?.name.to_string();
1557 if select_modifiers.matches_ilike(&name)
1558 && !select_modifiers.exclude.contains(&name)
1559 {
1560 projection_heights |= ExprSqlProjectionHeightBehavior::identify_from_expr(p);
1561
1562 retained_cols.push(if have_order_by {
1563 col(name.as_str())
1564 } else {
1565 p.clone()
1566 });
1567 retained_names.push(col(name));
1568 }
1569 }
1570
1571 if have_order_by {
1573 if projection_heights.contains(ExprSqlProjectionHeightBehavior::MaintainsColumn)
1578 || projection_heights == ExprSqlProjectionHeightBehavior::InheritsContext
1579 {
1580 lf = lf.with_columns(projections);
1581 } else {
1582 const NAME: PlSmallStr = PlSmallStr::from_static("__PL_INDEX");
1591 lf = lf
1592 .clone()
1593 .select(projections)
1594 .with_row_index(NAME, None)
1595 .join(
1596 lf.with_row_index(NAME, None),
1597 [col(NAME)],
1598 [col(NAME)],
1599 JoinArgs {
1600 how: JoinType::Left,
1601 validation: Default::default(),
1602 suffix: None,
1603 slice: None,
1604 nulls_equal: false,
1605 coalesce: Default::default(),
1606 maintain_order: MaintainOrderJoin::Left,
1607 build_side: None,
1608 },
1609 );
1610 }
1611 }
1612 if !select_modifiers.replace.is_empty() {
1613 lf = lf.with_columns(&select_modifiers.replace);
1614 }
1615 if !select_modifiers.rename.is_empty() {
1616 lf = lf.with_columns(select_modifiers.renamed_cols());
1617 }
1618 lf = self.process_order_by(lf, &query.order_by, Some(&retained_cols))?;
1619
1620 if projection_heights == ExprSqlProjectionHeightBehavior::InheritsContext
1622 && !have_order_by
1623 {
1624 lf = lf.with_columns(retained_cols).select(retained_names);
1626 } else {
1627 lf = lf.select(retained_cols);
1628 }
1629 if !select_modifiers.rename.is_empty() {
1630 lf = lf.rename(
1631 select_modifiers.rename.keys(),
1632 select_modifiers.rename.values(),
1633 true,
1634 );
1635 };
1636 lf
1637 } else {
1638 let having = select_stmt
1639 .having
1640 .as_ref()
1641 .map(|expr| parse_sql_expr(expr, self, Some(&schema)))
1642 .transpose()?;
1643 lf = self.process_group_by(lf, &group_by_keys, &projections, having)?;
1644 lf = self.process_order_by(lf, &query.order_by, None)?;
1645
1646 let output_cols: Vec<_> = projections
1648 .iter()
1649 .map(|p| p.to_field(&schema))
1650 .collect::<PolarsResult<Vec<_>>>()?
1651 .into_iter()
1652 .map(|f| col(f.name))
1653 .collect();
1654
1655 lf.select(&output_cols)
1656 };
1657
1658 lf = self.process_qualify(lf, &select_stmt.qualify, &window_fn_columns)?;
1660
1661 lf = match &select_stmt.distinct {
1663 Some(Distinct::Distinct) => lf.unique_stable(None, UniqueKeepStrategy::Any),
1664 Some(Distinct::On(exprs)) => {
1665 let schema = Some(self.get_frame_schema(&mut lf)?);
1667 let cols = exprs
1668 .iter()
1669 .map(|e| {
1670 let expr = parse_sql_expr(e, self, schema.as_deref())?;
1671 if let Expr::Column(name) = expr {
1672 Ok(name)
1673 } else {
1674 Err(polars_err!(SQLSyntax:"DISTINCT ON only supports column names"))
1675 }
1676 })
1677 .collect::<PolarsResult<Vec<_>>>()?;
1678
1679 lf = self.process_order_by(lf, &query.order_by, None)?;
1681 return Ok(lf.unique_stable(
1682 Some(Selector::ByName {
1683 names: cols.into(),
1684 strict: true,
1685 }),
1686 UniqueKeepStrategy::First,
1687 ));
1688 },
1689 Some(Distinct::All) | None => lf,
1691 };
1692 Ok(lf)
1693 }
1694
1695 fn column_projections(
1696 &mut self,
1697 select_stmt: &Select,
1698 schema: &SchemaRef,
1699 select_modifiers: &mut SelectModifiers,
1700 ) -> PolarsResult<Vec<Expr>> {
1701 if select_stmt.projection.is_empty()
1702 && select_stmt.flavor == SelectFlavor::FromFirstNoSelect
1703 {
1704 return Ok(schema.iter_names().map(|name| col(name.clone())).collect());
1706 }
1707 let mut items: Vec<ProjectionItem> = Vec::with_capacity(select_stmt.projection.len());
1708 let mut has_qualified_wildcard = false;
1709
1710 for select_item in &select_stmt.projection {
1711 match select_item {
1712 SelectItem::UnnamedExpr(expr) => {
1713 items.push(ProjectionItem::Exprs(vec![parse_sql_expr(
1714 expr,
1715 self,
1716 Some(schema),
1717 )?]));
1718 },
1719 SelectItem::ExprWithAlias { expr, alias } => {
1720 let expr = parse_sql_expr(expr, self, Some(schema))?;
1721 items.push(ProjectionItem::Exprs(vec![
1722 expr.alias(PlSmallStr::from_str(alias.value.as_str())),
1723 ]));
1724 },
1725 SelectItem::ExprWithAliases { .. } => {
1726 polars_bail!(SQLSyntax: "multiple aliases per expression are not supported: {:?}", select_item)
1727 },
1728 SelectItem::QualifiedWildcard(kind, wildcard_options) => match kind {
1729 SelectItemQualifiedWildcardKind::ObjectName(obj_name) => {
1730 let tbl_name = obj_name
1731 .0
1732 .last()
1733 .and_then(|p| p.as_ident())
1734 .map(|i| PlSmallStr::from_str(&i.value))
1735 .unwrap_or_default();
1736 let exprs = self.process_qualified_wildcard(
1737 obj_name,
1738 wildcard_options,
1739 select_modifiers,
1740 Some(schema),
1741 )?;
1742 items.push(ProjectionItem::QualifiedExprs(tbl_name, exprs));
1743 has_qualified_wildcard = true;
1744 },
1745 SelectItemQualifiedWildcardKind::Expr(_) => {
1746 polars_bail!(SQLSyntax: "qualified wildcard on expressions not yet supported: {:?}", select_item)
1747 },
1748 },
1749 SelectItem::Wildcard(wildcard_options) => {
1750 let cols = schema.iter_names().map(|name| col(name.clone())).collect();
1751 items.push(ProjectionItem::Exprs(
1752 self.process_wildcard_additional_options(
1753 cols,
1754 wildcard_options,
1755 select_modifiers,
1756 Some(schema),
1757 )?,
1758 ));
1759 },
1760 }
1761 }
1762
1763 let exprs = if has_qualified_wildcard {
1765 disambiguate_projection_cols(items, schema)?
1766 } else {
1767 items
1768 .into_iter()
1769 .flat_map(|item| match item {
1770 ProjectionItem::Exprs(exprs) | ProjectionItem::QualifiedExprs(_, exprs) => {
1771 exprs
1772 },
1773 })
1774 .collect()
1775 };
1776 let flattened_exprs = exprs
1777 .into_iter()
1778 .flat_map(|expr| expand_exprs(expr, schema))
1779 .collect();
1780
1781 Ok(flattened_exprs)
1782 }
1783
1784 fn process_where(
1785 &mut self,
1786 mut lf: LazyFrame,
1787 expr: &Option<SQLExpr>,
1788 filter_mode: FilterMode,
1789 schema: Option<SchemaRef>,
1790 ) -> PolarsResult<LazyFrame> {
1791 if let Some(expr) = expr {
1792 let schema = match schema {
1793 None => self.get_frame_schema(&mut lf)?,
1794 Some(s) => s,
1795 };
1796
1797 let (all_true, all_false) = match expr {
1799 SQLExpr::Value(ValueWithSpan {
1800 value: SQLValue::Boolean(b),
1801 ..
1802 }) => (*b, !*b),
1803 SQLExpr::BinaryOp { left, op, right } => match (&**left, &**right, op) {
1804 (SQLExpr::Value(a), SQLExpr::Value(b), SQLBinaryOperator::Eq) => {
1805 (a.value == b.value, a.value != b.value)
1806 },
1807 (SQLExpr::Value(a), SQLExpr::Value(b), SQLBinaryOperator::NotEq) => {
1808 (a.value != b.value, a.value == b.value)
1809 },
1810 _ => (false, false),
1811 },
1812 _ => (false, false),
1813 };
1814 let removing = filter_mode == FilterMode::RemoveTrue;
1815 if (all_true && !removing) || (all_false && removing) {
1816 return Ok(lf);
1817 } else if (all_false && !removing) || (all_true && removing) {
1818 return Ok(lf.clear());
1819 }
1820
1821 let residual_exprs: Vec<&SQLExpr>;
1825 (lf, residual_exprs) =
1826 self.rewrite_subquery_conjuncts(lf, expr, filter_mode, &schema)?;
1827
1828 let Some(parsed_residual) = residual_exprs
1829 .iter()
1830 .map(|e| parse_sql_expr(e, self, Some(&*schema)))
1831 .reduce(|a, b| Ok(a?.and(b?)))
1832 else {
1833 return Ok(lf);
1835 };
1836 let mut filter_expression = parsed_residual?;
1837 if filter_expression.clone().meta().has_multiple_outputs() {
1838 filter_expression = all_horizontal([filter_expression])?;
1839 }
1840 lf = self.process_subqueries(lf, vec![&mut filter_expression])?;
1841 lf = match filter_mode {
1842 FilterMode::KeepTrue => lf.filter(filter_expression),
1843 FilterMode::RemoveTrue => lf.remove(filter_expression),
1844 };
1845 }
1846 Ok(lf)
1847 }
1848
1849 pub(super) fn process_join(
1850 &mut self,
1851 tbl_left: &TableInfo,
1852 tbl_right: &TableInfo,
1853 constraint: &JoinConstraint,
1854 join_type: JoinType,
1855 ) -> PolarsResult<LazyFrame> {
1856 let (left_on, right_on, predicates) =
1857 process_join_constraint(constraint, tbl_left, tbl_right, self)?;
1858 let coalesce_type = match constraint {
1859 JoinConstraint::Natural => JoinCoalesce::CoalesceColumns,
1861 _ => JoinCoalesce::KeepColumns,
1862 };
1863 let suffix = format!(":{}", tbl_right.name);
1864
1865 let joined = if predicates.is_empty() {
1866 tbl_left
1868 .frame
1869 .clone()
1870 .join_builder()
1871 .with(tbl_right.frame.clone())
1872 .left_on(left_on)
1873 .right_on(right_on)
1874 .how(join_type)
1875 .suffix(suffix)
1876 .coalesce(coalesce_type)
1877 .finish()
1878 } else {
1879 let mut all_predicates = predicates;
1883 for (l, r) in left_on.into_iter().zip(right_on) {
1884 let r_suffixed = suffix_conflicting_columns(r, tbl_left, tbl_right, &suffix);
1885 all_predicates.push(l.eq(r_suffixed));
1886 }
1887 tbl_left
1888 .frame
1889 .clone()
1890 .join_builder()
1891 .with(tbl_right.frame.clone())
1892 .how(join_type)
1893 .suffix(suffix)
1894 .coalesce(coalesce_type)
1895 .join_where(all_predicates)
1896 };
1897
1898 Ok(joined)
1899 }
1900
1901 fn process_implicit_joins(
1907 &mut self,
1908 lf: &mut LazyFrame,
1909 from: &[TableWithJoins],
1910 where_clause: &Option<SQLExpr>,
1911 ) -> PolarsResult<Option<SQLExpr>> {
1912 let first = from.first().unwrap();
1913 let base_name = get_table_name(&first.relation).unwrap_or_default();
1914 let mut remaining_where = where_clause.clone();
1915 let mut joined_table_names: Vec<String> = vec![base_name];
1916
1917 for join in &first.joins {
1919 if let Some(name) = get_table_name(&join.relation) {
1920 joined_table_names.push(name);
1921 }
1922 }
1923 for tbl_expr in from.iter().skip(1) {
1924 let mut rf = self.execute_from_statement(tbl_expr)?;
1925 let r_name = get_table_name(&tbl_expr.relation).unwrap_or_default();
1926 polars_ensure!(
1927 !r_name.is_empty(),
1928 SQLInterface: "implicit joins require named tables; please provide an alias"
1929 );
1930 let left_schema = self.get_frame_schema(lf)?;
1931 let right_schema = self.get_frame_schema(&mut rf)?;
1932 let (join_expr, residual) =
1933 extract_join_predicates(&remaining_where, &joined_table_names, &r_name);
1934
1935 *lf = if let Some(on_expr) = join_expr {
1936 self.process_join(
1937 &TableInfo {
1938 frame: lf.clone(),
1939 name: PlSmallStr::from_str(
1940 &joined_table_names.first().cloned().unwrap_or_default(),
1941 ),
1942 schema: left_schema.clone(),
1943 },
1944 &TableInfo {
1945 frame: rf,
1946 name: PlSmallStr::from_str(&r_name),
1947 schema: right_schema.clone(),
1948 },
1949 &JoinConstraint::On(on_expr),
1950 JoinType::Inner,
1951 )?
1952 } else {
1953 lf.clone()
1954 .cross_join(rf, Some(format_pl_smallstr!(":{}", r_name)))
1955 };
1956 remaining_where = residual;
1957
1958 let joined_schema = self.get_frame_schema(lf)?;
1960 self.joined_aliases.insert(
1961 r_name.clone(),
1962 right_schema
1963 .iter_names()
1964 .filter_map(|name| {
1965 let aliased_name = format!("{name}:{r_name}");
1966 if left_schema.contains(name)
1967 && joined_schema.contains(aliased_name.as_str())
1968 {
1969 Some((name.to_string(), aliased_name))
1970 } else {
1971 None
1972 }
1973 })
1974 .collect::<PlHashMap<String, String>>(),
1975 );
1976 joined_table_names.push(r_name);
1977 for join in &tbl_expr.joins {
1978 if let Some(name) = get_table_name(&join.relation) {
1979 joined_table_names.push(name);
1980 }
1981 }
1982 }
1983 Ok(remaining_where)
1984 }
1985
1986 fn process_qualify(
1987 &mut self,
1988 mut lf: LazyFrame,
1989 qualify_expr: &Option<SQLExpr>,
1990 window_fn_columns: &PlHashSet<String>,
1991 ) -> PolarsResult<LazyFrame> {
1992 if let Some(expr) = qualify_expr {
1993 let (has_window_fns, column_refs) = QualifyExpression::analyze(expr);
1996 let references_window_alias = column_refs.iter().any(|c| window_fn_columns.contains(c));
1997 if !has_window_fns && !references_window_alias {
1998 polars_bail!(
1999 SQLSyntax:
2000 "QUALIFY clause must reference window functions either explicitly or via SELECT aliases"
2001 );
2002 }
2003 let schema = self.get_frame_schema(&mut lf)?;
2004 let mut filter_expression = parse_sql_expr(expr, self, Some(&schema))?;
2005 if filter_expression.clone().meta().has_multiple_outputs() {
2006 filter_expression = all_horizontal([filter_expression])?;
2007 }
2008 lf = self.process_subqueries(lf, vec![&mut filter_expression])?;
2009 lf = lf.filter(filter_expression);
2010 }
2011 Ok(lf)
2012 }
2013
2014 fn process_subqueries(
2015 &mut self,
2016 lf: LazyFrame,
2017 exprs: Vec<&mut Expr>,
2018 ) -> PolarsResult<LazyFrame> {
2019 let mut subplans = vec![];
2020
2021 for e in exprs {
2022 *e = e.clone().try_map_expr(|e| {
2023 if let Expr::SubPlan(lp, names) = e {
2024 assert_eq!(
2025 names.len(),
2026 1,
2027 "multiple columns in subqueries not yet supported"
2028 );
2029
2030 let select_expr = names[0].1.clone();
2031 let mut lf = LazyFrame::from((**lp).clone());
2032 let schema = self.get_frame_schema(&mut lf)?;
2033 polars_ensure!(schema.len() == 1, SQLSyntax: "SQL subquery returns more than one column");
2034 let lf = lf.select([select_expr.clone()]);
2035
2036 subplans.push(lf);
2037 Ok(Expr::Column(names[0].0.clone()).first())
2038 } else {
2039 Ok(e)
2040 }
2041 })?;
2042 }
2043
2044 if subplans.is_empty() {
2045 Ok(lf)
2046 } else {
2047 subplans.insert(0, lf);
2048 concat_lf_horizontal(
2049 subplans,
2050 HConcatOptions {
2051 broadcast_unit_length: true,
2052 ..Default::default()
2053 },
2054 )
2055 }
2056 }
2057
2058 fn execute_create_table(&mut self, stmt: &Statement) -> PolarsResult<LazyFrame> {
2059 if let Statement::CreateTable(create_table) = stmt {
2060 self.validate_create_table(create_table)?;
2061 let CreateTable {
2062 if_not_exists,
2063 name,
2064 query,
2065 columns,
2066 like,
2067 ..
2068 } = create_table;
2069
2070 let tbl_name = name.0.first().unwrap().as_ident().unwrap().value.as_str();
2071 if *if_not_exists && self.table_map.read().unwrap().contains_key(tbl_name) {
2072 polars_bail!(SQLInterface: "relation '{}' already exists", tbl_name);
2073 }
2074 let lf = match (query, columns.is_empty(), like) {
2075 (Some(query), true, None) => {
2076 self.execute_query(query)?
2080 },
2081 (None, false, None) => {
2082 let mut schema = Schema::with_capacity(columns.len());
2086 for col in columns {
2087 let col_name = col.name.value.as_str();
2088 let dtype = map_sql_dtype_to_polars(&col.data_type)?;
2089 schema.insert_at_index(schema.len(), col_name.into(), dtype)?;
2090 }
2091 DataFrame::empty_with_schema(&schema).lazy()
2092 },
2093 (None, true, Some(like_kind)) => {
2094 let like_name = match like_kind {
2098 CreateTableLikeKind::Plain(like)
2099 | CreateTableLikeKind::Parenthesized(like) => &like.name,
2100 };
2101 let like_table = like_name
2102 .0
2103 .first()
2104 .unwrap()
2105 .as_ident()
2106 .unwrap()
2107 .value
2108 .as_str();
2109 if let Some(table) = self.table_map.read().unwrap().get(like_table).cloned() {
2110 table.clear()
2111 } else {
2112 polars_bail!(SQLInterface: "table given in LIKE does not exist: {}", like_table)
2113 }
2114 },
2115 (None, true, None) => {
2117 polars_bail!(SQLInterface: "CREATE TABLE expected a query, column definitions, or LIKE clause")
2118 },
2119 _ => {
2121 polars_bail!(
2122 SQLInterface: "CREATE TABLE received mutually exclusive options:\nquery = {:?}\ncolumns = {:?}\nlike = {:?}",
2123 query,
2124 columns,
2125 like,
2126 )
2127 },
2128 };
2129 self.register(tbl_name, lf);
2130
2131 let df_created = df! { "Response" => [format!("CREATE TABLE {}", name.0.first().unwrap().as_ident().unwrap().value)] };
2132 Ok(df_created.unwrap().lazy())
2133 } else {
2134 unreachable!()
2135 }
2136 }
2137
2138 fn get_table(&mut self, relation: &TableFactor) -> PolarsResult<(String, LazyFrame)> {
2139 match relation {
2140 TableFactor::Table {
2141 name, alias, args, ..
2142 } => {
2143 self.validate_table_factor(relation)?;
2144 if let Some(args) = args {
2145 return self.execute_table_function(name, alias, &args.args);
2146 }
2147 let tbl_name = name.0.first().unwrap().as_ident().unwrap().value.as_str();
2148 if let Some(lf) = self.get_table_from_current_scope(tbl_name) {
2149 match alias {
2150 Some(alias) => {
2151 self.table_aliases
2152 .insert(alias.name.value.clone(), tbl_name.to_string());
2153 Ok((alias.name.value.clone(), lf))
2154 },
2155 None => Ok((tbl_name.to_string(), lf)),
2156 }
2157 } else {
2158 polars_bail!(SQLInterface: "relation '{}' was not found", tbl_name);
2159 }
2160 },
2161 TableFactor::Derived {
2162 lateral,
2163 subquery,
2164 alias,
2165 sample,
2166 } => {
2167 polars_ensure!(!(*lateral), SQLInterface: "`LATERAL` clause is not supported");
2168 polars_ensure!(sample.is_none(), SQLInterface: "table `SAMPLE` clause is not supported");
2169
2170 if let Some(alias) = alias {
2173 let mut lf =
2174 self.execute_isolated(|ctx| ctx.execute_query_no_ctes(subquery))?;
2175 lf = self.rename_columns_from_table_alias(lf, alias)?;
2176 self.table_map
2177 .write()
2178 .unwrap()
2179 .insert(alias.name.value.clone(), lf.clone());
2180 Ok((alias.name.value.clone(), lf))
2181 } else {
2182 let lf = self.execute_isolated(|ctx| ctx.execute_query_no_ctes(subquery))?;
2183 Ok(("".to_string(), lf))
2184 }
2185 },
2186 TableFactor::UNNEST {
2187 alias,
2188 array_exprs,
2189 with_offset,
2190 with_offset_alias: _,
2191 with_ordinality,
2192 } => {
2193 if let Some(alias) = alias {
2194 let column_names: Vec<Option<PlSmallStr>> = alias
2195 .columns
2196 .iter()
2197 .map(|c| {
2198 if c.name.value.is_empty() {
2199 None
2200 } else {
2201 Some(PlSmallStr::from_str(c.name.value.as_str()))
2202 }
2203 })
2204 .collect();
2205
2206 let column_values: Vec<Series> = array_exprs
2207 .iter()
2208 .map(|arr| parse_sql_array(arr, self))
2209 .collect::<Result<_, _>>()?;
2210
2211 polars_ensure!(!column_names.is_empty(),
2212 SQLSyntax:
2213 "UNNEST table alias must also declare column names, eg: {} (a,b,c)", alias.name.to_string()
2214 );
2215 if column_names.len() != column_values.len() {
2216 let plural = if column_values.len() > 1 { "s" } else { "" };
2217 polars_bail!(
2218 SQLSyntax:
2219 "UNNEST table alias requires {} column name{}, found {}", column_values.len(), plural, column_names.len()
2220 );
2221 }
2222 let column_series: Vec<Column> = column_values
2223 .into_iter()
2224 .zip(column_names)
2225 .map(|(s, name)| {
2226 if let Some(name) = name {
2227 s.with_name(name)
2228 } else {
2229 s
2230 }
2231 })
2232 .map(Column::from)
2233 .collect();
2234
2235 let lf = DataFrame::new_infer_height(column_series)?.lazy();
2236
2237 if *with_offset || *with_ordinality {
2238 polars_bail!(SQLInterface: "UNNEST tables do not (yet) support WITH ORDINALITY|OFFSET");
2240 }
2241 let table_name = alias.name.value.clone();
2242 self.table_map
2243 .write()
2244 .unwrap()
2245 .insert(table_name.clone(), lf.clone());
2246 Ok((table_name, lf))
2247 } else {
2248 polars_bail!(SQLSyntax: "UNNEST table must have an alias");
2249 }
2250 },
2251 TableFactor::NestedJoin {
2252 table_with_joins,
2253 alias,
2254 } => {
2255 let lf =
2256 self.execute_isolated(|ctx| ctx.execute_from_statement(table_with_joins))?;
2257 match alias {
2258 Some(a) => Ok((a.name.value.clone(), lf)),
2259 None => Ok(("".to_string(), lf)),
2260 }
2261 },
2262 _ => polars_bail!(SQLInterface: "not yet implemented: {}", relation),
2264 }
2265 }
2266
2267 fn execute_table_function(
2268 &mut self,
2269 name: &ObjectName,
2270 alias: &Option<TableAlias>,
2271 args: &[FunctionArg],
2272 ) -> PolarsResult<(String, LazyFrame)> {
2273 let tbl_fn = name.0.first().unwrap().as_ident().unwrap().value.as_str();
2274 let read_fn = tbl_fn.parse::<PolarsTableFunctions>()?;
2275 let (tbl_name, lf) = read_fn.execute(args)?;
2276 #[allow(clippy::useless_asref)]
2277 let tbl_name = alias
2278 .as_ref()
2279 .map(|a| a.name.value.clone())
2280 .unwrap_or_else(|| tbl_name.to_string());
2281
2282 self.table_map
2283 .write()
2284 .unwrap()
2285 .insert(tbl_name.clone(), lf.clone());
2286 Ok((tbl_name, lf))
2287 }
2288
2289 fn process_order_by(
2290 &mut self,
2291 mut lf: LazyFrame,
2292 order_by: &Option<OrderBy>,
2293 selected: Option<&[Expr]>,
2294 ) -> PolarsResult<LazyFrame> {
2295 if order_by.as_ref().is_none_or(|ob| match &ob.kind {
2296 OrderByKind::Expressions(exprs) => exprs.is_empty(),
2297 OrderByKind::All(_) => false,
2298 }) {
2299 return Ok(lf);
2300 }
2301 let schema = self.get_frame_schema(&mut lf)?;
2302 let columns_iter = schema.iter_names().map(|e| col(e.clone()));
2303 let (order_by, order_by_all, n_order_cols) = match &order_by.as_ref().unwrap().kind {
2304 OrderByKind::Expressions(exprs) => {
2305 if exprs.len() == 1
2308 && matches!(&exprs[0].expr, SQLExpr::Identifier(ident)
2309 if ident.value.to_uppercase() == "ALL"
2310 && !schema.iter_names().any(|name| name.to_uppercase() == "ALL"))
2311 {
2312 let n_cols = if let Some(selected) = selected {
2314 selected.len()
2315 } else {
2316 schema.len()
2317 };
2318 (vec![], Some(&exprs[0].options), n_cols)
2319 } else {
2320 (exprs.clone(), None, exprs.len())
2321 }
2322 },
2323 OrderByKind::All(opts) => {
2324 let n_cols = if let Some(selected) = selected {
2325 selected.len()
2326 } else {
2327 schema.len()
2328 };
2329 (vec![], Some(opts), n_cols)
2330 },
2331 };
2332 let mut descending = Vec::with_capacity(n_order_cols);
2333 let mut nulls_last = Vec::with_capacity(n_order_cols);
2334 let mut by: Vec<Expr> = Vec::with_capacity(n_order_cols);
2335
2336 if let Some(opts) = order_by_all {
2337 if let Some(selected) = selected {
2338 by.extend(selected.iter().cloned());
2339 } else {
2340 by.extend(columns_iter);
2341 };
2342 let desc_order = !opts.asc.unwrap_or(true);
2343 nulls_last.resize(by.len(), !opts.nulls_first.unwrap_or(desc_order));
2344 descending.resize(by.len(), desc_order);
2345 } else {
2346 let columns = &columns_iter.collect::<Vec<_>>();
2347 for ob in order_by {
2348 let desc_order = !ob.options.asc.unwrap_or(true);
2351 nulls_last.push(!ob.options.nulls_first.unwrap_or(desc_order));
2352 descending.push(desc_order);
2353
2354 by.push(self.expr_or_ordinal(
2356 &ob.expr,
2357 columns,
2358 selected,
2359 Some(&schema),
2360 "ORDER BY",
2361 )?)
2362 }
2363 }
2364 Ok(lf.sort_by_exprs(
2365 &by,
2366 SortMultipleOptions::default()
2367 .with_order_descending_multi(descending)
2368 .with_nulls_last_multi(nulls_last),
2369 ))
2370 }
2371
2372 fn process_group_by(
2373 &mut self,
2374 mut lf: LazyFrame,
2375 group_by_keys: &[Expr],
2376 projections: &[Expr],
2377 having: Option<Expr>,
2378 ) -> PolarsResult<LazyFrame> {
2379 let schema_before = self.get_frame_schema(&mut lf)?;
2380 let group_by_keys_schema =
2381 expressions_to_schema(group_by_keys, &schema_before, |duplicate_name: &str| {
2382 format!("group_by keys contained duplicate output name '{duplicate_name}'")
2383 })?;
2384
2385 let mut aliased_aggregations: PlHashMap<PlSmallStr, PlSmallStr> = PlHashMap::new();
2387 let mut aggregation_projection = Vec::with_capacity(projections.len());
2388 let mut projection_overrides = PlHashMap::with_capacity(projections.len());
2389 let mut projection_aliases = PlHashSet::new();
2390 let mut group_key_aliases = PlHashSet::new();
2391
2392 let group_key_data: Vec<_> = group_by_keys
2395 .iter()
2396 .map(|gk| {
2397 (
2398 strip_outer_alias(gk),
2399 gk.to_field(&schema_before).ok().map(|f| f.name),
2400 )
2401 })
2402 .collect();
2403
2404 let projection_group_key: Vec<Option<PlSmallStr>> = projections
2405 .iter()
2406 .map(|p| {
2407 let p_stripped = strip_outer_alias(p);
2408 group_key_data.iter().find_map(|(gk_stripped, gk_name)| {
2409 (*gk_stripped == p_stripped)
2410 .then(|| gk_name.clone())
2411 .flatten()
2412 })
2413 })
2414 .collect();
2415
2416 for (e, group_key) in projections.iter().zip(&projection_group_key) {
2417 let matches_group_key = group_key.is_some();
2418 let is_non_group_key_expr =
2420 !matches_group_key && expr_reduces_group(e, &group_by_keys_schema);
2421
2422 let mut e_inner = e;
2425 if let Expr::Alias(expr, alias) = e {
2426 if e.clone().meta().is_simple_projection(Some(&schema_before)) {
2427 group_key_aliases.insert(alias.as_ref());
2428 e_inner = expr
2429 } else if let Expr::Function {
2430 function: FunctionExpr::StructExpr(StructFunction::FieldByName(name)),
2431 ..
2432 } = expr.deref()
2433 {
2434 projection_overrides
2435 .insert(alias.as_ref(), col(name.clone()).alias(alias.clone()));
2436 } else if !is_non_group_key_expr && !group_by_keys_schema.contains(alias) {
2437 projection_aliases.insert(alias.as_ref());
2438 }
2439 }
2440 let field = e_inner.to_field(&schema_before)?;
2441 if is_non_group_key_expr {
2442 let mut e = e.clone();
2443 if let Expr::Agg(AggExpr::Implode {
2444 input: expr,
2445 maintain_order: _,
2446 }) = &e
2447 {
2448 e = (**expr).clone();
2449 } else if let Expr::Alias(expr, name) = &e {
2450 if let Expr::Agg(AggExpr::Implode {
2451 input: expr,
2452 maintain_order: _,
2453 }) = expr.as_ref()
2454 {
2455 e = (**expr).clone().alias(name.clone());
2456 }
2457 }
2458 if group_by_keys_schema.get(&field.name).is_some() {
2461 let alias_name = format_pl_smallstr!("__POLARS_AGG_{}", field.name);
2462 e = e.alias(alias_name.clone());
2463 aliased_aggregations.insert(field.name.clone(), alias_name);
2464 }
2465 aggregation_projection.push(e);
2466 } else if !matches_group_key {
2467 if let Expr::Column(_)
2469 | Expr::Function {
2470 function: FunctionExpr::StructExpr(StructFunction::FieldByName(_)),
2471 ..
2472 } = e_inner
2473 {
2474 if !group_by_keys_schema.contains(&field.name) {
2475 polars_bail!(SQLSyntax: "'{}' should participate in the GROUP BY clause or an aggregate function", &field.name);
2476 }
2477 }
2478 }
2479 }
2480
2481 let having = having.map(|having_expr| {
2484 having_expr.map_expr(|e| match &e {
2485 Expr::Column(name) => resolve_select_alias(name, projections, &schema_before)
2486 .map_or(e, |resolved| strip_outer_alias(&resolved)),
2487 _ => e,
2488 })
2489 });
2490
2491 let group_by = lf.group_by(group_by_keys);
2492 let aggregated = match having {
2493 Some(having) => group_by.having(having),
2494 None => group_by,
2495 }
2496 .agg(&aggregation_projection);
2497
2498 let projection_schema =
2499 expressions_to_schema(projections, &schema_before, |duplicate_name: &str| {
2500 format!("group_by aggregations contained duplicate output name '{duplicate_name}'")
2501 })?;
2502
2503 let final_projection = projection_schema
2506 .iter_names()
2507 .zip(projections.iter().zip(&projection_group_key))
2508 .map(|(name, (projection_expr, group_key))| {
2509 if let Some(expr) = projection_overrides.get(name.as_str()) {
2510 expr.clone()
2511 } else if let Some(aliased_name) = aliased_aggregations.get(name) {
2512 col(aliased_name.clone()).alias(name.clone())
2513 } else if let Some(key_name) = group_key {
2514 if key_name == name {
2517 col(name.clone())
2518 } else {
2519 col(key_name.clone()).alias(name.clone())
2520 }
2521 } else if group_by_keys_schema.get(name).is_some()
2522 || projection_aliases.contains(name.as_str())
2523 || group_key_aliases.contains(name.as_str())
2524 {
2525 if has_expr(projection_expr, |e| {
2526 matches!(e, Expr::Agg(_) | Expr::Len | Expr::Over { .. })
2527 }) {
2528 col(name.clone())
2529 } else {
2530 projection_expr.clone()
2531 }
2532 } else {
2533 col(name.clone())
2534 }
2535 })
2536 .collect::<Vec<_>>();
2537
2538 let mut output_projection = final_projection;
2540 for key_name in group_by_keys_schema.iter_names() {
2541 if !projection_schema.contains(key_name) {
2542 output_projection.push(col(key_name.clone()));
2544 } else if group_by_keys.iter().any(|k| is_simple_col_ref(k, key_name)) {
2545 let is_cross_aliased = projection_schema
2547 .iter_names()
2548 .zip(projections.iter())
2549 .any(|(name, p)| name == key_name && !is_simple_col_ref(p, key_name));
2550 if is_cross_aliased {
2551 let internal_name = format_pl_smallstr!("__POLARS_ORIG_{}", key_name);
2553 output_projection.push(col(key_name.clone()).alias(internal_name));
2554 }
2555 }
2556 }
2557 Ok(aggregated.select(&output_projection))
2558 }
2559
2560 fn process_limit_offset(
2561 &self,
2562 lf: LazyFrame,
2563 limit_clause: &Option<LimitClause>,
2564 fetch: &Option<Fetch>,
2565 ) -> PolarsResult<LazyFrame> {
2566 let (limit, offset) = match limit_clause {
2568 Some(LimitClause::LimitOffset {
2569 limit,
2570 offset,
2571 limit_by,
2572 }) => {
2573 if !limit_by.is_empty() {
2574 polars_bail!(SQLSyntax: "`LIMIT <n> BY <exprs>` clause is not supported");
2577 }
2578 (limit.as_ref(), offset.as_ref().map(|o| &o.value))
2579 },
2580 Some(LimitClause::OffsetCommaLimit { offset, limit }) => (Some(limit), Some(offset)),
2581 None => (None, None),
2582 };
2583
2584 let limit = match (fetch, limit) {
2586 (Some(fetch), None) => fetch.quantity.as_ref(),
2587 (Some(_), Some(_)) => {
2588 polars_bail!(SQLSyntax: "cannot use both `LIMIT` and `FETCH` in the same query")
2589 },
2590 (None, limit) => limit,
2591 };
2592
2593 match (offset, limit) {
2595 (
2596 Some(SQLExpr::Value(ValueWithSpan {
2597 value: SQLValue::Number(offset, _),
2598 ..
2599 })),
2600 Some(SQLExpr::Value(ValueWithSpan {
2601 value: SQLValue::Number(limit, _),
2602 ..
2603 })),
2604 ) => Ok(lf.slice(
2605 offset
2606 .parse()
2607 .map_err(|e| polars_err!(SQLInterface: "OFFSET conversion error: {}", e))?,
2608 limit.parse().map_err(
2609 |e| polars_err!(SQLInterface: "LIMIT/FETCH conversion error: {}", e),
2610 )?,
2611 )),
2612 (
2613 Some(SQLExpr::Value(ValueWithSpan {
2614 value: SQLValue::Number(offset, _),
2615 ..
2616 })),
2617 None,
2618 ) => Ok(lf.slice(
2619 offset
2620 .parse()
2621 .map_err(|e| polars_err!(SQLInterface: "OFFSET conversion error: {}", e))?,
2622 IdxSize::MAX,
2623 )),
2624 (
2625 None,
2626 Some(SQLExpr::Value(ValueWithSpan {
2627 value: SQLValue::Number(limit, _),
2628 ..
2629 })),
2630 ) => {
2631 Ok(lf.limit(limit.parse().map_err(
2632 |e| polars_err!(SQLInterface: "LIMIT/FETCH conversion error: {}", e),
2633 )?))
2634 },
2635 (None, None) => Ok(lf),
2636 _ => polars_bail!(
2637 SQLSyntax: "non-numeric arguments for LIMIT/OFFSET/FETCH are not supported",
2638 ),
2639 }
2640 }
2641
2642 fn process_qualified_wildcard(
2643 &mut self,
2644 ObjectName(idents): &ObjectName,
2645 options: &WildcardAdditionalOptions,
2646 modifiers: &mut SelectModifiers,
2647 schema: Option<&Schema>,
2648 ) -> PolarsResult<Vec<Expr>> {
2649 let mut idents_with_wildcard: Vec<Ident> = idents
2650 .iter()
2651 .filter_map(|p| p.as_ident().cloned())
2652 .collect();
2653 idents_with_wildcard.push(Ident::new("*"));
2654
2655 let exprs = resolve_compound_identifier(self, &idents_with_wildcard, schema)?;
2656 self.process_wildcard_additional_options(exprs, options, modifiers, schema)
2657 }
2658
2659 fn process_wildcard_additional_options(
2660 &mut self,
2661 exprs: Vec<Expr>,
2662 options: &WildcardAdditionalOptions,
2663 modifiers: &mut SelectModifiers,
2664 schema: Option<&Schema>,
2665 ) -> PolarsResult<Vec<Expr>> {
2666 if options.opt_except.is_some() && options.opt_exclude.is_some() {
2667 polars_bail!(SQLInterface: "EXCLUDE and EXCEPT wildcard options cannot be used together (prefer EXCLUDE)")
2668 } else if options.opt_exclude.is_some() && options.opt_ilike.is_some() {
2669 polars_bail!(SQLInterface: "EXCLUDE and ILIKE wildcard options cannot be used together")
2670 }
2671
2672 if let Some(items) = &options.opt_exclude {
2674 match items {
2675 ExcludeSelectItem::Single(name) => {
2676 modifiers.exclude.insert(object_name_to_string(name));
2677 },
2678 ExcludeSelectItem::Multiple(names) => {
2679 modifiers
2680 .exclude
2681 .extend(names.iter().map(object_name_to_string));
2682 },
2683 };
2684 }
2685
2686 if let Some(items) = &options.opt_except {
2688 modifiers.exclude.insert(items.first_element.value.clone());
2689 modifiers
2690 .exclude
2691 .extend(items.additional_elements.iter().map(|i| i.value.clone()));
2692 }
2693
2694 if let Some(item) = &options.opt_ilike {
2696 let rx = regex::escape(item.pattern.as_str())
2697 .replace('%', ".*")
2698 .replace('_', ".");
2699
2700 modifiers.ilike = Some(
2701 polars_utils::regex_cache::compile_regex(format!("^(?is){rx}$").as_str()).unwrap(),
2702 );
2703 }
2704
2705 if let Some(items) = &options.opt_rename {
2707 let renames = match items {
2708 RenameSelectItem::Single(rename) => std::slice::from_ref(rename),
2709 RenameSelectItem::Multiple(renames) => renames.as_slice(),
2710 };
2711 for rn in renames {
2712 let before = PlSmallStr::from_str(rn.ident.value.as_str());
2713 let after = PlSmallStr::from_str(rn.alias.value.as_str());
2714 if before != after {
2715 modifiers.rename.insert(before, after);
2716 }
2717 }
2718 }
2719
2720 if let Some(replacements) = &options.opt_replace {
2722 for rp in &replacements.items {
2723 let replacement_expr = parse_sql_expr(&rp.expr, self, schema);
2724 modifiers
2725 .replace
2726 .push(replacement_expr?.alias(rp.column_name.value.as_str()));
2727 }
2728 }
2729 Ok(exprs)
2730 }
2731
2732 fn rename_columns_from_table_alias(
2733 &mut self,
2734 mut lf: LazyFrame,
2735 alias: &TableAlias,
2736 ) -> PolarsResult<LazyFrame> {
2737 if alias.columns.is_empty() {
2738 Ok(lf)
2739 } else {
2740 let schema = self.get_frame_schema(&mut lf)?;
2741 if alias.columns.len() != schema.len() {
2742 polars_bail!(
2743 SQLSyntax: "number of columns ({}) in alias '{}' does not match the number of columns in the table/query ({})",
2744 alias.columns.len(), alias.name.value, schema.len()
2745 )
2746 } else {
2747 let existing_columns: Vec<_> = schema.iter_names().collect();
2748 let new_columns: Vec<_> =
2749 alias.columns.iter().map(|c| c.name.value.clone()).collect();
2750 Ok(lf.rename(existing_columns, new_columns, true))
2751 }
2752 }
2753 }
2754}
2755
2756impl SQLContext {
2757 pub fn new_from_table_map(table_map: PlHashMap<String, LazyFrame>) -> Self {
2759 Self {
2760 table_map: Arc::new(RwLock::new(table_map)),
2761 ..Default::default()
2762 }
2763 }
2764}
2765
2766fn expand_exprs(expr: Expr, schema: &SchemaRef) -> Vec<Expr> {
2767 match expr {
2768 Expr::Column(nm) if is_regex_colname(nm.as_str()) => {
2769 let re = polars_utils::regex_cache::compile_regex(&nm).unwrap();
2770 schema
2771 .iter_names()
2772 .filter(|name| re.is_match(name))
2773 .map(|name| col(name.clone()))
2774 .collect::<Vec<_>>()
2775 },
2776 Expr::Selector(s) => s
2777 .into_columns(schema, &Default::default())
2778 .unwrap()
2779 .into_iter()
2780 .map(col)
2781 .collect::<Vec<_>>(),
2782 _ => vec![expr],
2783 }
2784}
2785
2786fn is_regex_colname(nm: &str) -> bool {
2787 nm.starts_with('^') && nm.ends_with('$')
2788}
2789
2790fn object_name_to_string(name: &ObjectName) -> String {
2792 name.0
2793 .iter()
2794 .filter_map(|p| p.as_ident())
2795 .map(|i| i.value.as_str())
2796 .collect::<Vec<_>>()
2797 .join(".")
2798}
2799
2800fn get_using_cols(op: &JoinOperator) -> Option<impl Iterator<Item = String> + '_> {
2802 use JoinOperator::*;
2803 match op {
2804 Join(JoinConstraint::Using(cols))
2805 | Inner(JoinConstraint::Using(cols))
2806 | Left(JoinConstraint::Using(cols))
2807 | LeftOuter(JoinConstraint::Using(cols))
2808 | Right(JoinConstraint::Using(cols))
2809 | RightOuter(JoinConstraint::Using(cols))
2810 | FullOuter(JoinConstraint::Using(cols))
2811 | Semi(JoinConstraint::Using(cols))
2812 | Anti(JoinConstraint::Using(cols))
2813 | LeftSemi(JoinConstraint::Using(cols))
2814 | LeftAnti(JoinConstraint::Using(cols))
2815 | RightSemi(JoinConstraint::Using(cols))
2816 | RightAnti(JoinConstraint::Using(cols)) => Some(cols.iter().filter_map(|c| {
2817 c.0.first()
2818 .and_then(|p| p.as_ident())
2819 .map(|i| i.value.clone())
2820 })),
2821 _ => None,
2822 }
2823}
2824
2825pub(crate) fn get_table_name(factor: &TableFactor) -> Option<String> {
2827 match factor {
2828 TableFactor::Table { name, alias, .. } => {
2829 alias.as_ref().map(|a| a.name.value.clone()).or_else(|| {
2830 name.0
2831 .last()
2832 .and_then(|p| p.as_ident())
2833 .map(|i| i.value.clone())
2834 })
2835 },
2836 TableFactor::Derived { alias, .. }
2837 | TableFactor::NestedJoin { alias, .. }
2838 | TableFactor::TableFunction { alias, .. } => alias.as_ref().map(|a| a.name.value.clone()),
2839 _ => None,
2840 }
2841}
2842
2843fn is_simple_col_ref(expr: &Expr, col_name: &PlSmallStr) -> bool {
2845 match expr {
2846 Expr::Column(n) => n == col_name,
2847 Expr::Alias(inner, _) => matches!(inner.as_ref(), Expr::Column(n) if n == col_name),
2848 _ => false,
2849 }
2850}
2851
2852fn strip_outer_alias(expr: &Expr) -> Expr {
2854 if let Expr::Alias(inner, _) = expr {
2855 inner.as_ref().clone()
2856 } else {
2857 expr.clone()
2858 }
2859}
2860
2861fn resolve_select_alias(name: &str, projections: &[Expr], schema: &Schema) -> Option<Expr> {
2866 if schema.contains(name) {
2868 return None;
2869 }
2870 projections.iter().find_map(|p| match p {
2872 Expr::Alias(inner, alias) if alias.as_str() == name => {
2873 Some(inner.as_ref().clone().alias(alias.clone()))
2874 },
2875 _ => None,
2876 })
2877}
2878
2879fn expr_cols_all_in_schema(expr: &Expr, schema: &Schema) -> bool {
2881 let mut found_cols = false;
2882 let mut all_in_schema = true;
2883 for e in expr.into_iter() {
2884 if let Expr::Column(name) = e {
2885 found_cols = true;
2886 if !schema.contains(name.as_str()) {
2887 all_in_schema = false;
2888 break;
2889 }
2890 }
2891 }
2892 found_cols && all_in_schema
2893}
2894
2895fn determine_left_right_join_on(
2903 ctx: &mut SQLContext,
2904 expr_left: &SQLExpr,
2905 expr_right: &SQLExpr,
2906 tbl_left: &TableInfo,
2907 tbl_right: &TableInfo,
2908 join_schema: &Schema,
2909) -> PolarsResult<(Vec<Expr>, Vec<Expr>)> {
2910 let left_on = match parse_sql_expr(expr_left, ctx, Some(join_schema))? {
2913 Expr::Alias(inner, _) => Arc::unwrap_or_clone(inner),
2914 e => e,
2915 };
2916 let right_on = match parse_sql_expr(expr_right, ctx, Some(join_schema))? {
2917 Expr::Alias(inner, _) => Arc::unwrap_or_clone(inner),
2918 e => e,
2919 };
2920
2921 let left_refs = (
2925 expr_refers_to_table(expr_left, &tbl_left.name),
2926 expr_refers_to_table(expr_left, &tbl_right.name),
2927 );
2928 let right_refs = (
2929 expr_refers_to_table(expr_right, &tbl_left.name),
2930 expr_refers_to_table(expr_right, &tbl_right.name),
2931 );
2932 match (left_refs, right_refs) {
2934 ((true, false), (false, true)) => return Ok((vec![left_on], vec![right_on])),
2936 ((false, true), (true, false)) => return Ok((vec![right_on], vec![left_on])),
2938 ((true, true), _) | (_, (true, true)) if tbl_left.name != tbl_right.name => {
2940 polars_bail!(
2941 SQLInterface: "unsupported join condition: {} side references both '{}' and '{}'",
2942 if left_refs.0 && left_refs.1 {
2943 "left"
2944 } else {
2945 "right"
2946 }, tbl_left.name, tbl_right.name
2947 )
2948 },
2949 _ => {},
2951 }
2952
2953 let left_on_cols_in = (
2958 expr_cols_all_in_schema(&left_on, &tbl_left.schema),
2959 expr_cols_all_in_schema(&left_on, &tbl_right.schema),
2960 );
2961 let right_on_cols_in = (
2962 expr_cols_all_in_schema(&right_on, &tbl_left.schema),
2963 expr_cols_all_in_schema(&right_on, &tbl_right.schema),
2964 );
2965 match (left_on_cols_in, right_on_cols_in) {
2966 ((true, false), (false, true)) => Ok((vec![left_on], vec![right_on])),
2968 ((false, true), (true, false)) => Ok((vec![right_on], vec![left_on])),
2969 ((true, true), (true, false)) => Ok((vec![right_on], vec![left_on])),
2971 ((true, true), (false, true)) => Ok((vec![left_on], vec![right_on])),
2972 ((true, false), (true, true)) => Ok((vec![left_on], vec![right_on])),
2973 ((false, true), (true, true)) => Ok((vec![right_on], vec![left_on])),
2974 _ => Ok((vec![left_on], vec![right_on])),
2976 }
2977}
2978
2979fn process_join_on(
2986 ctx: &mut SQLContext,
2987 sql_expr: &SQLExpr,
2988 tbl_left: &TableInfo,
2989 tbl_right: &TableInfo,
2990) -> PolarsResult<(Vec<Expr>, Vec<Expr>, Vec<Expr>)> {
2991 match sql_expr {
2992 SQLExpr::BinaryOp { left, op, right } => match op {
2993 SQLBinaryOperator::And => {
2994 let (mut left_i, mut right_i, mut preds_i) =
2995 process_join_on(ctx, left, tbl_left, tbl_right)?;
2996 let (mut left_j, mut right_j, mut preds_j) =
2997 process_join_on(ctx, right, tbl_left, tbl_right)?;
2998 left_i.append(&mut left_j);
2999 right_i.append(&mut right_j);
3000 preds_i.append(&mut preds_j);
3001 Ok((left_i, right_i, preds_i))
3002 },
3003 SQLBinaryOperator::Eq => {
3004 let join_schema = build_join_schema(tbl_left, tbl_right)?;
3005 let (l, r) = determine_left_right_join_on(
3006 ctx,
3007 left,
3008 right,
3009 tbl_left,
3010 tbl_right,
3011 &join_schema,
3012 )?;
3013 Ok((l, r, vec![]))
3014 },
3015 SQLBinaryOperator::Lt
3016 | SQLBinaryOperator::LtEq
3017 | SQLBinaryOperator::Gt
3018 | SQLBinaryOperator::GtEq
3019 | SQLBinaryOperator::NotEq => {
3020 let join_schema = build_join_schema(tbl_left, tbl_right)?;
3021 let suffix = format!(":{}", tbl_right.name);
3022
3023 let lhs = suffix_if_right_table(
3026 parse_sql_expr(left, ctx, Some(&join_schema))?,
3027 left,
3028 tbl_left,
3029 tbl_right,
3030 &suffix,
3031 );
3032 let rhs = suffix_if_right_table(
3033 parse_sql_expr(right, ctx, Some(&join_schema))?,
3034 right,
3035 tbl_left,
3036 tbl_right,
3037 &suffix,
3038 );
3039
3040 let polars_op = match op {
3041 SQLBinaryOperator::Lt => Operator::Lt,
3042 SQLBinaryOperator::LtEq => Operator::LtEq,
3043 SQLBinaryOperator::Gt => Operator::Gt,
3044 SQLBinaryOperator::GtEq => Operator::GtEq,
3045 SQLBinaryOperator::NotEq => Operator::NotEq,
3046 _ => unreachable!(),
3047 };
3048 let predicate = Expr::BinaryExpr {
3049 left: Arc::new(lhs),
3050 op: polars_op,
3051 right: Arc::new(rhs),
3052 };
3053 Ok((vec![], vec![], vec![predicate]))
3054 },
3055 _ => polars_bail!(
3056 SQLInterface: "unsupported join constraint operator '{:?}'", op
3057 ),
3058 },
3059 SQLExpr::Nested(expr) => process_join_on(ctx, expr, tbl_left, tbl_right),
3060 _ => polars_bail!(
3061 SQLInterface: "unsupported join constraint expression: {:?}", sql_expr
3062 ),
3063 }
3064}
3065
3066fn expr_reduces_group(expr: &Expr, group_by_keys_schema: &Schema) -> bool {
3069 has_expr(expr, |e| match e {
3070 Expr::Agg(_) | Expr::Len | Expr::Over { .. } => true,
3071 #[cfg(feature = "dynamic_group_by")]
3072 Expr::Rolling { .. } => true,
3073 Expr::AnonymousFunction { options, .. } => options.returns_scalar(),
3074 Expr::Function { function: func, .. } if !matches!(func, FunctionExpr::StructExpr(_)) => {
3075 has_expr(
3077 e,
3078 |e| matches!(e, Expr::Column(name) if !group_by_keys_schema.contains(name)),
3079 )
3080 },
3081 _ => false,
3082 })
3083}
3084
3085fn build_join_schema(tbl_left: &TableInfo, tbl_right: &TableInfo) -> PolarsResult<Schema> {
3088 let mut join_schema = Schema::with_capacity(tbl_left.schema.len() + tbl_right.schema.len());
3089 for (name, dtype) in tbl_left.schema.iter() {
3090 join_schema.insert_at_index(join_schema.len(), name.clone(), dtype.clone())?;
3091 }
3092 for (name, dtype) in tbl_right.schema.iter() {
3093 if !join_schema.contains(name) {
3094 join_schema.insert_at_index(join_schema.len(), name.clone(), dtype.clone())?;
3095 }
3096 }
3097 Ok(join_schema)
3098}
3099
3100fn suffix_conflicting_columns(
3103 expr: Expr,
3104 tbl_left: &TableInfo,
3105 tbl_right: &TableInfo,
3106 suffix: &str,
3107) -> Expr {
3108 expr.map_expr(|e| match e {
3109 Expr::Column(ref name)
3110 if tbl_left.schema.contains(name.as_str())
3111 && tbl_right.schema.contains(name.as_str()) =>
3112 {
3113 Expr::Column(PlSmallStr::from_string(format!("{name}{suffix}")))
3114 },
3115 other => other,
3116 })
3117}
3118
3119fn suffix_if_right_table(
3122 expr: Expr,
3123 sql_expr: &SQLExpr,
3124 tbl_left: &TableInfo,
3125 tbl_right: &TableInfo,
3126 suffix: &str,
3127) -> Expr {
3128 let expr = match expr {
3130 Expr::Alias(inner, _) => Arc::unwrap_or_clone(inner),
3131 e => e,
3132 };
3133
3134 let refs_left = expr_refers_to_table(sql_expr, &tbl_left.name);
3135 let refs_right = expr_refers_to_table(sql_expr, &tbl_right.name);
3136
3137 let is_right = if refs_right && !refs_left {
3138 true
3139 } else if refs_left {
3140 false
3141 } else {
3142 !expr_cols_all_in_schema(&expr, &tbl_left.schema)
3144 && expr_cols_all_in_schema(&expr, &tbl_right.schema)
3145 };
3146
3147 if is_right {
3148 suffix_conflicting_columns(expr, tbl_left, tbl_right, suffix)
3149 } else {
3150 expr
3151 }
3152}
3153
3154fn process_join_constraint(
3155 constraint: &JoinConstraint,
3156 tbl_left: &TableInfo,
3157 tbl_right: &TableInfo,
3158 ctx: &mut SQLContext,
3159) -> PolarsResult<(Vec<Expr>, Vec<Expr>, Vec<Expr>)> {
3160 match constraint {
3161 JoinConstraint::On(expr @ SQLExpr::BinaryOp { .. }) => {
3162 process_join_on(ctx, expr, tbl_left, tbl_right)
3163 },
3164 JoinConstraint::Using(idents) if !idents.is_empty() => {
3165 let using: Vec<Expr> = idents
3166 .iter()
3167 .map(|ObjectName(parts)| {
3168 if parts.len() != 1 {
3169 polars_bail!(SQLSyntax: "JOIN \"USING\" clause expects simple column names, not qualified names");
3170 }
3171 match parts[0].as_ident() {
3172 Some(ident) => Ok(col(ident.value.as_str())),
3173 None => polars_bail!(SQLSyntax: "JOIN \"USING\" clause expects identifiers, not functions"),
3174 }
3175 })
3176 .collect::<PolarsResult<Vec<_>>>()?;
3177 Ok((using.clone(), using, vec![]))
3178 },
3179 JoinConstraint::Natural => {
3180 let left_names = tbl_left.schema.iter_names().collect::<PlHashSet<_>>();
3181 let right_names = tbl_right.schema.iter_names().collect::<PlHashSet<_>>();
3182 let on: Vec<Expr> = left_names
3183 .intersection(&right_names)
3184 .map(|&name| col(name.clone()))
3185 .collect();
3186 if on.is_empty() {
3187 polars_bail!(SQLInterface: "no common columns found for NATURAL JOIN")
3188 }
3189 Ok((on.clone(), on, vec![]))
3190 },
3191 _ => polars_bail!(SQLInterface: "unsupported SQL join constraint:\n{:?}", constraint),
3192 }
3193}
3194
3195fn flatten_and_conditions(expr: &SQLExpr) -> Vec<&SQLExpr> {
3197 match expr {
3198 SQLExpr::BinaryOp {
3199 left,
3200 op: SQLBinaryOperator::And,
3201 right,
3202 } => {
3203 let mut conditions = flatten_and_conditions(left);
3204 conditions.extend(flatten_and_conditions(right));
3205 conditions
3206 },
3207 SQLExpr::Nested(inner) => flatten_and_conditions(inner),
3208 _ => vec![expr],
3209 }
3210}
3211
3212fn combine_and_conditions(conditions: Vec<SQLExpr>) -> Option<SQLExpr> {
3214 conditions
3215 .into_iter()
3216 .reduce(|left, right| SQLExpr::BinaryOp {
3217 left: Box::new(left),
3218 op: SQLBinaryOperator::And,
3219 right: Box::new(right),
3220 })
3221}
3222
3223fn is_join_comparison(expr: &SQLExpr, left_tables: &[String], right_table: &str) -> bool {
3226 if let SQLExpr::BinaryOp {
3227 left,
3228 op:
3229 SQLBinaryOperator::Eq
3230 | SQLBinaryOperator::Lt
3231 | SQLBinaryOperator::LtEq
3232 | SQLBinaryOperator::Gt
3233 | SQLBinaryOperator::GtEq
3234 | SQLBinaryOperator::NotEq,
3235 right,
3236 } = expr
3237 {
3238 let left_refs_right = expr_refers_to_table(left, right_table);
3239 let right_refs_right = expr_refers_to_table(right, right_table);
3240
3241 let left_refs_any_left = left_tables
3242 .iter()
3243 .any(|t| expr_refers_to_table(left, t.as_str()));
3244 let right_refs_any_left = left_tables
3245 .iter()
3246 .any(|t| expr_refers_to_table(right, t.as_str()));
3247
3248 (left_refs_right && right_refs_any_left) || (right_refs_right && left_refs_any_left)
3250 } else {
3251 false
3252 }
3253}
3254
3255fn extract_join_predicates(
3258 where_expr: &Option<SQLExpr>,
3259 left_tables: &[String],
3260 right_table: &str,
3261) -> (Option<SQLExpr>, Option<SQLExpr>) {
3262 let Some(expr) = where_expr else {
3263 return (None, None);
3264 };
3265 let conditions = flatten_and_conditions(expr);
3266 let mut join_conds = Vec::new();
3267 let mut filter_conds = Vec::new();
3268 for cond in conditions {
3269 if is_join_comparison(cond, left_tables, right_table) {
3270 join_conds.push(cond.clone());
3271 } else {
3272 filter_conds.push(cond.clone());
3273 }
3274 }
3275 (
3276 combine_and_conditions(join_conds),
3277 combine_and_conditions(filter_conds),
3278 )
3279}
3280
3281pub fn extract_table_identifiers(
3285 query: &str,
3286 include_schema: bool,
3287 unique: bool,
3288) -> PolarsResult<Vec<String>> {
3289 let mut parser = Parser::new(&GenericDialect);
3290 parser = parser.with_options(ParserOptions {
3291 trailing_commas: true,
3292 ..Default::default()
3293 });
3294 let ast = parser
3295 .try_with_sql(query)
3296 .map_err(to_sql_interface_err)?
3297 .parse_statements()
3298 .map_err(to_sql_interface_err)?;
3299
3300 let mut collector = TableIdentifierCollector {
3301 include_schema,
3302 ..Default::default()
3303 };
3304 for stmt in &ast {
3305 let _ = stmt.visit(&mut collector);
3306 }
3307 Ok(if unique {
3308 collector
3309 .tables
3310 .into_iter()
3311 .collect::<PlIndexSet<_>>()
3312 .into_iter()
3313 .collect()
3314 } else {
3315 collector.tables
3316 })
3317}
3318
3319bitflags::bitflags! {
3320 #[derive(PartialEq)]
3325 struct ExprSqlProjectionHeightBehavior: u8 {
3326 const MaintainsColumn = 1 << 0;
3328 const Independent = 1 << 1;
3332 const InheritsContext = 1 << 2;
3335 }
3336}
3337
3338impl ExprSqlProjectionHeightBehavior {
3339 fn identify_from_expr(expr: &Expr) -> Self {
3340 let mut has_column = false;
3341 let mut has_independent = false;
3342
3343 for e in expr.into_iter() {
3344 use Expr::*;
3345 has_column |= matches!(e, Column(_) | Selector(_));
3346 has_independent |= match e {
3347 AnonymousFunction { options, .. } => {
3349 options.returns_scalar() || !options.is_length_preserving()
3350 },
3351 Literal(v) => !v.is_scalar(),
3352 Explode { .. } | Filter { .. } | Gather { .. } | Slice { .. } => true,
3353 Agg { .. } | Len => true,
3354 _ => false,
3355 }
3356 }
3357 if has_independent {
3358 Self::Independent
3359 } else if has_column {
3360 Self::MaintainsColumn
3361 } else {
3362 Self::InheritsContext
3363 }
3364 }
3365}