You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2022/09/12 12:32:13 UTC
[arrow-datafusion] branch master updated: Pass `return_type` to `AccumulatorFunctionImplementation ` for user defined aggregates (#3428)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/master by this push:
new 447141cb8 Pass `return_type` to `AccumulatorFunctionImplementation ` for user defined aggregates (#3428)
447141cb8 is described below
commit 447141cb8878331e5efcc39ff16e8e45dffa9ae4
Author: Andrew Lamb <an...@nerdnetworks.org>
AuthorDate: Mon Sep 12 08:32:07 2022 -0400
Pass `return_type` to `AccumulatorFunctionImplementation ` for user defined aggregates (#3428)
* Add return_type to accumulator factor signature
* fixup
* update for new API
---
datafusion-examples/examples/simple_udaf.rs | 2 +-
datafusion/core/src/execution/context.rs | 2 +-
datafusion/core/src/physical_plan/udaf.rs | 3 ++-
datafusion/core/tests/sql/udf.rs | 2 +-
datafusion/core/tests/user_defined_aggregates.rs | 2 +-
datafusion/expr/src/function.rs | 9 +++++----
datafusion/proto/src/lib.rs | 2 +-
7 files changed, 12 insertions(+), 10 deletions(-)
diff --git a/datafusion-examples/examples/simple_udaf.rs b/datafusion-examples/examples/simple_udaf.rs
index 378d2548e..0afb57271 100644
--- a/datafusion-examples/examples/simple_udaf.rs
+++ b/datafusion-examples/examples/simple_udaf.rs
@@ -169,7 +169,7 @@ async fn main() -> Result<()> {
Arc::new(DataType::Float64),
Volatility::Immutable,
// This is the accumulator factory; DataFusion uses it to create new accumulators.
- Arc::new(|| Ok(Box::new(GeometricMean::new()))),
+ Arc::new(|_| Ok(Box::new(GeometricMean::new()))),
// This is the description of the state. `state()` must match the types here.
Arc::new(vec![DataType::Float64, DataType::UInt32]),
);
diff --git a/datafusion/core/src/execution/context.rs b/datafusion/core/src/execution/context.rs
index aaed710df..64f75dfa2 100644
--- a/datafusion/core/src/execution/context.rs
+++ b/datafusion/core/src/execution/context.rs
@@ -2017,7 +2017,7 @@ mod tests {
DataType::Float64,
Arc::new(DataType::Float64),
Volatility::Immutable,
- Arc::new(|| Ok(Box::new(AvgAccumulator::try_new(&DataType::Float64)?))),
+ Arc::new(|_| Ok(Box::new(AvgAccumulator::try_new(&DataType::Float64)?))),
Arc::new(vec![DataType::UInt64, DataType::Float64]),
);
diff --git a/datafusion/core/src/physical_plan/udaf.rs b/datafusion/core/src/physical_plan/udaf.rs
index 71e7e0657..e017bb5ad 100644
--- a/datafusion/core/src/physical_plan/udaf.rs
+++ b/datafusion/core/src/physical_plan/udaf.rs
@@ -64,6 +64,7 @@ pub fn create_aggregate_expr(
pub struct AggregateFunctionExpr {
fun: AggregateUDF,
args: Vec<Arc<dyn PhysicalExpr>>,
+ /// Output / return type of this aggregate
data_type: DataType,
name: String,
}
@@ -99,7 +100,7 @@ impl AggregateExpr for AggregateFunctionExpr {
}
fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
- (self.fun.accumulator)()
+ (self.fun.accumulator)(&self.data_type)
}
fn name(&self) -> &str {
diff --git a/datafusion/core/tests/sql/udf.rs b/datafusion/core/tests/sql/udf.rs
index 705f0abe1..9d13d5f59 100644
--- a/datafusion/core/tests/sql/udf.rs
+++ b/datafusion/core/tests/sql/udf.rs
@@ -166,7 +166,7 @@ async fn simple_udaf() -> Result<()> {
DataType::Float64,
Arc::new(DataType::Float64),
Volatility::Immutable,
- Arc::new(|| Ok(Box::new(AvgAccumulator::try_new(&DataType::Float64)?))),
+ Arc::new(|_| Ok(Box::new(AvgAccumulator::try_new(&DataType::Float64)?))),
Arc::new(vec![DataType::UInt64, DataType::Float64]),
);
diff --git a/datafusion/core/tests/user_defined_aggregates.rs b/datafusion/core/tests/user_defined_aggregates.rs
index 33e5d0a0d..fe84a7e8d 100644
--- a/datafusion/core/tests/user_defined_aggregates.rs
+++ b/datafusion/core/tests/user_defined_aggregates.rs
@@ -114,7 +114,7 @@ fn register_aggregate(ctx: &mut SessionContext) {
let signatures = vec![TypeSignature::Exact(FirstSelector::input_datatypes())];
let accumulator: AccumulatorFunctionImplementation =
- Arc::new(|| Ok(Box::new(FirstSelector::new())));
+ Arc::new(|_| Ok(Box::new(FirstSelector::new())));
let volatility = Volatility::Immutable;
diff --git a/datafusion/expr/src/function.rs b/datafusion/expr/src/function.rs
index b8b17bf40..973ea8f83 100644
--- a/datafusion/expr/src/function.rs
+++ b/datafusion/expr/src/function.rs
@@ -39,15 +39,16 @@ use std::sync::Arc;
pub type ScalarFunctionImplementation =
Arc<dyn Fn(&[ColumnarValue]) -> Result<ColumnarValue> + Send + Sync>;
-/// A function's return type
+/// Factory that returns the functions's return type given the input argument types
pub type ReturnTypeFunction =
Arc<dyn Fn(&[DataType]) -> Result<Arc<DataType>> + Send + Sync>;
-/// the implementation of an aggregate function
+/// Factory that returns an accumulator for the given aggregate, given
+/// its return datatype.
pub type AccumulatorFunctionImplementation =
- Arc<dyn Fn() -> Result<Box<dyn Accumulator>> + Send + Sync>;
+ Arc<dyn Fn(&DataType) -> Result<Box<dyn Accumulator>> + Send + Sync>;
-/// This signature corresponds to which types an aggregator serializes
+/// Factory that returns the types used by an aggregator to serialize
/// its state, given its return datatype.
pub type StateTypeFunction =
Arc<dyn Fn(&DataType) -> Result<Arc<Vec<DataType>>> + Send + Sync>;
diff --git a/datafusion/proto/src/lib.rs b/datafusion/proto/src/lib.rs
index f0da13793..8e9475329 100644
--- a/datafusion/proto/src/lib.rs
+++ b/datafusion/proto/src/lib.rs
@@ -1089,7 +1089,7 @@ mod roundtrip_tests {
Arc::new(DataType::Float64),
Volatility::Immutable,
// This is the accumulator factory; DataFusion uses it to create new accumulators.
- Arc::new(|| Ok(Box::new(Dummy {}))),
+ Arc::new(|_| Ok(Box::new(Dummy {}))),
// This is the description of the state. `state()` must match the types here.
Arc::new(vec![DataType::Float64, DataType::UInt32]),
);