Implement equals for stateful functions (apache#16781)

* Implement equals for stateful functions

Default implementation of `ScalarUDFImpl::equals`,
`AggregateUDFImpl::equals` and `WindowUDFImpl::equals` is correct for
stateless functions and those which only state is the `Signature`, which
is most of the functions.

This implements `equals` and `hash_value` for functions that have state
other than `Signature` object.

This fixes correctness issues which could occur when such stateful
functions are used together in one query.

* downgrade for MSRV

* Improve doc

* Update default UDF:: equals to compare aliases too

* Update default UDF:: equals to compare type too (‼️)

* remove now-obsoleted UDF equals/hash customizations

remove these which compare signature, aliases, as the default handles
these now

* remove equals impl which compares name, signature -- default does that

* cleanup imports

(cherry picked from commit afd8235)

Author: findepi

datafusion-examples/examples/function_factory.rs

@@ -28,6 +28,7 @@ use datafusion::logical_expr::{
     ColumnarValue, CreateFunction, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl,
     Signature, Volatility,
 };
+use std::hash::{DefaultHasher, Hash, Hasher};
 use std::result::Result as RResult;
 use std::sync::Arc;
 
@@ -157,6 +158,38 @@ impl ScalarUDFImpl for ScalarFunctionWrapper {
     fn output_ordering(&self, _input: &[ExprProperties]) -> Result<SortProperties> {
         Ok(SortProperties::Unordered)
     }
+
+    fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+        let Some(other) = other.as_any().downcast_ref::<Self>() else {
+            return false;
+        };
+        let Self {
+            name,
+            expr,
+            signature,
+            return_type,
+        } = self;
+        name == &other.name
+            && expr == &other.expr
+            && signature == &other.signature
+            && return_type == &other.return_type
+    }
+
+    fn hash_value(&self) -> u64 {
+        let Self {
+            name,
+            expr,
+            signature,
+            return_type,
+        } = self;
+        let mut hasher = DefaultHasher::new();
+        std::any::type_name::<Self>().hash(&mut hasher);
+        name.hash(&mut hasher);
+        expr.hash(&mut hasher);
+        signature.hash(&mut hasher);
+        return_type.hash(&mut hasher);
+        hasher.finish()
+    }
 }
 
 impl ScalarFunctionWrapper {

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

@@ -215,6 +215,34 @@ impl ScalarUDFImpl for Simple0ArgsScalarUDF {
     ) -> Result<ColumnarValue> {
         Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(100))))
     }
+
+    fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+        let Some(other) = other.as_any().downcast_ref::<Self>() else {
+            return false;
+        };
+        let Self {
+            name,
+            signature,
+            return_type,
+        } = self;
+        name == &other.name
+            && signature == &other.signature
+            && return_type == &other.return_type
+    }
+
+    fn hash_value(&self) -> u64 {
+        let Self {
+            name,
+            signature,
+            return_type,
+        } = self;
+        let mut hasher = DefaultHasher::new();
+        std::any::type_name::<Self>().hash(&mut hasher);
+        name.hash(&mut hasher);
+        signature.hash(&mut hasher);
+        return_type.hash(&mut hasher);
+        hasher.finish()
+    }
 }
 
 #[tokio::test]
@@ -556,6 +584,34 @@ impl ScalarUDFImpl for AddIndexToStringVolatileScalarUDF {
         };
         Ok(ColumnarValue::Array(Arc::new(StringArray::from(answer))))
     }
+
+    fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+        let Some(other) = other.as_any().downcast_ref::<Self>() else {
+            return false;
+        };
+        let Self {
+            name,
+            signature,
+            return_type,
+        } = self;
+        name == &other.name
+            && signature == &other.signature
+            && return_type == &other.return_type
+    }
+
+    fn hash_value(&self) -> u64 {
+        let Self {
+            name,
+            signature,
+            return_type,
+        } = self;
+        let mut hasher = DefaultHasher::new();
+        std::any::type_name::<Self>().hash(&mut hasher);
+        name.hash(&mut hasher);
+        signature.hash(&mut hasher);
+        return_type.hash(&mut hasher);
+        hasher.finish()
+    }
 }
 
 #[tokio::test]
