1use datafusion::logical_expr::AggregateUDFImpl;
2use datafusion::{arrow, common, error, functions_aggregate, logical_expr};
3use std::fmt;
4use std::ops::Deref;
5
6make_udaf_expr_and_func!(
7 MaxByFunction,
8 max_by,
9 x y,
10 "Returns the value of the first column corresponding to the maximum value in the second column.",
11 max_by_udaf
12);
13
14#[derive(Eq, Hash, PartialEq)]
15pub struct MaxByFunction {
16 null_first: bool,
17 signature: logical_expr::Signature,
18}
19
20impl fmt::Debug for MaxByFunction {
21 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
22 f.debug_struct("MaxBy")
23 .field("name", &self.name())
24 .field("signature", &self.signature)
25 .field("accumulator", &"<FUNC>")
26 .finish()
27 }
28}
29impl Default for MaxByFunction {
30 fn default() -> Self {
31 Self::new(true)
32 }
33}
34
35impl MaxByFunction {
36 pub fn new(null_first: bool) -> Self {
37 Self {
38 null_first,
39 signature: logical_expr::Signature::user_defined(logical_expr::Volatility::Immutable),
40 }
41 }
42}
43
44fn get_min_max_by_result_type(
45 input_types: &[arrow::datatypes::DataType],
46) -> error::Result<Vec<arrow::datatypes::DataType>> {
47 match &input_types[0] {
48 arrow::datatypes::DataType::Dictionary(_, dict_value_type) => {
49 let mut result = vec![dict_value_type.deref().clone()];
51 result.extend_from_slice(&input_types[1..]);
53 Ok(result)
54 }
55 _ => Ok(input_types.to_vec()),
56 }
57}
58
59impl logical_expr::AggregateUDFImpl for MaxByFunction {
60 fn name(&self) -> &str {
61 "max_by"
62 }
63
64 fn signature(&self) -> &logical_expr::Signature {
65 &self.signature
66 }
67
68 fn return_type(
69 &self,
70 arg_types: &[arrow::datatypes::DataType],
71 ) -> error::Result<arrow::datatypes::DataType> {
72 Ok(arg_types[0].to_owned())
73 }
74
75 fn accumulator(
76 &self,
77 _acc_args: logical_expr::function::AccumulatorArgs,
78 ) -> error::Result<Box<dyn logical_expr::Accumulator>> {
79 common::exec_err!("should not reach here")
80 }
81
82 fn coerce_types(
83 &self,
84 arg_types: &[arrow::datatypes::DataType],
85 ) -> error::Result<Vec<arrow::datatypes::DataType>> {
86 get_min_max_by_result_type(arg_types)
87 }
88
89 fn simplify(&self) -> Option<logical_expr::function::AggregateFunctionSimplification> {
90 let null_first = self.null_first;
91 let simplify = move |mut aggr_func: logical_expr::expr::AggregateFunction,
92 _: &logical_expr::simplify::SimplifyContext| {
93 let mut order_by = aggr_func.params.order_by;
94 let (second_arg, first_arg) = (
95 aggr_func.params.args.remove(1),
96 aggr_func.params.args.remove(0),
97 );
98 let sort = logical_expr::expr::Sort::new(second_arg, true, null_first);
99 order_by.push(sort);
100 let func = logical_expr::expr::AggregateFunction::new_udf(
101 functions_aggregate::first_last::last_value_udaf(),
102 vec![first_arg],
103 aggr_func.params.distinct,
104 aggr_func.params.filter,
105 order_by,
106 aggr_func.params.null_treatment,
107 );
108 let func = logical_expr::expr::Expr::AggregateFunction(func);
109 Ok(func)
110 };
111 Some(Box::new(simplify))
112 }
113}
114
115make_udaf_expr_and_func!(
116 MinByFunction,
117 min_by,
118 x y,
119 "Returns the value of the first column corresponding to the minimum value in the second column.",
120 min_by_udaf
121);
122
123#[derive(Eq, Hash, PartialEq)]
124pub struct MinByFunction {
125 null_first: bool,
126 signature: logical_expr::Signature,
127}
128
129impl fmt::Debug for MinByFunction {
130 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
131 f.debug_struct("MinBy")
132 .field("name", &self.name())
133 .field("signature", &self.signature)
134 .field("accumulator", &"<FUNC>")
135 .finish()
136 }
137}
138
139impl Default for MinByFunction {
140 fn default() -> Self {
141 Self::new(true)
142 }
143}
144
145impl MinByFunction {
146 pub fn new(null_first: bool) -> Self {
147 Self {
148 null_first,
149 signature: logical_expr::Signature::user_defined(logical_expr::Volatility::Immutable),
150 }
151 }
152}
153
154impl logical_expr::AggregateUDFImpl for MinByFunction {
155 fn name(&self) -> &str {
156 "min_by"
157 }
158
159 fn signature(&self) -> &logical_expr::Signature {
160 &self.signature
161 }
162
163 fn return_type(
164 &self,
165 arg_types: &[arrow::datatypes::DataType],
166 ) -> error::Result<arrow::datatypes::DataType> {
167 Ok(arg_types[0].to_owned())
168 }
169
170 fn accumulator(
171 &self,
172 _acc_args: logical_expr::function::AccumulatorArgs,
173 ) -> error::Result<Box<dyn logical_expr::Accumulator>> {
174 common::exec_err!("should not reach here")
175 }
176
177 fn coerce_types(
178 &self,
179 arg_types: &[arrow::datatypes::DataType],
180 ) -> error::Result<Vec<arrow::datatypes::DataType>> {
181 get_min_max_by_result_type(arg_types)
182 }
183
184 fn simplify(&self) -> Option<logical_expr::function::AggregateFunctionSimplification> {
185 let null_first = self.null_first;
186 let simplify = move |mut aggr_func: logical_expr::expr::AggregateFunction,
187 _: &logical_expr::simplify::SimplifyContext| {
188 let mut order_by = aggr_func.params.order_by;
189 let (second_arg, first_arg) = (
190 aggr_func.params.args.remove(1),
191 aggr_func.params.args.remove(0),
192 );
193
194 let sort = logical_expr::expr::Sort::new(second_arg, false, null_first);
195 order_by.push(sort); let func = logical_expr::expr::AggregateFunction::new_udf(
197 functions_aggregate::first_last::last_value_udaf(),
198 vec![first_arg],
199 aggr_func.params.distinct,
200 aggr_func.params.filter,
201 order_by,
202 aggr_func.params.null_treatment,
203 );
204 let func = logical_expr::expr::Expr::AggregateFunction(func);
205 Ok(func)
206 };
207 Some(Box::new(simplify))
208 }
209}
210
211#[cfg(test)]
212mod tests {
213 use super::*;
214
215 use datafusion::arrow::array::ArrayAccessor;
216 use datafusion::{arrow, datasource, error, prelude};
217 use std::sync;
218
219 const TEST_TABLE_NAME: &str = "types";
220 const STRING_COLUMN_NAME: &str = "string";
221 const DICTIONARY_COLUMN_NAME: &str = "dict_string";
222 const INT64_COLUMN_NAME: &str = "int64";
223 const FLOAT64_COLUMN_NAME: &str = "float64";
224
225 const MIN_STRING_VALUE: &str = "a";
226 const MID_STRING_VALUE: &str = "b";
227 const MAX_STRING_VALUE: &str = "c";
228 const MIN_FLOAT_VALUE: f64 = 0.25;
229 const MID_FLOAT_VALUE: f64 = 0.5;
230 const MAX_FLOAT_VALUE: f64 = 0.75;
231 const MIN_INT_VALUE: i64 = -1;
232 const MID_INT_VALUE: i64 = 0;
233 const MAX_INT_VALUE: i64 = 1;
234 const MIN_DICTIONARY_VALUE: &str = "a";
235 const MID_DICTIONARY_VALUE: &str = "b";
236 const MAX_DICTIONARY_VALUE: &str = "c";
237
238 fn test_schema() -> sync::Arc<arrow::datatypes::Schema> {
239 sync::Arc::new(arrow::datatypes::Schema::new(vec![
240 arrow::datatypes::Field::new(
241 STRING_COLUMN_NAME,
242 arrow::datatypes::DataType::Utf8,
243 false,
244 ),
245 arrow::datatypes::Field::new_dictionary(
246 DICTIONARY_COLUMN_NAME,
247 arrow::datatypes::DataType::Int32,
248 arrow::datatypes::DataType::Utf8,
249 false,
250 ),
251 arrow::datatypes::Field::new(
252 INT64_COLUMN_NAME,
253 arrow::datatypes::DataType::Int64,
254 false,
255 ),
256 arrow::datatypes::Field::new(
257 FLOAT64_COLUMN_NAME,
258 arrow::datatypes::DataType::Float64,
259 false,
260 ),
261 ]))
262 }
263
264 fn test_data(
265 schema: sync::Arc<arrow::datatypes::Schema>,
266 ) -> Vec<arrow::record_batch::RecordBatch> {
267 vec![
268 arrow::record_batch::RecordBatch::try_new(
269 schema,
270 vec![
271 sync::Arc::new(arrow::array::StringArray::from(vec![
272 MID_STRING_VALUE,
273 MIN_STRING_VALUE,
274 MAX_STRING_VALUE,
275 ])),
276 sync::Arc::new(
277 vec![
278 Some(MID_DICTIONARY_VALUE),
279 Some(MIN_DICTIONARY_VALUE),
280 Some(MAX_DICTIONARY_VALUE),
281 ]
282 .into_iter()
283 .collect::<arrow::array::DictionaryArray<arrow::datatypes::Int32Type>>(),
284 ),
285 sync::Arc::new(arrow::array::Int64Array::from(vec![
286 MID_INT_VALUE,
287 MIN_INT_VALUE,
288 MAX_INT_VALUE,
289 ])),
290 sync::Arc::new(arrow::array::Float64Array::from(vec![
291 MID_FLOAT_VALUE,
292 MIN_FLOAT_VALUE,
293 MAX_FLOAT_VALUE,
294 ])),
295 ],
296 )
297 .unwrap(),
298 ]
299 }
300
301 fn test_ctx() -> datafusion::common::Result<prelude::SessionContext> {
302 let schema = test_schema();
303 let data = test_data(schema.clone());
304 let table = datasource::MemTable::try_new(schema, vec![data])?;
305 let ctx = prelude::SessionContext::new();
306 ctx.register_table(TEST_TABLE_NAME, sync::Arc::new(table))?;
307 Ok(ctx)
308 }
309
310 async fn extract_single_value<T, A>(df: prelude::DataFrame) -> error::Result<T>
311 where
312 A: arrow::array::Array + 'static,
313 for<'a> &'a A: arrow::array::ArrayAccessor,
314 for<'a> <&'a A as arrow::array::ArrayAccessor>::Item: Into<T>,
315 {
316 let results = df.collect().await?;
317 let col = results[0].column(0);
318 let v1 = col.as_any().downcast_ref::<A>().unwrap();
319 let value = v1.value(0).into();
320 Ok(value)
321 }
322
323 #[cfg(test)]
324 mod max_by {
325
326 use super::*;
327
328 #[tokio::test]
329 async fn test_max_by_string_int() -> error::Result<()> {
330 let query = format!(
331 "SELECT max_by({}, {}) FROM {}",
332 STRING_COLUMN_NAME, INT64_COLUMN_NAME, TEST_TABLE_NAME
333 );
334 let df = ctx()?.sql(&query).await?;
335 let result = extract_single_value::<String, arrow::array::StringArray>(df).await?;
336 assert_eq!(result, MAX_STRING_VALUE);
337 Ok(())
338 }
339
340 #[tokio::test]
341 async fn test_max_by_string_float() -> error::Result<()> {
342 let query = format!(
343 "SELECT max_by({}, {}) FROM {}",
344 STRING_COLUMN_NAME, FLOAT64_COLUMN_NAME, TEST_TABLE_NAME
345 );
346 let df = ctx()?.sql(&query).await?;
347 let result = extract_single_value::<String, arrow::array::StringArray>(df).await?;
348 assert_eq!(result, MAX_STRING_VALUE);
349 Ok(())
350 }
351
352 #[tokio::test]
353 async fn test_max_by_float_string() -> error::Result<()> {
354 let query = format!(
355 "SELECT max_by({}, {}) FROM {}",
356 FLOAT64_COLUMN_NAME, STRING_COLUMN_NAME, TEST_TABLE_NAME
357 );
358 let df = ctx()?.sql(&query).await?;
359 let result = extract_single_value::<f64, arrow::array::Float64Array>(df).await?;
360 assert_eq!(result, MAX_FLOAT_VALUE);
361 Ok(())
362 }
363
364 #[tokio::test]
365 async fn test_max_by_int_string() -> error::Result<()> {
366 let query = format!(
367 "SELECT max_by({}, {}) FROM {}",
368 INT64_COLUMN_NAME, STRING_COLUMN_NAME, TEST_TABLE_NAME
369 );
370 let df = ctx()?.sql(&query).await?;
371 let result = extract_single_value::<i64, arrow::array::Int64Array>(df).await?;
372 assert_eq!(result, MAX_INT_VALUE);
373 Ok(())
374 }
375
376 #[tokio::test]
377 async fn test_max_by_dictionary_int() -> error::Result<()> {
378 let query = format!(
379 "SELECT max_by({}, {}) FROM {}",
380 DICTIONARY_COLUMN_NAME, INT64_COLUMN_NAME, TEST_TABLE_NAME
381 );
382 let df = ctx()?.sql(&query).await?;
383 let result = extract_single_value::<String, arrow::array::StringArray>(df).await?;
384 assert_eq!(result, MAX_DICTIONARY_VALUE);
385 Ok(())
386 }
387
388 #[tokio::test]
389 async fn test_max_by_ignores_nulls() -> error::Result<()> {
390 let query = r#"
391 SELECT max_by(v, k)
392 FROM (
393 VALUES
394 ('a', 1),
395 ('b', CAST(NULL AS INT)),
396 ('c', 2)
397 ) AS t(v, k)
398 "#;
399 let df = ctx()?.sql(query).await?;
400 let result = extract_single_value::<String, arrow::array::StringArray>(df).await?;
401 assert_eq!(result, "c", "max_by should ignore NULLs");
402 Ok(())
403 }
404
405 fn ctx() -> error::Result<prelude::SessionContext> {
406 let ctx = test_ctx()?;
407 let max_by_udaf = MaxByFunction::default();
408 ctx.register_udaf(max_by_udaf.into());
409 Ok(ctx)
410 }
411 }
412
413 #[cfg(test)]
414 mod min_by {
415
416 use super::*;
417
418 #[tokio::test]
419 async fn test_min_by_string_int() -> error::Result<()> {
420 let query = format!(
421 "SELECT min_by({}, {}) FROM {}",
422 STRING_COLUMN_NAME, INT64_COLUMN_NAME, TEST_TABLE_NAME
423 );
424 let df = ctx()?.sql(&query).await?;
425 let result = extract_single_value::<String, arrow::array::StringArray>(df).await?;
426 assert_eq!(result, MIN_STRING_VALUE);
427 Ok(())
428 }
429
430 #[tokio::test]
431 async fn test_min_by_string_float() -> error::Result<()> {
432 let query = format!(
433 "SELECT min_by({}, {}) FROM {}",
434 STRING_COLUMN_NAME, FLOAT64_COLUMN_NAME, TEST_TABLE_NAME
435 );
436 let df = ctx()?.sql(&query).await?;
437 let result = extract_single_value::<String, arrow::array::StringArray>(df).await?;
438 assert_eq!(result, MIN_STRING_VALUE);
439 Ok(())
440 }
441
442 #[tokio::test]
443 async fn test_min_by_float_string() -> error::Result<()> {
444 let query = format!(
445 "SELECT min_by({}, {}) FROM {}",
446 FLOAT64_COLUMN_NAME, STRING_COLUMN_NAME, TEST_TABLE_NAME
447 );
448 let df = ctx()?.sql(&query).await?;
449 let result = extract_single_value::<f64, arrow::array::Float64Array>(df).await?;
450 assert_eq!(result, MIN_FLOAT_VALUE);
451 Ok(())
452 }
453
454 #[tokio::test]
455 async fn test_min_by_int_string() -> error::Result<()> {
456 let query = format!(
457 "SELECT min_by({}, {}) FROM {}",
458 INT64_COLUMN_NAME, STRING_COLUMN_NAME, TEST_TABLE_NAME
459 );
460 let df = ctx()?.sql(&query).await?;
461 let result = extract_single_value::<i64, arrow::array::Int64Array>(df).await?;
462 assert_eq!(result, MIN_INT_VALUE);
463 Ok(())
464 }
465
466 #[tokio::test]
467 async fn test_min_by_dictionary_int() -> error::Result<()> {
468 let query = format!(
469 "SELECT min_by({}, {}) FROM {}",
470 DICTIONARY_COLUMN_NAME, INT64_COLUMN_NAME, TEST_TABLE_NAME
471 );
472 let df = ctx()?.sql(&query).await?;
473 let result = extract_single_value::<String, arrow::array::StringArray>(df).await?;
474 assert_eq!(result, MIN_DICTIONARY_VALUE);
475 Ok(())
476 }
477
478 #[tokio::test]
479 async fn test_min_by_ignores_nulls() -> error::Result<()> {
480 let query = r#"
481 SELECT min_by(v, k)
482 FROM (
483 VALUES
484 ('a', 1),
485 ('b', CAST(NULL AS INT)),
486 ('c', 2)
487 ) AS t(v, k)
488 "#;
489 let df = ctx()?.sql(query).await?;
490 let result = extract_single_value::<String, arrow::array::StringArray>(df).await?;
491 assert_eq!(result, "a", "min_by should ignore NULLs");
492 Ok(())
493 }
494
495 fn ctx() -> error::Result<prelude::SessionContext> {
496 let ctx = test_ctx()?;
497 let min_by_udaf = MinByFunction::default();
498 ctx.register_udaf(min_by_udaf.into());
499 Ok(ctx)
500 }
501 }
502}