fix: UDF, UDAF, UDWF with_alias(..) should wrap the inner function fully (#12098)

* fix: UDF, UDAF, UDWF with_alias(..) should wrap the inner function fully

* revert back to having Arc<Self>

* add notes about adding stuff into Aliased impls

* fix clippy

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: Blizzara

datafusion/expr/src/udaf.rs

@@ -337,6 +337,9 @@ where
 /// let expr = geometric_mean.call(vec![col("a")]);
 /// ```
 pub trait AggregateUDFImpl: Debug + Send + Sync {
+    // Note: When adding any methods (with default implementations), remember to add them also
+    // into the AliasedAggregateUDFImpl below!
+
     /// Returns this object as an [`Any`] trait object
     fn as_any(&self) -> &dyn Any;
 
@@ -635,6 +638,60 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl {
         &self.aliases
     }
 
+    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
+        self.inner.state_fields(args)
+    }
+
+    fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
+        self.inner.groups_accumulator_supported(args)
+    }
+
+    fn create_groups_accumulator(
+        &self,
+        args: AccumulatorArgs,
+    ) -> Result<Box<dyn GroupsAccumulator>> {
+        self.inner.create_groups_accumulator(args)
+    }
+
+    fn create_sliding_accumulator(
+        &self,
+        args: AccumulatorArgs,
+    ) -> Result<Box<dyn Accumulator>> {
+        self.inner.accumulator(args)
+    }
+
+    fn with_beneficial_ordering(
+        self: Arc<Self>,
+        beneficial_ordering: bool,
+    ) -> Result<Option<Arc<dyn AggregateUDFImpl>>> {
+        Arc::clone(&self.inner)
+            .with_beneficial_ordering(beneficial_ordering)
+            .map(|udf| {
+                udf.map(|udf| {
+                    Arc::new(AliasedAggregateUDFImpl {
+                        inner: udf,
+                        aliases: self.aliases.clone(),
+                    }) as Arc<dyn AggregateUDFImpl>
+                })
+            })
+    }
+
+    fn order_sensitivity(&self) -> AggregateOrderSensitivity {
+        self.inner.order_sensitivity()
+    }
+
+    fn simplify(&self) -> Option<AggregateFunctionSimplification> {
+        self.inner.simplify()
+    }
+
+    fn reverse_expr(&self) -> ReversedUDAF {
+        self.inner.reverse_expr()
+    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        self.inner.coerce_types(arg_types)
+    }
+
     fn equals(&self, other: &dyn AggregateUDFImpl) -> bool {
         if let Some(other) = other.as_any().downcast_ref::<AliasedAggregateUDFImpl>() {
             self.inner.equals(other.inner.as_ref()) && self.aliases == other.aliases
@@ -649,6 +706,10 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl {
         self.aliases.hash(hasher);
         hasher.finish()
     }
+
+    fn is_descending(&self) -> Option<bool> {
+        self.inner.is_descending()
+    }
 }
 
 /// Implementation of [`AggregateUDFImpl`] that wraps the function style pointers

datafusion/expr/src/udf.rs

@@ -346,6 +346,9 @@ where
 /// let expr = add_one.call(vec![col("a")]);
 /// ```
 pub trait ScalarUDFImpl: Debug + Send + Sync {
+    // Note: When adding any methods (with default implementations), remember to add them also
+    // into the AliasedScalarUDFImpl below!
+
     /// Returns this object as an [`Any`] trait object
     fn as_any(&self) -> &dyn Any;
 
@@ -632,6 +635,14 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl {
         self.inner.name()
     }
 
+    fn display_name(&self, args: &[Expr]) -> Result<String> {
+        self.inner.display_name(args)
+    }
+
+    fn schema_name(&self, args: &[Expr]) -> Result<String> {
+        self.inner.schema_name(args)
+    }
+
     fn signature(&self) -> &Signature {
         self.inner.signature()
     }
@@ -640,12 +651,57 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl {
         self.inner.return_type(arg_types)
     }
 
+    fn aliases(&self) -> &[String] {
+        &self.aliases
+    }
+
+    fn return_type_from_exprs(
+        &self,
+        args: &[Expr],
+        schema: &dyn ExprSchema,
+        arg_types: &[DataType],
+    ) -> Result<DataType> {
+        self.inner.return_type_from_exprs(args, schema, arg_types)
+    }
+
     fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
         self.inner.invoke(args)
     }
 
-    fn aliases(&self) -> &[String] {
-        &self.aliases
+    fn invoke_no_args(&self, number_rows: usize) -> Result<ColumnarValue> {
+        self.inner.invoke_no_args(number_rows)
+    }
+
+    fn simplify(
+        &self,
+        args: Vec<Expr>,
+        info: &dyn SimplifyInfo,
+    ) -> Result<ExprSimplifyResult> {
+        self.inner.simplify(args, info)
+    }
+
+    fn short_circuits(&self) -> bool {
+        self.inner.short_circuits()
+    }
+
+    fn evaluate_bounds(&self, input: &[&Interval]) -> Result<Interval> {
+        self.inner.evaluate_bounds(input)
+    }
+
+    fn propagate_constraints(
+        &self,
+        interval: &Interval,
+        inputs: &[&Interval],
+    ) -> Result<Option<Vec<Interval>>> {
+        self.inner.propagate_constraints(interval, inputs)
+    }
+
+    fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> {
+        self.inner.output_ordering(inputs)
+    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        self.inner.coerce_types(arg_types)
     }
 
     fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {

datafusion/expr/src/udwf.rs

@@ -266,6 +266,9 @@ where
 ///     .unwrap();
 /// ```
 pub trait WindowUDFImpl: Debug + Send + Sync {
+    // Note: When adding any methods (with default implementations), remember to add them also
+    // into the AliasedWindowUDFImpl below!
+
     /// Returns this object as an [`Any`] trait object
     fn as_any(&self) -> &dyn Any;
 
@@ -428,6 +431,10 @@ impl WindowUDFImpl for AliasedWindowUDFImpl {
         &self.aliases
     }
 
+    fn simplify(&self) -> Option<WindowFunctionSimplification> {
+        self.inner.simplify()
+    }
+
     fn equals(&self, other: &dyn WindowUDFImpl) -> bool {
         if let Some(other) = other.as_any().downcast_ref::<AliasedWindowUDFImpl>() {
             self.inner.equals(other.inner.as_ref()) && self.aliases == other.aliases
@@ -442,6 +449,18 @@ impl WindowUDFImpl for AliasedWindowUDFImpl {
         self.aliases.hash(hasher);
         hasher.finish()
     }
+
+    fn nullable(&self) -> bool {
+        self.inner.nullable()
+    }
+
+    fn sort_options(&self) -> Option<SortOptions> {
+        self.inner.sort_options()
+    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        self.inner.coerce_types(arg_types)
+    }
 }
 
 /// Implementation of [`WindowUDFImpl`] that wraps the function style pointers