@@ -25,9 +25,12 @@ use datafusion_common::{
25
25
config:: ConfigOptions , file_options:: file_type:: FileType , not_impl_err, DFSchema ,
26
26
Result , TableReference ,
27
27
} ;
28
- use sqlparser:: ast;
28
+ use sqlparser:: ast:: { self , NullTreatment } ;
29
29
30
- use crate :: { AggregateUDF , Expr , GetFieldAccess , ScalarUDF , TableSource , WindowUDF } ;
30
+ use crate :: {
31
+ AggregateUDF , Expr , GetFieldAccess , ScalarUDF , SortExpr , TableSource , WindowFrame ,
32
+ WindowFunctionDefinition , WindowUDF ,
33
+ } ;
31
34
32
35
/// Provides the `SQL` query planner meta-data about tables and
33
36
/// functions referenced in SQL statements, without a direct dependency on the
@@ -138,7 +141,7 @@ pub trait ExprPlanner: Debug + Send + Sync {
138
141
139
142
/// Plan an array literal, such as `[1, 2, 3]`
140
143
///
141
- /// Returns origin expression arguments if not possible
144
+ /// Returns original expression arguments if not possible
142
145
fn plan_array_literal (
143
146
& self ,
144
147
exprs : Vec < Expr > ,
@@ -149,14 +152,14 @@ pub trait ExprPlanner: Debug + Send + Sync {
149
152
150
153
/// Plan a `POSITION` expression, such as `POSITION(<expr> in <expr>)`
151
154
///
152
- /// returns origin expression arguments if not possible
155
+ /// Returns original expression arguments if not possible
153
156
fn plan_position ( & self , args : Vec < Expr > ) -> Result < PlannerResult < Vec < Expr > > > {
154
157
Ok ( PlannerResult :: Original ( args) )
155
158
}
156
159
157
160
/// Plan a dictionary literal, such as `{ key: value, ...}`
158
161
///
159
- /// Returns origin expression arguments if not possible
162
+ /// Returns original expression arguments if not possible
160
163
fn plan_dictionary_literal (
161
164
& self ,
162
165
expr : RawDictionaryExpr ,
@@ -167,14 +170,14 @@ pub trait ExprPlanner: Debug + Send + Sync {
167
170
168
171
/// Plan an extract expression, such as`EXTRACT(month FROM foo)`
169
172
///
170
- /// Returns origin expression arguments if not possible
173
+ /// Returns original expression arguments if not possible
171
174
fn plan_extract ( & self , args : Vec < Expr > ) -> Result < PlannerResult < Vec < Expr > > > {
172
175
Ok ( PlannerResult :: Original ( args) )
173
176
}
174
177
175
178
/// Plan an substring expression, such as `SUBSTRING(<expr> [FROM <expr>] [FOR <expr>])`
176
179
///
177
- /// Returns origin expression arguments if not possible
180
+ /// Returns original expression arguments if not possible
178
181
fn plan_substring ( & self , args : Vec < Expr > ) -> Result < PlannerResult < Vec < Expr > > > {
179
182
Ok ( PlannerResult :: Original ( args) )
180
183
}
@@ -195,14 +198,14 @@ pub trait ExprPlanner: Debug + Send + Sync {
195
198
196
199
/// Plans an overlay expression, such as `overlay(str PLACING substr FROM pos [FOR count])`
197
200
///
198
- /// Returns origin expression arguments if not possible
201
+ /// Returns original expression arguments if not possible
199
202
fn plan_overlay ( & self , args : Vec < Expr > ) -> Result < PlannerResult < Vec < Expr > > > {
200
203
Ok ( PlannerResult :: Original ( args) )
201
204
}
202
205
203
206
/// Plans a `make_map` expression, such as `make_map(key1, value1, key2, value2, ...)`
204
207
///
205
- /// Returns origin expression arguments if not possible
208
+ /// Returns original expression arguments if not possible
206
209
fn plan_make_map ( & self , args : Vec < Expr > ) -> Result < PlannerResult < Vec < Expr > > > {
207
210
Ok ( PlannerResult :: Original ( args) )
208
211
}
@@ -230,6 +233,23 @@ pub trait ExprPlanner: Debug + Send + Sync {
230
233
fn plan_any ( & self , expr : RawBinaryExpr ) -> Result < PlannerResult < RawBinaryExpr > > {
231
234
Ok ( PlannerResult :: Original ( expr) )
232
235
}
236
+
237
+ /// Plans aggregate functions, such as `COUNT(<expr>)`
238
+ ///
239
+ /// Returns original expression arguments if not possible
240
+ fn plan_aggregate (
241
+ & self ,
242
+ expr : RawAggregateExpr ,
243
+ ) -> Result < PlannerResult < RawAggregateExpr > > {
244
+ Ok ( PlannerResult :: Original ( expr) )
245
+ }
246
+
247
+ /// Plans window functions, such as `COUNT(<expr>)`
248
+ ///
249
+ /// Returns original expression arguments if not possible
250
+ fn plan_window ( & self , expr : RawWindowExpr ) -> Result < PlannerResult < RawWindowExpr > > {
251
+ Ok ( PlannerResult :: Original ( expr) )
252
+ }
233
253
}
234
254
235
255
/// An operator with two arguments to plan
@@ -266,6 +286,30 @@ pub struct RawDictionaryExpr {
266
286
pub values : Vec < Expr > ,
267
287
}
268
288
289
+ /// This structure is used by `AggregateFunctionPlanner` to plan operators with
290
+ /// custom expressions.
291
+ #[ derive( Debug , Clone ) ]
292
+ pub struct RawAggregateExpr {
293
+ pub func : Arc < AggregateUDF > ,
294
+ pub args : Vec < Expr > ,
295
+ pub distinct : bool ,
296
+ pub filter : Option < Box < Expr > > ,
297
+ pub order_by : Option < Vec < SortExpr > > ,
298
+ pub null_treatment : Option < NullTreatment > ,
299
+ }
300
+
301
+ /// This structure is used by `WindowFunctionPlanner` to plan operators with
302
+ /// custom expressions.
303
+ #[ derive( Debug , Clone ) ]
304
+ pub struct RawWindowExpr {
305
+ pub func_def : WindowFunctionDefinition ,
306
+ pub args : Vec < Expr > ,
307
+ pub partition_by : Vec < Expr > ,
308
+ pub order_by : Vec < SortExpr > ,
309
+ pub window_frame : WindowFrame ,
310
+ pub null_treatment : Option < NullTreatment > ,
311
+ }
312
+
269
313
/// Result of planning a raw expr with [`ExprPlanner`]
270
314
#[ derive( Debug , Clone ) ]
271
315
pub enum PlannerResult < T > {
0 commit comments