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}