Skip to content

Commit b2d48a2

Browse files
PokIsemainePokIsemaine
committed
refactor: orderyby all
1 parent a7262a3 commit b2d48a2

File tree

5 files changed

+65
-104
lines changed

5 files changed

+65
-104
lines changed

‎datafusion/expr/src/expr.rs‎

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ use datafusion_common::{
3939
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
4040
use sqlparser::ast::{
4141
display_comma_separated, ExceptSelectItem, ExcludeSelectItem, IlikeSelectItem,
42-
NullTreatment, OrderByExpr, OrderByOptions, RenameSelectItem, ReplaceSelectElement,
42+
NullTreatment, RenameSelectItem, ReplaceSelectElement,
4343
};
4444

4545
/// Represents logical expressions such as `A + 1`, or `CAST(c1 AS int)`.
@@ -701,24 +701,6 @@ impl TryCast {
701701
}
702702
}
703703

704-
/// OrderBy Expressions
705-
pub enum OrderByExprs {
706-
OrderByExprVec(Vec<OrderByExpr>),
707-
All {
708-
exprs: Vec<Expr>,
709-
options: OrderByOptions,
710-
},
711-
}
712-
713-
impl OrderByExprs {
714-
pub fn is_empty(&self) -> bool {
715-
match self {
716-
OrderByExprs::OrderByExprVec(exprs) => exprs.is_empty(),
717-
OrderByExprs::All { exprs, .. } => exprs.is_empty(),
718-
}
719-
}
720-
}
721-
722704
/// SORT expression
723705
#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
724706
pub struct Sort {

‎datafusion/sql/src/expr/function.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ use datafusion_common::{
2222
internal_datafusion_err, internal_err, not_impl_err, plan_datafusion_err, plan_err,
2323
DFSchema, Dependency, Diagnostic, Result, Span,
2424
};
25-
use datafusion_expr::expr::{OrderByExprs, ScalarFunction, Unnest, WildcardOptions};
25+
use datafusion_expr::expr::{ScalarFunction, Unnest, WildcardOptions};
2626
use datafusion_expr::planner::{PlannerResult, RawAggregateExpr, RawWindowExpr};
2727
use datafusion_expr::{
2828
expr, Expr, ExprFunctionExt, ExprSchemable, WindowFrame, WindowFunctionDefinition,
@@ -291,7 +291,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
291291
.map(|e| self.sql_expr_to_logical_expr(e, schema, planner_context))
292292
.collect::<Result<Vec<_>>>()?;
293293
let mut order_by = self.order_by_to_sort_expr(
294-
OrderByExprs::OrderByExprVec(window.order_by),
294+
window.order_by,
295295
schema,
296296
planner_context,
297297
// Numeric literals in window function ORDER BY are treated as constants
@@ -405,7 +405,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
405405
(!within_group.is_empty()).then_some(within_group)
406406
} else {
407407
let order_by = self.order_by_to_sort_expr(
408-
OrderByExprs::OrderByExprVec(order_by),
408+
order_by,
409409
schema,
410410
planner_context,
411411
true,

‎datafusion/sql/src/expr/order_by.rs‎

Lines changed: 38 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
1919
use datafusion_common::{
2020
not_impl_err, plan_datafusion_err, plan_err, Column, DFSchema, Result,
2121
};
22-
use datafusion_expr::expr::{OrderByExprs, Sort};
22+
use datafusion_expr::expr::Sort;
2323
use datafusion_expr::{Expr, SortExpr};
2424
use sqlparser::ast::{
2525
Expr as SQLExpr, OrderByExpr, OrderByOptions, Value, ValueWithSpan,
@@ -41,7 +41,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
4141
/// If false, interpret numeric literals as constant values.
4242
pub(crate) fn order_by_to_sort_expr(
4343
&self,
44-
order_by_exprs: OrderByExprs,
44+
order_by_exprs: Vec<OrderByExpr>,
4545
input_schema: &DFSchema,
4646
planner_context: &mut PlannerContext,
4747
literal_to_column: bool,
@@ -72,62 +72,48 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
7272
Sort::new(expr, asc, nulls_first)
7373
};
7474

75-
match order_by_exprs {
76-
OrderByExprs::OrderByExprVec(expressions) => {
77-
for e in expressions {
78-
let OrderByExpr {
79-
expr,
80-
options: OrderByOptions { asc, nulls_first },
81-
with_fill,
82-
} = e;
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;
8381

84-
if let Some(with_fill) = with_fill {
85-
return not_impl_err!(
86-
"ORDER BY WITH FILL is not supported: {with_fill}"
87-
);
88-
}
82+
if let Some(with_fill) = with_fill {
83+
return not_impl_err!("ORDER BY WITH FILL is not supported: {with_fill}");
84+
}
8985

90-
let expr = match expr {
91-
SQLExpr::Value(ValueWithSpan {
92-
value: Value::Number(v, _),
93-
span: _,
94-
}) if literal_to_column => {
95-
let field_index = v
96-
.parse::<usize>()
97-
.map_err(|err| plan_datafusion_err!("{}", err))?;
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))?;
9894

99-
if field_index == 0 {
100-
return plan_err!(
101-
"Order by index starts at 1 for column indexes"
102-
);
103-
} else if input_schema.fields().len() < field_index {
104-
return plan_err!(
105-
"Order by column out of bounds, specified: {}, max: {}",
106-
field_index,
107-
input_schema.fields().len()
108-
);
109-
}
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+
}
110106

111-
Expr::Column(Column::from(
112-
input_schema.qualified_field(field_index - 1),
113-
))
114-
}
115-
e => self.sql_expr_to_logical_expr(
116-
e,
117-
order_by_schema,
118-
planner_context,
119-
)?,
120-
};
121-
sort_expr_vec.push(make_sort_expr(expr, asc, nulls_first));
107+
Expr::Column(Column::from(
108+
input_schema.qualified_field(field_index - 1),
109+
))
122110
}
123-
}
124-
OrderByExprs::All { exprs, options } => {
125-
let OrderByOptions { asc, nulls_first } = options;
126-
for expr in exprs {
127-
sort_expr_vec.push(make_sort_expr(expr, asc, nulls_first));
111+
e => {
112+
self.sql_expr_to_logical_expr(e, order_by_schema, planner_context)?
128113
}
129-
}
130-
};
114+
};
115+
sort_expr_vec.push(make_sort_expr(expr, asc, nulls_first));
116+
}
131117

132118
Ok(sort_expr_vec)
133119
}

