Derive UDWF equality from PartialEq, Hash (#17057)

* Simplify WindowUDFImpl::equals ForeignWindowUDF's impl

`ForeignWindowUDF` contains `FFI_WindowUDF` struct which was compared by
pointer only. This means that effectively `ForeignWindowUDF` was also
compared by pointer only, with all other equality checks being
redundant. This commit simplifies the implementation to make it more
obvious and more performant.

* Add PtrEq wrapper for pointer-based equality

* Derive UDWF equality from PartialEq, Hash

* Document moved functions

Author: findepi

datafusion/core/tests/user_defined/user_defined_window_functions.rs

@@ -28,8 +28,10 @@ use datafusion::common::test_util::batches_to_string;
 use datafusion::common::{Result, ScalarValue};
 use datafusion::prelude::SessionContext;
 use datafusion_common::exec_datafusion_err;
+use datafusion_expr::ptr_eq::PtrEq;
 use datafusion_expr::{
-    PartitionEvaluator, Signature, TypeSignature, Volatility, WindowUDF, WindowUDFImpl,
+    udf_equals_hash, PartitionEvaluator, Signature, TypeSignature, Volatility, WindowUDF,
+    WindowUDFImpl,
 };
 use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
 use datafusion_functions_window_common::{
@@ -523,10 +525,10 @@ impl OddCounter {
     }
 
     fn register(ctx: &mut SessionContext, test_state: Arc<TestState>) {
-        #[derive(Debug, Clone)]
+        #[derive(Debug, Clone, PartialEq, Hash)]
         struct SimpleWindowUDF {
             signature: Signature,
-            test_state: Arc<TestState>,
+            test_state: PtrEq<Arc<TestState>>,
             aliases: Vec<String>,
         }
 
@@ -536,7 +538,7 @@ impl OddCounter {
                     Signature::exact(vec![DataType::Float64], Volatility::Immutable);
                 Self {
                     signature,
-                    test_state,
+                    test_state: test_state.into(),
                     aliases: vec!["odd_counter_alias".to_string()],
                 }
             }
@@ -570,32 +572,7 @@ impl OddCounter {
                 Ok(Field::new(field_args.name(), DataType::Int64, true).into())
             }
 
-            fn equals(&self, other: &dyn WindowUDFImpl) -> bool {
-                let Some(other) = other.as_any().downcast_ref::<Self>() else {
-                    return false;
-                };
-                let Self {
-                    signature,
-                    test_state,
-                    aliases,
-                } = self;
-                signature == &other.signature
-                    && Arc::ptr_eq(test_state, &other.test_state)
-                    && aliases == &other.aliases
-            }
-
-            fn hash_value(&self) -> u64 {
-                let Self {
-                    signature,
-                    test_state,
-                    aliases,
-                } = self;
-                let mut hasher = DefaultHasher::new();
-                signature.hash(&mut hasher);
-                Arc::as_ptr(test_state).hash(&mut hasher);
-                aliases.hash(&mut hasher);
-                hasher.finish()
-            }
+            udf_equals_hash!(WindowUDFImpl);
         }
 
         ctx.register_udwf(WindowUDF::from(SimpleWindowUDF::new(test_state)))

datafusion/expr/src/async_udf.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::utils::{arc_ptr_eq, arc_ptr_hash};
+use crate::ptr_eq::{arc_ptr_eq, arc_ptr_hash};
 use crate::{
     udf_equals_hash, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl,
 };

datafusion/expr/src/expr_fn.rs

@@ -25,8 +25,8 @@ use crate::function::{
     AccumulatorArgs, AccumulatorFactoryFunction, PartitionEvaluatorFactory,
     StateFieldsArgs,
 };
+use crate::ptr_eq::PtrEq;
 use crate::select_expr::SelectExpr;
-use crate::utils::{arc_ptr_eq, arc_ptr_hash};
 use crate::{
     conditional_expressions::CaseBuilder, expr::Sort, logical_plan::Subquery,
     udf_equals_hash, AggregateUDF, Expr, LogicalPlan, Operator, PartitionEvaluator,
@@ -403,41 +403,12 @@ pub fn create_udf(
 
 /// Implements [`ScalarUDFImpl`] for functions that have a single signature and
 /// return type.
+#[derive(PartialEq, Hash)]
 pub struct SimpleScalarUDF {
     name: String,
     signature: Signature,
     return_type: DataType,
-    fun: ScalarFunctionImplementation,
-}
-
-impl PartialEq for SimpleScalarUDF {
-    fn eq(&self, other: &Self) -> bool {
-        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)
-    }
-}
-
-impl Hash for SimpleScalarUDF {
-    fn hash<H: Hasher>(&self, state: &mut H) {
-        let Self {
-            name,
-            signature,
-            return_type,
-            fun,
-        } = self;
-        name.hash(state);
-        signature.hash(state);
-        return_type.hash(state);
-        arc_ptr_hash(fun, state);
-    }
+    fun: PtrEq<ScalarFunctionImplementation>,
 }
 
 impl Debug for SimpleScalarUDF {
@@ -481,7 +452,7 @@ impl SimpleScalarUDF {
             name: name.into(),
             signature,
             return_type,
-            fun,
+            fun: fun.into(),
         }
     }
 }
@@ -690,11 +661,12 @@ pub fn create_udwf(
 
 /// Implements [`WindowUDFImpl`] for functions that have a single signature and
 /// return type.
+#[derive(PartialEq, Hash)]
 pub struct SimpleWindowUDF {
     name: String,
     signature: Signature,
     return_type: DataType,
-    partition_evaluator_factory: PartitionEvaluatorFactory,
+    partition_evaluator_factory: PtrEq<PartitionEvaluatorFactory>,
 }
 
 impl Debug for SimpleWindowUDF {
@@ -724,7 +696,7 @@ impl SimpleWindowUDF {
             name,
             signature,
             return_type,
-            partition_evaluator_factory,
+            partition_evaluator_factory: partition_evaluator_factory.into(),
         }
     }
 }
@@ -757,40 +729,7 @@ impl WindowUDFImpl for SimpleWindowUDF {
         )))
     }
 