@@ -985,6 +1041,38 @@ impl ScalarUDFImpl for ScalarFunctionWrapper {
     fn aliases(&self) -> &[String] {
         &[]
     }
+
+    fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+        let Some(other) = other.as_any().downcast_ref::<Self>() else {
+            return false;
+        };
+        let Self {
+            name,
+            expr,
+            signature,
+            return_type,
+        } = self;
+        name == &other.name
+            && expr == &other.expr
+            && signature == &other.signature
+            && return_type == &other.return_type
+    }
+
+    fn hash_value(&self) -> u64 {
+        let Self {
+            name,
+            expr,
+            signature,
+            return_type,
+        } = self;
+        let mut hasher = DefaultHasher::new();
+        std::any::type_name::<Self>().hash(&mut hasher);
+        name.hash(&mut hasher);
+        expr.hash(&mut hasher);
+        signature.hash(&mut hasher);
+        return_type.hash(&mut hasher);
+        hasher.finish()
+    }
 }
 
 impl ScalarFunctionWrapper {

datafusion/doc/src/lib.rs

@@ -33,7 +33,7 @@
 /// thus all text should be in English.
 ///
 /// [SQL function documentation]: https://datafusion.apache.org/user-guide/sql/index.html
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Hash)]
 pub struct Documentation {
     /// The section in the documentation where the UDF will be documented
     pub doc_section: DocSection,
@@ -153,7 +153,7 @@ impl Documentation {
     }
 }
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, Hash)]
 pub struct DocSection {
     /// True to include this doc section in the public
     /// documentation, false otherwise

datafusion/expr/src/expr_fn.rs

@@ -43,6 +43,7 @@ use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
 use sqlparser::ast::NullTreatment;
 use std::any::Any;
 use std::fmt::Debug;
+use std::hash::{DefaultHasher, Hash, Hasher};
 use std::ops::Not;
 use std::sync::Arc;
 
@@ -484,6 +485,38 @@ impl ScalarUDFImpl for SimpleScalarUDF {
     ) -> Result<ColumnarValue> {
         (self.fun)(args)
     }
+
+    fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+        let Some(other) = other.as_any().downcast_ref::<Self>() else {
+            return false;
+        };
+        let Self {
+            name,
+            signature,
+            return_type,
+            fun,
+        } = self;
+        name == &other.name
+            && signature == &other.signature
+            && return_type == &other.return_type
+            && Arc::ptr_eq(fun, &other.fun)
+    }
+
+    fn hash_value(&self) -> u64 {
+        let Self {
+            name,
+            signature,
+            return_type,
+            fun,
+        } = self;
+        let mut hasher = DefaultHasher::new();
+        std::any::type_name::<Self>().hash(&mut hasher);
+        name.hash(&mut hasher);
+        signature.hash(&mut hasher);
+        return_type.hash(&mut hasher);
+        Arc::as_ptr(fun).hash(&mut hasher);
+        hasher.finish()
+    }
 }
 
 /// Creates a new UDAF with a specific signature, state type and return type.
@@ -603,6 +636,42 @@ impl AggregateUDFImpl for SimpleAggregateUDF {
     fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
         Ok(self.state_fields.clone())
     }
+
+    fn equals(&self, other: &dyn AggregateUDFImpl) -> bool {
+        let Some(other) = other.as_any().downcast_ref::<Self>() else {
+            return false;
+        };
+        let Self {
+            name,
+            signature,
+            return_type,
+            accumulator,
+            state_fields,
+        } = self;
+        name == &other.name
+            && signature == &other.signature
+            && return_type == &other.return_type
+            && Arc::ptr_eq(accumulator, &other.accumulator)
+            && state_fields == &other.state_fields
+    }
+
+    fn hash_value(&self) -> u64 {
+        let Self {
+            name,
+            signature,
+            return_type,
+            accumulator,
+            state_fields,
+        } = self;
+        let mut hasher = DefaultHasher::new();
+        std::any::type_name::<Self>().hash(&mut hasher);
+        name.hash(&mut hasher);
+        signature.hash(&mut hasher);
+        return_type.hash(&mut hasher);
+        Arc::as_ptr(accumulator).hash(&mut hasher);
+        state_fields.hash(&mut hasher);
+        hasher.finish()
+    }
 }
 
 /// Creates a new UDWF with a specific signature, state type and return type.
@@ -695,6 +764,41 @@ impl WindowUDFImpl for SimpleWindowUDF {
             true,
         ))
     }
