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