datafusion_sql/expr/order_by.rs
1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
19use datafusion_common::{
20 not_impl_err, plan_datafusion_err, plan_err, Column, DFSchema, Result,
21};
22use datafusion_expr::expr::Sort;
23use datafusion_expr::{Expr, SortExpr};
24use sqlparser::ast::{
25 Expr as SQLExpr, OrderByExpr, OrderByOptions, Value, ValueWithSpan,
26};
27
28impl<S: ContextProvider> SqlToRel<'_, S> {
29 /// Convert sql [OrderByExpr] to `Vec<Expr>`.
30 ///
31 /// `input_schema` and `additional_schema` are used to resolve column references in the order-by expressions.
32 /// `input_schema` is the schema of the input logical plan, typically derived from the SELECT list.
33 ///
34 /// Usually order-by expressions can only reference the input plan's columns.
35 /// But the `SELECT ... FROM ... ORDER BY ...` syntax is a special case. Besides the input schema,
36 /// it can reference an `additional_schema` derived from the `FROM` clause.
37 ///
38 /// If `literal_to_column` is true, treat any numeric literals (e.g. `2`) as a 1 based index into the
39 /// SELECT list (e.g. `SELECT a, b FROM table ORDER BY 2`). Literals only reference the `input_schema`.
40 ///
41 /// If false, interpret numeric literals as constant values.
42 pub(crate) fn order_by_to_sort_expr(
43 &self,
44 order_by_exprs: Vec<OrderByExpr>,
45 input_schema: &DFSchema,
46 planner_context: &mut PlannerContext,
47 literal_to_column: bool,
48 additional_schema: Option<&DFSchema>,
49 ) -> Result<Vec<SortExpr>> {
50 if order_by_exprs.is_empty() {
51 return Ok(vec![]);
52 }
53
54 let mut combined_schema;
55 let order_by_schema = match additional_schema {
56 Some(schema) => {
57 combined_schema = input_schema.clone();
58 combined_schema.merge(schema);
59 &combined_schema
60 }
61 None => input_schema,
62 };
63
64 let mut sort_expr_vec = Vec::with_capacity(order_by_exprs.len());
65
66 let make_sort_expr = |expr: Expr,
67 asc: Option<bool>,
68 nulls_first: Option<bool>| {
69 let asc = asc.unwrap_or(true);
70 let nulls_first = nulls_first
71 .unwrap_or_else(|| self.options.default_null_ordering.nulls_first(asc));
72 Sort::new(expr, asc, nulls_first)
73 };
74
75 for order_by_expr in order_by_exprs {
76 let OrderByExpr {
77 expr,
78 options: OrderByOptions { asc, nulls_first },
79 with_fill,
80 } = order_by_expr;
81
82 if let Some(with_fill) = with_fill {
83 return not_impl_err!("ORDER BY WITH FILL is not supported: {with_fill}");
84 }
85
86 let expr = match expr {
87 SQLExpr::Value(ValueWithSpan {
88 value: Value::Number(v, _),
89 span: _,
90 }) if literal_to_column => {
91 let field_index = v
92 .parse::<usize>()
93 .map_err(|err| plan_datafusion_err!("{}", err))?;
94
95 if field_index == 0 {
96 return plan_err!(
97 "Order by index starts at 1 for column indexes"
98 );
99 } else if input_schema.fields().len() < field_index {
100 return plan_err!(
101 "Order by column out of bounds, specified: {}, max: {}",
102 field_index,
103 input_schema.fields().len()
104 );
105 }
106
107 Expr::Column(Column::from(
108 input_schema.qualified_field(field_index - 1),
109 ))
110 }
111 e => {
112 self.sql_expr_to_logical_expr(e, order_by_schema, planner_context)?
113 }
114 };
115 sort_expr_vec.push(make_sort_expr(expr, asc, nulls_first));
116 }
117
118 Ok(sort_expr_vec)
119 }
120}