Skip to content

Implement equals for stateful functions #16781

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Jul 18, 2025
Merged

Implement equals for stateful functions #16781

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
4c877f6
Implement equals for stateful functions
findepi Jul 15, 2025
a0db32d
downgrade for MSRV
findepi Jul 15, 2025
bc55bed
Improve doc
findepi Jul 16, 2025
492c446
Update default UDF:: equals to compare aliases too
findepi Jul 16, 2025
6016142
Update default UDF:: equals to compare type too (‼️)
findepi Jul 16, 2025
cb18579
remove now-obsoleted UDF equals/hash customizations
findepi Jul 16, 2025
f27480a
remove equals impl which compares name, signature -- default does that
findepi Jul 16, 2025
4c954ec
cleanup imports
findepi Jul 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions datafusion-examples/examples/function_factory.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -153,6 +154,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 {
Expand Down
27 changes: 27 additions & 0 deletions datafusion/core/tests/user_defined/user_defined_aggregates.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -957,6 +957,33 @@ impl AggregateUDFImpl for MetadataBasedAggregateUdf {
curr_sum: 0,
}))
}

fn equals(&self, other: &dyn AggregateUDFImpl) -> bool {
let Some(other) = other.as_any().downcast_ref::<Self>() else {
return false;
};
let Self {
name,
signature,
metadata,
} = self;
name == &other.name
&& signature == &other.signature
&& metadata == &other.metadata
}

fn hash_value(&self) -> u64 {
let Self {
name,
signature,
metadata: _, // unhashable
} = self;
let mut hasher = DefaultHasher::new();
std::any::type_name::<Self>().hash(&mut hasher);
name.hash(&mut hasher);
signature.hash(&mut hasher);
hasher.finish()
}
}

#[derive(Debug)]
Expand Down
117 changes: 112 additions & 5 deletions datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,34 @@ impl ScalarUDFImpl for Simple0ArgsScalarUDF {
fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> 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]
Expand Down Expand Up @@ -557,6 +585,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]
Expand Down Expand Up @@ -974,6 +1030,38 @@ impl ScalarUDFImpl for ScalarFunctionWrapper {

Ok(ExprSimplifyResult::Simplified(replacement))
}

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 {
Expand Down Expand Up @@ -1450,7 +1538,30 @@ impl ScalarUDFImpl for MetadataBasedUdf {
}

fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
self.name == other.name()
let Some(other) = other.as_any().downcast_ref::<Self>() else {
return false;
};
let Self {
name,
signature,
metadata,
} = self;
name == &other.name
&& signature == &other.signature
&& metadata == &other.metadata
}

fn hash_value(&self) -> u64 {
let Self {
name,
signature,
metadata: _, // unhashable
} = self;
let mut hasher = DefaultHasher::new();
std::any::type_name::<Self>().hash(&mut hasher);
name.hash(&mut hasher);
signature.hash(&mut hasher);
hasher.finish()
}
}

Expand Down Expand Up @@ -1669,10 +1780,6 @@ impl ScalarUDFImpl for ExtensionBasedUdf {
}
}
}

fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
self.name == other.name()
}
}

struct MyUserExtentionType {}
Expand Down
55 changes: 55 additions & 0 deletions datafusion/core/tests/user_defined/user_defined_window_functions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ use datafusion_physical_expr::{
PhysicalExpr,
};
use std::collections::HashMap;
use std::hash::{DefaultHasher, Hash, Hasher};
use std::{
any::Any,
ops::Range,
Expand Down Expand Up @@ -568,6 +569,33 @@ impl OddCounter {
fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
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()
}
}

ctx.register_udwf(WindowUDF::from(SimpleWindowUDF::new(test_state)))
Expand Down Expand Up @@ -815,6 +843,33 @@ impl WindowUDFImpl for MetadataBasedWindowUdf {
.with_metadata(self.metadata.clone())
.into())
}

fn equals(&self, other: &dyn WindowUDFImpl) -> bool {
let Some(other) = other.as_any().downcast_ref::<Self>() else {
return false;
};
let Self {
name,
signature,
metadata,
} = self;
name == &other.name
&& signature == &other.signature
&& metadata == &other.metadata
}

fn hash_value(&self) -> u64 {
let Self {
name,
signature,
metadata: _, // unhashable
} = self;
let mut hasher = DefaultHasher::new();
std::any::type_name::<Self>().hash(&mut hasher);
name.hash(&mut hasher);
signature.hash(&mut hasher);
hasher.finish()
}
}

#[derive(Debug)]
Expand Down
4 changes: 2 additions & 2 deletions datafusion/doc/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,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,
Expand Down Expand Up @@ -158,7 +158,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
Expand Down
17 changes: 17 additions & 0 deletions datafusion/expr/src/async_udf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ use datafusion_expr_common::columnar_value::ColumnarValue;
use datafusion_expr_common::signature::Signature;
use std::any::Any;
use std::fmt::{Debug, Display};
use std::hash::{DefaultHasher, Hash, Hasher};
use std::sync::Arc;

/// A scalar UDF that can invoke using async methods
Expand Down Expand Up @@ -111,6 +112,22 @@ impl ScalarUDFImpl for AsyncScalarUDF {
fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
internal_err!("async functions should not be called directly")
}

fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
let Some(other) = other.as_any().downcast_ref::<Self>() else {
return false;
};
let Self { inner } = self;
// TODO when MSRV >= 1.86.0, switch to `inner.equals(other.inner.as_ref())` leveraging trait upcasting
Arc::ptr_eq(inner, &other.inner)
}

fn hash_value(&self) -> u64 {
let Self { inner } = self;
let mut hasher = DefaultHasher::new();
Arc::as_ptr(inner).hash(&mut hasher);
hasher.finish()
}
}

impl Display for AsyncScalarUDF {
Expand Down
Loading