-    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()
-    }
+    udf_equals_hash!(WindowUDFImpl);
 }
 
 pub fn interval_year_month_lit(value: &str) -> Expr {

datafusion/expr/src/lib.rs

@@ -67,6 +67,7 @@ pub mod async_udf;
 pub mod statistics {
     pub use datafusion_expr_common::statistics::*;
 }
+pub mod ptr_eq;
 pub mod test;
 pub mod tree_node;
 pub mod type_coercion;

datafusion/expr/src/ptr_eq.rs

@@ -0,0 +1,132 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::fmt::Debug;
+use std::hash::{Hash, Hasher};
+use std::ops::Deref;
+use std::sync::Arc;
+
+/// Compares two `Arc` pointers for equality based on their underlying pointers values.
+/// This is not equivalent to [`Arc::ptr_eq`] for fat pointers, see that method
+/// for more information.
+pub fn arc_ptr_eq<T: ?Sized>(a: &Arc<T>, b: &Arc<T>) -> bool {
+    std::ptr::eq(Arc::as_ptr(a), Arc::as_ptr(b))
+}
+
+/// Hashes an `Arc` pointer based on its underlying pointer value.
+/// The general contract for this function is that if [`arc_ptr_eq`] returns `true`
+/// for two `Arc`s, then this function should return the same hash value for both.
+pub fn arc_ptr_hash<T: ?Sized>(a: &Arc<T>, hasher: &mut impl Hasher) {
+    std::ptr::hash(Arc::as_ptr(a), hasher)
+}
+
+/// A wrapper around a pointer that implements `PartialEq` and `Hash` comparing
+/// the underlying pointer address.
+#[derive(Clone)]
+#[allow(private_bounds)] // This is so that PtrEq can only be used with allowed pointer types (e.g. Arc), without allowing misuse.
+pub struct PtrEq<Ptr: PointerType>(Ptr);
+
+impl<T> PartialEq for PtrEq<Arc<T>>
+where
+    T: ?Sized,
+{
+    fn eq(&self, other: &Self) -> bool {
+        arc_ptr_eq(&self.0, &other.0)
+    }
+}
+
+impl<T> Hash for PtrEq<Arc<T>>
+where
+    T: ?Sized,
+{
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        arc_ptr_hash(&self.0, state);
+    }
+}
+
+impl<Ptr> From<Ptr> for PtrEq<Ptr>
+where
+    Ptr: PointerType,
+{
+    fn from(ptr: Ptr) -> Self {
+        PtrEq(ptr)
+    }
+}
+
+impl<T> From<PtrEq<Arc<T>>> for Arc<T>
+where
+    T: ?Sized,
+{
+    fn from(wrapper: PtrEq<Arc<T>>) -> Self {
+        wrapper.0
+    }
+}
+
+impl<Ptr> Debug for PtrEq<Ptr>
+where
+    Ptr: PointerType + Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl<Ptr> Deref for PtrEq<Ptr>
+where
+    Ptr: PointerType,
+{
+    type Target = Ptr;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+trait PointerType {}
+impl<T> PointerType for Arc<T> where T: ?Sized {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::hash::DefaultHasher;
+
+    #[test]
+    pub fn test_ptr_eq_wrapper() {
+        let a = Arc::new("Hello".to_string());
+        let b = Arc::new(a.deref().clone());
+        let c = Arc::new("world".to_string());
+
+        let wrapper = PtrEq(Arc::clone(&a));
+        assert_eq!(wrapper, wrapper);
+
+        // same address (equal)
+        assert_eq!(PtrEq(Arc::clone(&a)), PtrEq(Arc::clone(&a)));
+        assert_eq!(hash(PtrEq(Arc::clone(&a))), hash(PtrEq(Arc::clone(&a))));
+
+        // different address, same content (not equal)
+        assert_ne!(PtrEq(Arc::clone(&a)), PtrEq(Arc::clone(&b)));
+
+        // different address, different content (not equal)
+        assert_ne!(PtrEq(Arc::clone(&a)), PtrEq(Arc::clone(&c)));
+    }
+
+    fn hash<T: Hash>(value: T) -> u64 {
+        let hasher = &mut DefaultHasher::new();
+        value.hash(hasher);
+        hasher.finish()
+    }
+}