Support WITHIN GROUP syntax to standardize certain existing aggregate functions

datafusion-examples/examples/simplify_udaf_expression.rs

@@ -100,6 +100,7 @@ impl AggregateUDFImpl for BetterAvgUdaf {
                 aggregate_function.filter,
                 aggregate_function.order_by,
                 aggregate_function.null_treatment,
+                aggregate_function.within_group,
             )))
         };
 

datafusion-examples/examples/sql_frontend.rs

@@ -157,6 +157,10 @@ impl ContextProvider for MyContextProvider {
         None
     }
 
+    fn get_ordered_set_aggregate_meta(&self, _name: &str) -> Option<Arc<AggregateUDF>> {
+        None
+    }
+
     fn get_variable_type(&self, _variable_names: &[String]) -> Option<DataType> {
         None
     }

datafusion/catalog/src/session.rs

@@ -110,6 +110,11 @@ pub trait Session: Send + Sync {
     /// Return reference to aggregate_functions
     fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>>;
 
+    /// Return reference to ordered_set_aggregate_functions
+    ///
+    /// Note : ordered_set_aggregate_functions are a subset of aggregate_functions.
+    fn ordered_set_aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>>;
+
     /// Return reference to window functions
     fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>>;
 
@@ -132,6 +137,7 @@ impl From<&dyn Session> for TaskContext {
             state.config().clone(),
             state.scalar_functions().clone(),
             state.aggregate_functions().clone(),
+            state.ordered_set_aggregate_functions().clone(),
             state.window_functions().clone(),
             state.runtime_env().clone(),
         )

datafusion/core/src/execution/context/mod.rs

@@ -1580,6 +1580,10 @@ impl FunctionRegistry for SessionContext {
         self.state.read().udaf(name)
     }
 
+    fn ordered_set_udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {
+        self.state.read().ordered_set_udaf(name)
+    }
+
     fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>> {
         self.state.read().udwf(name)
     }
@@ -1595,6 +1599,15 @@ impl FunctionRegistry for SessionContext {
         self.state.write().register_udaf(udaf)
     }
 
+    fn register_ordered_set_udaf(
+        &mut self,
+        ordered_set_udaf: Arc<AggregateUDF>,
+    ) -> Result<Option<Arc<AggregateUDF>>> {
+        self.state
+            .write()
+            .register_ordered_set_udaf(ordered_set_udaf)
+    }
+
     fn register_udwf(&mut self, udwf: Arc<WindowUDF>) -> Result<Option<Arc<WindowUDF>>> {
         self.state.write().register_udwf(udwf)
     }

datafusion/core/src/execution/session_state.rs

@@ -144,6 +144,10 @@ pub struct SessionState {
     scalar_functions: HashMap<String, Arc<ScalarUDF>>,
     /// Aggregate functions registered in the context
     aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
+    /// Ordered-set aggregate functions registered in the context
+    ///
+    /// Note : ordered_set_aggregate_functions are a subset of aggregate_functions.
+    ordered_set_aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
     /// Window functions registered in the context
     window_functions: HashMap<String, Arc<WindowUDF>>,
     /// Deserializer registry for extensions.
@@ -202,6 +206,7 @@ impl Debug for SessionState {
             .field("table_functions", &self.table_functions)
             .field("scalar_functions", &self.scalar_functions)
             .field("aggregate_functions", &self.aggregate_functions)
+            .field("ordered_set_aggregate_functions", &self.ordered_set_aggregate_functions)
             .field("window_functions", &self.window_functions)
             .field("prepared_plans", &self.prepared_plans)
             .finish()
@@ -241,6 +246,10 @@ impl Session for SessionState {
         &self.aggregate_functions
     }
 
+    fn ordered_set_aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
+        &self.ordered_set_aggregate_functions
+    }
+
     fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>> {
         &self.window_functions
     }
@@ -880,6 +889,13 @@ impl SessionState {
         &self.aggregate_functions
     }
 
+    /// Return reference to ordered_set_aggregate_functions
+    ///
+    /// Note : ordered_set_aggregate_functions are a subset of aggregate_functions.
+    pub fn ordered_set_aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
+        &self.ordered_set_aggregate_functions
+    }
+
     /// Return reference to window functions
     pub fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>> {
         &self.window_functions
@@ -968,6 +984,7 @@ pub struct SessionStateBuilder {
     table_functions: Option<HashMap<String, Arc<TableFunction>>>,
     scalar_functions: Option<Vec<Arc<ScalarUDF>>>,
     aggregate_functions: Option<Vec<Arc<AggregateUDF>>>,
+    ordered_set_aggregate_functions: Option<Vec<Arc<AggregateUDF>>>,
     window_functions: Option<Vec<Arc<WindowUDF>>>,
     serializer_registry: Option<Arc<dyn SerializerRegistry>>,
     file_formats: Option<Vec<Arc<dyn FileFormatFactory>>>,
@@ -998,6 +1015,7 @@ impl SessionStateBuilder {
             table_functions: None,
             scalar_functions: None,
             aggregate_functions: None,
+            ordered_set_aggregate_functions: None,
             window_functions: None,
             serializer_registry: None,
             file_formats: None,
@@ -1048,6 +1066,9 @@ impl SessionStateBuilder {
             aggregate_functions: Some(
                 existing.aggregate_functions.into_values().collect_vec(),
             ),
+            ordered_set_aggregate_functions: Some(
+                existing.ordered_set_aggregate_functions.into_values().collect_vec(),
+            ),
             window_functions: Some(existing.window_functions.into_values().collect_vec()),
             serializer_registry: Some(existing.serializer_registry),
             file_formats: Some(existing.file_formats.into_values().collect_vec()),
@@ -1073,6 +1094,7 @@ impl SessionStateBuilder {
             .with_expr_planners(SessionStateDefaults::default_expr_planners())
             .with_scalar_functions(SessionStateDefaults::default_scalar_functions())
             .with_aggregate_functions(SessionStateDefaults::default_aggregate_functions())
+            .with_ordered_set_aggregate_functions(SessionStateDefaults::default_ordered_set_aggregate_functions())
             .with_window_functions(SessionStateDefaults::default_window_functions())
     }
 
@@ -1206,6 +1228,17 @@ impl SessionStateBuilder {
         self
     }
 
+    /// Set the map of ordered-set [`AggregateUDF`]s
+    ///
+    /// Note : ordered_set_aggregate_functions are a subset of aggregate_functions.
+    pub fn with_ordered_set_aggregate_functions(
+        mut self,
+        ordered_set_aggregate_functions: Vec<Arc<AggregateUDF>>,
+    ) -> Self {
+        self.ordered_set_aggregate_functions = Some(ordered_set_aggregate_functions);
+        self
+    }
+
     /// Set the map of [`WindowUDF`]s
     pub fn with_window_functions(
         mut self,
@@ -1340,6 +1373,7 @@ impl SessionStateBuilder {
             table_functions,
             scalar_functions,
             aggregate_functions,
+            ordered_set_aggregate_functions,
             window_functions,
             serializer_registry,
             file_formats,
@@ -1371,6 +1405,7 @@ impl SessionStateBuilder {
             table_functions: table_functions.unwrap_or_default(),
             scalar_functions: HashMap::new(),
             aggregate_functions: HashMap::new(),
+            ordered_set_aggregate_functions: HashMap::new(),
             window_functions: HashMap::new(),
             serializer_registry: serializer_registry
                 .unwrap_or(Arc::new(EmptySerializerRegistry)),
@@ -1411,6 +1446,20 @@ impl SessionStateBuilder {
             });
         }
 
+        if let Some(ordered_set_aggregate_functions) = ordered_set_aggregate_functions {
+            ordered_set_aggregate_functions
+                .into_iter()
+                .for_each(|udaf| {
+                    let existing_ordered_set_udf = state.register_ordered_set_udaf(udaf);
+                    if let Ok(Some(existing_ordered_set_udf)) = existing_ordered_set_udf {
+                        debug!(
+                            "Overwrote an ordered set existing UDF: {}",
+                            existing_ordered_set_udf.name()
+                        );
+                    }
+                });
+        }
+
         if let Some(window_functions) = window_functions {
             window_functions.into_iter().for_each(|udwf| {
                 let existing_udf = state.register_udwf(udwf);
@@ -1610,6 +1659,7 @@ impl Debug for SessionStateBuilder {
             .field("table_functions", &self.table_functions)
             .field("scalar_functions", &self.scalar_functions)
             .field("aggregate_functions", &self.aggregate_functions)
+            .field("ordered_set_aggregate_functions", &self.ordered_set_aggregate_functions)
             .field("window_functions", &self.window_functions)
             .finish()
     }
@@ -1696,6 +1746,13 @@ impl<'a> ContextProvider for SessionContextProvider<'a> {
         self.state.aggregate_functions().get(name).cloned()
     }
 
+    fn get_ordered_set_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
+        self.state
+            .ordered_set_aggregate_functions()
+            .get(name)
+            .cloned()
+    }
+
     fn get_window_meta(&self, name: &str) -> Option<Arc<WindowUDF>> {
         self.state.window_functions().get(name).cloned()
     }
@@ -1766,6 +1823,17 @@ impl FunctionRegistry for SessionState {
         })
     }
 
+    fn ordered_set_udaf(
+        &self,
+        name: &str,
+    ) -> datafusion_common::Result<Arc<AggregateUDF>> {
+        let result = self.ordered_set_aggregate_functions.get(name);
+
+        result.cloned().ok_or_else(|| {
+            plan_datafusion_err!("There is no ordered set UDAF named \"{name}\" in the registry")
+        })
+    }
+
     fn udwf(&self, name: &str) -> datafusion_common::Result<Arc<WindowUDF>> {
         let result = self.window_functions.get(name);
 
@@ -1796,6 +1864,18 @@ impl FunctionRegistry for SessionState {
         Ok(self.aggregate_functions.insert(udaf.name().into(), udaf))
     }
 
+    fn register_ordered_set_udaf(
+        &mut self,
+        ordered_set_udaf: Arc<AggregateUDF>,
+    ) -> datafusion_common::Result<Option<Arc<AggregateUDF>>> {
+        ordered_set_udaf.aliases().iter().for_each(|alias| {
+            self.ordered_set_aggregate_functions
+                .insert(alias.clone(), Arc::clone(&ordered_set_udaf));
+        });
+        Ok(self.ordered_set_aggregate_functions
+            .insert(ordered_set_udaf.name().into(), ordered_set_udaf))
+    }
+
     fn register_udwf(
         &mut self,
         udwf: Arc<WindowUDF>,
@@ -1833,6 +1913,19 @@ impl FunctionRegistry for SessionState {
         Ok(udaf)
     }
 
+    fn deregister_ordered_set_udaf(
+        &mut self,
+        name: &str,
+    ) -> datafusion_common::Result<Option<Arc<AggregateUDF>>> {
+        let ordered_set_udaf = self.ordered_set_aggregate_functions.remove(name);
+        if let Some(ordered_set_udaf) = &ordered_set_udaf {
+            for alias in ordered_set_udaf.aliases() {
+                self.ordered_set_aggregate_functions.remove(alias);
+            }
+        }
+        Ok(ordered_set_udaf)
+    }
+
     fn deregister_udwf(
         &mut self,
         name: &str,
@@ -1895,6 +1988,7 @@ impl From<&SessionState> for TaskContext {
             state.config.clone(),
             state.scalar_functions.clone(),
             state.aggregate_functions.clone(),
+            state.ordered_set_aggregate_functions.clone(),
             state.window_functions.clone(),
             Arc::clone(&state.runtime_env),
         )

datafusion/core/src/execution/session_state_defaults.rs

@@ -114,6 +114,13 @@ impl SessionStateDefaults {
         functions_aggregate::all_default_aggregate_functions()
     }
 
+    /// returns the list of default ordered-set [`AggregateUDF']'s
+    ///
+    /// Note : default_ordered_set_aggregate_functions are a subset of default_aggregate_functions.
+    pub fn default_ordered_set_aggregate_functions() -> Vec<Arc<AggregateUDF>> {
+        functions_aggregate::all_default_ordered_set_aggregate_functions()
+    }
+
     /// returns the list of default [`WindowUDF']'s
     pub fn default_window_functions() -> Vec<Arc<WindowUDF>> {
         functions_window::all_default_window_functions()

datafusion/core/src/physical_planner.rs

@@ -1594,6 +1594,7 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
             filter,
             order_by,
             null_treatment,
+            within_group,
         }) => {
             let name = if let Some(name) = name {
                 name
@@ -1616,13 +1617,20 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
                 == NullTreatment::IgnoreNulls;
 
             let (agg_expr, filter, order_by) = {
-                let physical_sort_exprs = match order_by {
-                    Some(exprs) => Some(create_physical_sort_exprs(
-                        exprs,
+                let physical_sort_exprs = match within_group {
+                    Some(within_group) => Some(create_physical_sort_exprs(
+                        within_group,
                         logical_input_schema,
                         execution_props,
                     )?),
-                    None => None,
+                    None => match order_by {
+                        Some(order_by) => Some(create_physical_sort_exprs(
+                            order_by,
+                            logical_input_schema,
+                            execution_props,
+                        )?),
+                        None => None,
+                    },
                 };
 
                 let ordering_reqs: LexOrdering =

datafusion/core/tests/dataframe/dataframe_functions.rs

@@ -360,14 +360,14 @@ async fn test_fn_approx_median() -> Result<()> {
 
 #[tokio::test]
 async fn test_fn_approx_percentile_cont() -> Result<()> {
-    let expr = approx_percentile_cont(col("b"), lit(0.5), None);
+    let expr = approx_percentile_cont(vec![col("b").sort(true, false)], lit(0.5), None);
 
     let expected = [
-        "+---------------------------------------------+",
-        "| approx_percentile_cont(test.b,Float64(0.5)) |",
-        "+---------------------------------------------+",
-        "| 10                                          |",
-        "+---------------------------------------------+",
+        "+-----------------------------------------------------------------------------+",
+        "| approx_percentile_cont(test.b,Float64(0.5)) WITHIN GROUP [b ASC NULLS LAST] |",
+        "+-----------------------------------------------------------------------------+",
+        "| 10                                                                          |",
+        "+-----------------------------------------------------------------------------+",
     ];
 
     let df = create_test_table().await?;
@@ -381,27 +381,27 @@ async fn test_fn_approx_percentile_cont() -> Result<()> {
         None::<&str>,
         "arg_2".to_string(),
     ));
-    let expr = approx_percentile_cont(col("b"), alias_expr, None);
+    let expr = approx_percentile_cont(vec![col("b").sort(true, false)], alias_expr, None);
     let df = create_test_table().await?;
     let expected = [
-        "+--------------------------------------+",
-        "| approx_percentile_cont(test.b,arg_2) |",
-        "+--------------------------------------+",
-        "| 10                                   |",
-        "+--------------------------------------+",
+        "+----------------------------------------------------------------------+",
+        "| approx_percentile_cont(test.b,arg_2) WITHIN GROUP [b ASC NULLS LAST] |",
+        "+----------------------------------------------------------------------+",
+        "| 10                                                                   |",
+        "+----------------------------------------------------------------------+",
     ];
     let batches = df.aggregate(vec![], vec![expr]).unwrap().collect().await?;
 
     assert_batches_eq!(expected, &batches);
 
     // with number of centroids set
-    let expr = approx_percentile_cont(col("b"), lit(0.5), Some(lit(2)));
+    let expr = approx_percentile_cont(vec![col("b").sort(true, false)], lit(0.5), Some(lit(2)));
     let expected = [
-        "+------------------------------------------------------+",
-        "| approx_percentile_cont(test.b,Float64(0.5),Int32(2)) |",
-        "+------------------------------------------------------+",
-        "| 30                                                   |",
-        "+------------------------------------------------------+",
+        "+--------------------------------------------------------------------------------------+",
+        "| approx_percentile_cont(test.b,Float64(0.5),Int32(2)) WITHIN GROUP [b ASC NULLS LAST] |",
+        "+--------------------------------------------------------------------------------------+",
+        "| 30                                                                                   |",
+        "+--------------------------------------------------------------------------------------+",
     ];
 
     let df = create_test_table().await?;