refactor: ordery by all

Author: PokIsemaine

datafusion/expr/src/expr.rs

@@ -39,7 +39,7 @@ use datafusion_common::{
 use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 use sqlparser::ast::{
     display_comma_separated, ExceptSelectItem, ExcludeSelectItem, IlikeSelectItem,
-    NullTreatment, OrderByExpr, OrderByOptions, RenameSelectItem, ReplaceSelectElement,
+    NullTreatment, RenameSelectItem, ReplaceSelectElement,
 };
 
 /// Represents logical expressions such as `A + 1`, or `CAST(c1 AS int)`.
@@ -701,24 +701,6 @@ impl TryCast {
     }
 }
 
-/// OrderBy Expressions
-pub enum OrderByExprs {
-    OrderByExprVec(Vec<OrderByExpr>),
-    All {
-        exprs: Vec<Expr>,
-        options: OrderByOptions,
-    },
-}
-
-impl OrderByExprs {
-    pub fn is_empty(&self) -> bool {
-        match self {
-            OrderByExprs::OrderByExprVec(exprs) => exprs.is_empty(),
-            OrderByExprs::All { exprs, .. } => exprs.is_empty(),
-        }
-    }
-}
-
 /// SORT expression
 #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
 pub struct Sort {
@@ -1765,38 +1747,23 @@ impl Expr {
     pub fn infer_placeholder_types(self, schema: &DFSchema) -> Result<(Expr, bool)> {
         let mut has_placeholder = false;
         self.transform(|mut expr| {
-            match &mut expr {
-                // Default to assuming the arguments are the same type
-                Expr::BinaryExpr(BinaryExpr { left, op: _, right }) => {
-                    rewrite_placeholder(left.as_mut(), right.as_ref(), schema)?;
-                    rewrite_placeholder(right.as_mut(), left.as_ref(), schema)?;
-                }
-                Expr::Between(Between {
-                    expr,
-                    negated: _,
-                    low,
-                    high,
-                }) => {
-                    rewrite_placeholder(low.as_mut(), expr.as_ref(), schema)?;
-                    rewrite_placeholder(high.as_mut(), expr.as_ref(), schema)?;
-                }
-                Expr::InList(InList {
-                    expr,
-                    list,
-                    negated: _,
-                }) => {
-                    for item in list.iter_mut() {
-                        rewrite_placeholder(item, expr.as_ref(), schema)?;
-                    }
-                }
-                Expr::Like(Like { expr, pattern, .. })
-                | Expr::SimilarTo(Like { expr, pattern, .. }) => {
-                    rewrite_placeholder(pattern.as_mut(), expr.as_ref(), schema)?;
-                }
-                Expr::Placeholder(_) => {
-                    has_placeholder = true;
-                }
-                _ => {}
+            // Default to assuming the arguments are the same type
+            if let Expr::BinaryExpr(BinaryExpr { left, op: _, right }) = &mut expr {
+                rewrite_placeholder(left.as_mut(), right.as_ref(), schema)?;
+                rewrite_placeholder(right.as_mut(), left.as_ref(), schema)?;
+            };
+            if let Expr::Between(Between {
+                expr,
+                negated: _,
+                low,
+                high,
+            }) = &mut expr
+            {
+                rewrite_placeholder(low.as_mut(), expr.as_ref(), schema)?;
+                rewrite_placeholder(high.as_mut(), expr.as_ref(), schema)?;
+            }
+            if let Expr::Placeholder(_) = &expr {
+                has_placeholder = true;
             }
             Ok(Transformed::yes(expr))
         })
@@ -3218,117 +3185,10 @@ mod test {
         case, lit, qualified_wildcard, wildcard, wildcard_with_options, ColumnarValue,
         ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Volatility,
     };
-    use arrow::datatypes::{Field, Schema};
     use sqlparser::ast;
     use sqlparser::ast::{Ident, IdentWithAlias};
     use std::any::Any;
 
-    #[test]
-    fn infer_placeholder_in_clause() {
-        // SELECT * FROM employees WHERE department_id IN ($1, $2, $3);
-        let column = col("department_id");
-        let param_placeholders = vec![
-            Expr::Placeholder(Placeholder {
-                id: "$1".to_string(),
-                data_type: None,
-            }),
-            Expr::Placeholder(Placeholder {
-                id: "$2".to_string(),
-                data_type: None,
-            }),
-            Expr::Placeholder(Placeholder {
-                id: "$3".to_string(),
-                data_type: None,
-            }),
-        ];
-        let in_list = Expr::InList(InList {
-            expr: Box::new(column),
-            list: param_placeholders,
-            negated: false,
-        });
-
-        let schema = Arc::new(Schema::new(vec![
-            Field::new("name", DataType::Utf8, true),
-            Field::new("department_id", DataType::Int32, true),
-        ]));
-        let df_schema = DFSchema::try_from(schema).unwrap();
-
-        let (inferred_expr, contains_placeholder) =
-            in_list.infer_placeholder_types(&df_schema).unwrap();
-
-        assert!(contains_placeholder);
-
-        match inferred_expr {
-            Expr::InList(in_list) => {
-                for expr in in_list.list {
-                    match expr {
-                        Expr::Placeholder(placeholder) => {
-                            assert_eq!(
-                                placeholder.data_type,
-                                Some(DataType::Int32),
-                                "Placeholder {} should infer Int32",
-                                placeholder.id
-                            );
-                        }
-                        _ => panic!("Expected Placeholder expression"),
-                    }
-                }
-            }
-            _ => panic!("Expected InList expression"),
-        }
-    }
-
-    #[test]
-    fn infer_placeholder_like_and_similar_to() {
-        // name LIKE $1
-        let schema =
-            Arc::new(Schema::new(vec![Field::new("name", DataType::Utf8, true)]));
-        let df_schema = DFSchema::try_from(schema).unwrap();
-
-        let like = Like {
-            expr: Box::new(col("name")),
-            pattern: Box::new(Expr::Placeholder(Placeholder {
-                id: "$1".to_string(),
-                data_type: None,
-            })),
-            negated: false,
-            case_insensitive: false,
-            escape_char: None,
-        };
-
-        let expr = Expr::Like(like.clone());
-
-        let (inferred_expr, _) = expr.infer_placeholder_types(&df_schema).unwrap();
-        match inferred_expr {
-            Expr::Like(like) => match *like.pattern {
-                Expr::Placeholder(placeholder) => {
-                    assert_eq!(placeholder.data_type, Some(DataType::Utf8));
-                }
-                _ => panic!("Expected Placeholder"),
-            },
-            _ => panic!("Expected Like"),
-        }
-
-        // name SIMILAR TO $1
-        let expr = Expr::SimilarTo(like);
-
-        let (inferred_expr, _) = expr.infer_placeholder_types(&df_schema).unwrap();
-        match inferred_expr {
-            Expr::SimilarTo(like) => match *like.pattern {
-                Expr::Placeholder(placeholder) => {
-                    assert_eq!(
-                        placeholder.data_type,
-                        Some(DataType::Utf8),
-                        "Placeholder {} should infer Utf8",
-                        placeholder.id
-                    );
-                }
-                _ => panic!("Expected Placeholder expression"),
-            },
-            _ => panic!("Expected SimilarTo expression"),
-        }
-    }
-
     #[test]
     #[allow(deprecated)]
     fn format_case_when() -> Result<()> {

datafusion/sql/src/expr/function.rs

@@ -22,7 +22,7 @@ use datafusion_common::{
     internal_datafusion_err, internal_err, not_impl_err, plan_datafusion_err, plan_err,
     DFSchema, Dependency, Diagnostic, Result, Span,
 };
-use datafusion_expr::expr::{OrderByExprs, ScalarFunction, Unnest, WildcardOptions};
+use datafusion_expr::expr::{ScalarFunction, Unnest, WildcardOptions};
 use datafusion_expr::planner::{PlannerResult, RawAggregateExpr, RawWindowExpr};
 use datafusion_expr::{
     expr, Expr, ExprFunctionExt, ExprSchemable, WindowFrame, WindowFunctionDefinition,
@@ -291,7 +291,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                 .map(|e| self.sql_expr_to_logical_expr(e, schema, planner_context))
                 .collect::<Result<Vec<_>>>()?;
             let mut order_by = self.order_by_to_sort_expr(
-                OrderByExprs::OrderByExprVec(window.order_by),
+                window.order_by,
                 schema,
                 planner_context,
                 // Numeric literals in window function ORDER BY are treated as constants
@@ -405,7 +405,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                     (!within_group.is_empty()).then_some(within_group)
                 } else {
                     let order_by = self.order_by_to_sort_expr(
-                        OrderByExprs::OrderByExprVec(order_by),
+                        order_by,
                         schema,
                         planner_context,
                         true,

datafusion/sql/src/expr/order_by.rs

@@ -19,7 +19,7 @@ use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
 use datafusion_common::{
     not_impl_err, plan_datafusion_err, plan_err, Column, DFSchema, Result,
 };
-use datafusion_expr::expr::{OrderByExprs, Sort};
+use datafusion_expr::expr::Sort;
 use datafusion_expr::{Expr, SortExpr};
 use sqlparser::ast::{
     Expr as SQLExpr, OrderByExpr, OrderByOptions, Value, ValueWithSpan,
@@ -41,7 +41,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
     /// If false, interpret numeric literals as constant values.
     pub(crate) fn order_by_to_sort_expr(
         &self,
-        order_by_exprs: OrderByExprs,
+        order_by_exprs: Vec<OrderByExpr>,
         input_schema: &DFSchema,
         planner_context: &mut PlannerContext,
         literal_to_column: bool,
@@ -72,62 +72,48 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                 Sort::new(expr, asc, nulls_first)
             };
 
-        match order_by_exprs {
-            OrderByExprs::OrderByExprVec(expressions) => {
-                for e in expressions {
-                    let OrderByExpr {
-                        expr,
-                        options: OrderByOptions { asc, nulls_first },
-                        with_fill,
-                    } = e;
+        for order_by_expr in order_by_exprs {
+            let OrderByExpr {
+                expr,
+                options: OrderByOptions { asc, nulls_first },
+                with_fill,
+            } = order_by_expr;
 
-                    if let Some(with_fill) = with_fill {
-                        return not_impl_err!(
-                            "ORDER BY WITH FILL is not supported: {with_fill}"
-                        );
-                    }
+            if let Some(with_fill) = with_fill {
+                return not_impl_err!("ORDER BY WITH FILL is not supported: {with_fill}");
+            }
 
-                    let expr = match expr {
-                        SQLExpr::Value(ValueWithSpan {
-                            value: Value::Number(v, _),
-                            span: _,
-                        }) if literal_to_column => {
-                            let field_index = v
-                                .parse::<usize>()
-                                .map_err(|err| plan_datafusion_err!("{}", err))?;
+            let expr = match expr {
+                SQLExpr::Value(ValueWithSpan {
+                    value: Value::Number(v, _),
+                    span: _,
+                }) if literal_to_column => {
+                    let field_index = v
+                        .parse::<usize>()
+                        .map_err(|err| plan_datafusion_err!("{}", err))?;
 
-                            if field_index == 0 {
-                                return plan_err!(
-                                    "Order by index starts at 1 for column indexes"
-                                );
-                            } else if input_schema.fields().len() < field_index {
-                                return plan_err!(
-                                    "Order by column out of bounds, specified: {}, max: {}",
-                                    field_index,
-                                    input_schema.fields().len()
-                                );
-                            }
+                    if field_index == 0 {
+                        return plan_err!(
+                            "Order by index starts at 1 for column indexes"
+                        );
+                    } else if input_schema.fields().len() < field_index {
+                        return plan_err!(
+                            "Order by column out of bounds, specified: {}, max: {}",
+                            field_index,
+                            input_schema.fields().len()
+                        );
+                    }
 
-                            Expr::Column(Column::from(
-                                input_schema.qualified_field(field_index - 1),
-                            ))
-                        }
-                        e => self.sql_expr_to_logical_expr(
-                            e,
-                            order_by_schema,
-                            planner_context,
-                        )?,
-                    };
-                    sort_expr_vec.push(make_sort_expr(expr, asc, nulls_first));
+                    Expr::Column(Column::from(
+                        input_schema.qualified_field(field_index - 1),
+                    ))
                 }
-            }
-            OrderByExprs::All { exprs, options } => {
-                let OrderByOptions { asc, nulls_first } = options;
-                for expr in exprs {
-                    sort_expr_vec.push(make_sort_expr(expr, asc, nulls_first));
+                e => {
+                    self.sql_expr_to_logical_expr(e, order_by_schema, planner_context)?
                 }
-            }
-        };
+            };
+            sort_expr_vec.push(make_sort_expr(expr, asc, nulls_first));
+        }
 
         Ok(sort_expr_vec)
     }