+
+    fn equals(&self, other: &dyn WindowUDFImpl) -> bool {
+        let Some(other) = other.as_any().downcast_ref::<Self>() else {
+            return false;
+        };
+        let Self {
+            name,
+            signature,
+            return_type,
+            partition_evaluator_factory,
+        } = self;
+        name == &other.name
+            && signature == &other.signature
+            && return_type == &other.return_type
+            && Arc::ptr_eq(
+                partition_evaluator_factory,
+                &other.partition_evaluator_factory,
+            )
+    }
+
+    fn hash_value(&self) -> u64 {
+        let Self {
+            name,
+            signature,
+            return_type,
+            partition_evaluator_factory,
+        } = self;
+        let mut hasher = DefaultHasher::new();
+        std::any::type_name::<Self>().hash(&mut hasher);
+        name.hash(&mut hasher);
+        signature.hash(&mut hasher);
+        return_type.hash(&mut hasher);
+        Arc::as_ptr(partition_evaluator_factory).hash(&mut hasher);
+        hasher.finish()
+    }
 }
 
 pub fn interval_year_month_lit(value: &str) -> Expr {

datafusion/expr/src/udaf.rs

@@ -574,26 +574,35 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {
     /// Return true if this aggregate UDF is equal to the other.
     ///
     /// Allows customizing the equality of aggregate UDFs.
+    /// *Must* be implemented explicitly if the UDF type has internal state.
     /// Must be consistent with [`Self::hash_value`] and follow the same rules as [`Eq`]:
     ///
     /// - reflexive: `a.equals(a)`;
     /// - symmetric: `a.equals(b)` implies `b.equals(a)`;
     /// - transitive: `a.equals(b)` and `b.equals(c)` implies `a.equals(c)`.
     ///
-    /// By default, compares [`Self::name`] and [`Self::signature`].
+    /// By default, compares type, [`Self::name`], [`Self::aliases`] and [`Self::signature`].
     fn equals(&self, other: &dyn AggregateUDFImpl) -> bool {
-        self.name() == other.name() && self.signature() == other.signature()
+        self.as_any().type_id() == other.as_any().type_id()
+            && self.name() == other.name()
+            && self.aliases() == other.aliases()
+            && self.signature() == other.signature()
     }
 
     /// Returns a hash value for this aggregate UDF.
     ///
-    /// Allows customizing the hash code of aggregate UDFs. Similarly to [`Hash`] and [`Eq`],
-    /// if [`Self::equals`] returns true for two UDFs, their `hash_value`s must be the same.
+    /// Allows customizing the hash code of aggregate UDFs.
+    /// *Must* be implemented explicitly whenever [`Self::equals`] is implemented.
     ///
-    /// By default, hashes [`Self::name`] and [`Self::signature`].
+    /// Similarly to [`Hash`] and [`Eq`], if [`Self::equals`] returns true for two UDFs,
+    /// their `hash_value`s must be the same.
+    ///
+    /// By default, it is consistent with default implementation of [`Self::equals`].
     fn hash_value(&self) -> u64 {
         let hasher = &mut DefaultHasher::new();
+        self.as_any().type_id().hash(hasher);
         self.name().hash(hasher);
+        self.aliases().hash(hasher);
         self.signature().hash(hasher);
         hasher.finish()
     }

datafusion/expr/src/udf.rs

@@ -801,26 +801,35 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
     /// Return true if this scalar UDF is equal to the other.
     ///
     /// Allows customizing the equality of scalar UDFs.
+    /// *Must* be implemented explicitly if the UDF type has internal state.
     /// Must be consistent with [`Self::hash_value`] and follow the same rules as [`Eq`]:
     ///
     /// - reflexive: `a.equals(a)`;
     /// - symmetric: `a.equals(b)` implies `b.equals(a)`;
     /// - transitive: `a.equals(b)` and `b.equals(c)` implies `a.equals(c)`.
     ///
-    /// By default, compares [`Self::name`] and [`Self::signature`].
+    /// By default, compares type, [`Self::name`], [`Self::aliases`] and [`Self::signature`].
     fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
-        self.name() == other.name() && self.signature() == other.signature()
+        self.as_any().type_id() == other.as_any().type_id()
+            && self.name() == other.name()
+            && self.aliases() == other.aliases()
+            && self.signature() == other.signature()
     }
 
     /// Returns a hash value for this scalar UDF.
     ///
-    /// Allows customizing the hash code of scalar UDFs. Similarly to [`Hash`] and [`Eq`],
-    /// if [`Self::equals`] returns true for two UDFs, their `hash_value`s must be the same.
+    /// Allows customizing the hash code of scalar UDFs.
+    /// *Must* be implemented explicitly whenever [`Self::equals`] is implemented.
     ///
-    /// By default, hashes [`Self::name`] and [`Self::signature`].
+    /// Similarly to [`Hash`] and [`Eq`], if [`Self::equals`] returns true for two UDFs,
+    /// their `hash_value`s must be the same.
+    ///
+    /// By default, it is consistent with default implementation of [`Self::equals`].
     fn hash_value(&self) -> u64 {
         let hasher = &mut DefaultHasher::new();
+        self.as_any().type_id().hash(hasher);
         self.name().hash(hasher);
+        self.aliases().hash(hasher);
         self.signature().hash(hasher);
         hasher.finish()
     }
@@ -951,6 +960,7 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl {
 
     fn hash_value(&self) -> u64 {
         let hasher = &mut DefaultHasher::new();
+        std::any::type_name::<Self>().hash(hasher);
         self.inner.hash_value().hash(hasher);
         self.aliases.hash(hasher);
         hasher.finish()