‎datafusion/sql/src/query.rs‎

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -21,15 +21,16 @@ use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
2121

2222
use crate::stack::StackGuard;
2323
use datafusion_common::{not_impl_err, Constraints, DFSchema, Result};
24-
use datafusion_expr::expr::{OrderByExprs, Sort};
24+
use datafusion_expr::expr::Sort;
2525

2626
use datafusion_expr::{
2727
CreateMemoryTable, DdlStatement, Distinct, Expr, LogicalPlan, LogicalPlanBuilder,
2828
};
2929
use sqlparser::ast::{
30-
Expr as SQLExpr, Offset as SQLOffset, OrderBy, OrderByKind, Query, SelectInto,
31-
SetExpr,
30+
Expr as SQLExpr, Ident, Offset as SQLOffset, OrderBy, OrderByExpr, OrderByKind,
31+
Query, SelectInto, SetExpr,
3232
};
33+
use sqlparser::tokenizer::Span;
3334

3435
impl<S: ContextProvider> SqlToRel<'_, S> {
3536
/// Generate a logical plan from an SQL query/subquery
@@ -151,46 +152,44 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
151152
}
152153

153154
/// Returns the order by expressions from the query.
154-
fn to_order_by_exprs(order_by: Option<OrderBy>) -> Result<OrderByExprs> {
155+
fn to_order_by_exprs(order_by: Option<OrderBy>) -> Result<Vec<OrderByExpr>> {
155156
to_order_by_exprs_with_select(order_by, None)
156157
}
157158

158159
/// Returns the order by expressions from the query with the select expressions.
159160
pub(crate) fn to_order_by_exprs_with_select(
160161
order_by: Option<OrderBy>,
161162
select_exprs: Option<&Vec<Expr>>,
162-
) -> Result<OrderByExprs> {
163+
) -> Result<Vec<OrderByExpr>> {
163164
let Some(OrderBy { kind, interpolate }) = order_by else {
164165
// If no order by, return an empty array.
165-
return Ok(OrderByExprs::OrderByExprVec(vec![]));
166+
return Ok(vec![]);
166167
};
167168
if let Some(_interpolate) = interpolate {
168169
return not_impl_err!("ORDER BY INTERPOLATE is not supported");
169170
}
170171
match kind {
171172
OrderByKind::All(order_by_options) => {
172173
let Some(exprs) = select_exprs else {
173-
return Ok(OrderByExprs::All {
174-
exprs: vec![],
175-
options: order_by_options,
176-
});
174+
return Ok(vec![]);
177175
};
178-
179-
let order_by_epxrs = exprs
176+
let order_by_exprs = exprs
180177
.iter()
181178
.filter_map(|select_expr| match select_expr {
182-
Expr::Column(_) => Some(select_expr.clone()),
179+
Expr::Column(column) => Some(OrderByExpr {
180+
expr: SQLExpr::Identifier(Ident {
181+
value: column.name.clone(),
182+
quote_style: None,
183+
span: Span::empty(),
184+
}),
185+
options: order_by_options.clone(),
186+
with_fill: None,
187+
}),
183188
_ => None,
184189
})
185190
.collect::<Vec<_>>();
186-
187-
Ok(OrderByExprs::All {
188-
exprs: order_by_epxrs,
189-
options: order_by_options,
190-
})
191-
}
192-
OrderByKind::Expressions(order_by_exprs) => {
193-
Ok(OrderByExprs::OrderByExprVec(order_by_exprs))
191+
Ok(order_by_exprs)
194192
}
193+
OrderByKind::Expressions(order_by_exprs) => Ok(order_by_exprs),
195194
}
196195
}

‎datafusion/sql/src/statement.rs‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@ use datafusion_common::{
3939
ToDFSchema,
4040
};
4141
use datafusion_expr::dml::{CopyTo, InsertOp};
42-
use datafusion_expr::expr::OrderByExprs;
4342
use datafusion_expr::expr_rewriter::normalize_col_with_schemas_and_ambiguity_check;
4443
use datafusion_expr::logical_plan::builder::project;
4544
use datafusion_expr::logical_plan::DdlStatement;
@@ -1241,7 +1240,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
12411240
.to_dfschema_ref()?;
12421241
let using: Option<String> = using.as_ref().map(ident_to_string);
12431242
let columns = self.order_by_to_sort_expr(
1244-
OrderByExprs::OrderByExprVec(columns),
1243+
columns,
12451244
&table_schema,
12461245
planner_context,
12471246
false,
@@ -1424,13 +1423,8 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
14241423
let mut all_results = vec![];
14251424
for expr in order_exprs {
14261425
// Convert each OrderByExpr to a SortExpr:
1427-
let expr_vec = self.order_by_to_sort_expr(
1428-
OrderByExprs::OrderByExprVec(expr),
1429-
schema,
1430-
planner_context,
1431-
true,
1432-
None,
1433-
)?;
1426+
let expr_vec =
1427+
self.order_by_to_sort_expr(expr, schema, planner_context, true, None)?;
14341428
// Verify that columns of all SortExprs exist in the schema:
14351429
for sort in expr_vec.iter() {
14361430
for column in sort.expr.column_refs().iter() {

0 commit comments

Comments
 (0)