alamb commented on code in PR #3768:
URL: https://github.com/apache/arrow-datafusion/pull/3768#discussion_r992463588
##########
datafusion/physical-expr/src/aggregate/coercion_rule.rs:
##########
@@ -1,54 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
👍
##########
datafusion/optimizer/src/type_coercion.rs:
##########
@@ -596,6 +659,123 @@ mod test {
Ok(())
}
+ #[test]
+ fn agg_udaf() -> Result<()> {
+ let empty = empty();
+ let my_avg = create_udaf(
+ "MY_AVG",
+ DataType::Float64,
+ Arc::new(DataType::Float64),
+ Volatility::Immutable,
+ Arc::new(|_|
Ok(Box::new(AvgAccumulator::try_new(&DataType::Float64)?))),
+ Arc::new(vec![DataType::UInt64, DataType::Float64]),
+ );
+ let udaf = Expr::AggregateUDF {
+ fun: Arc::new(my_avg),
+ args: vec![lit(10i64)],
Review Comment:
👍
##########
datafusion/optimizer/src/type_coercion.rs:
##########
@@ -596,6 +659,123 @@ mod test {
Ok(())
}
+ #[test]
+ fn agg_udaf() -> Result<()> {
+ let empty = empty();
+ let my_avg = create_udaf(
+ "MY_AVG",
+ DataType::Float64,
+ Arc::new(DataType::Float64),
+ Volatility::Immutable,
+ Arc::new(|_|
Ok(Box::new(AvgAccumulator::try_new(&DataType::Float64)?))),
+ Arc::new(vec![DataType::UInt64, DataType::Float64]),
+ );
+ let udaf = Expr::AggregateUDF {
+ fun: Arc::new(my_avg),
+ args: vec![lit(10i64)],
+ filter: None,
+ };
+ let plan = LogicalPlan::Projection(Projection::try_new(vec![udaf],
empty, None)?);
+ let rule = TypeCoercion::new();
+ let mut config = OptimizerConfig::default();
+ let plan = rule.optimize(&plan, &mut config)?;
+ assert_eq!(
+ "Projection: MY_AVG(CAST(Int64(10) AS Float64))\n EmptyRelation",
+ &format!("{:?}", plan)
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn agg_udaf_invalid_input() -> Result<()> {
+ let empty = empty();
+ let return_type: ReturnTypeFunction =
+ Arc::new(move |_| Ok(Arc::new(DataType::Float64)));
+ let state_type: StateTypeFunction =
+ Arc::new(move |_| Ok(Arc::new(vec![DataType::UInt64,
DataType::Float64])));
+ let accumulator: AccumulatorFunctionImplementation =
+ Arc::new(|_|
Ok(Box::new(AvgAccumulator::try_new(&DataType::Float64)?)));
+ let my_avg = AggregateUDF::new(
+ "MY_AVG",
+ &Signature::uniform(1, vec![DataType::Float64],
Volatility::Immutable),
+ &return_type,
+ &accumulator,
+ &state_type,
+ );
+ let udaf = Expr::AggregateUDF {
+ fun: Arc::new(my_avg),
+ args: vec![lit("10")],
+ filter: None,
+ };
+ let plan = LogicalPlan::Projection(Projection::try_new(vec![udaf],
empty, None)?);
+ let rule = TypeCoercion::new();
+ let mut config = OptimizerConfig::default();
+ let plan = rule.optimize(&plan, &mut config);
+ assert!(plan.is_err());
+ assert_eq!(
+ "Plan(\"Coercion from [Utf8] to the signature Uniform(1,
[Float64]) failed.\")",
+ &format!("{:?}", plan.err().unwrap())
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn agg_function_case() -> Result<()> {
+ let empty = empty();
+ let fun: AggregateFunction = AggregateFunction::Avg;
+ let agg_expr = Expr::AggregateFunction {
+ fun,
+ args: vec![lit(12i64)],
+ distinct: false,
+ filter: None,
+ };
+ let plan =
+ LogicalPlan::Projection(Projection::try_new(vec![agg_expr], empty,
None)?);
+ let rule = TypeCoercion::new();
+ let mut config = OptimizerConfig::default();
+ let plan = rule.optimize(&plan, &mut config)?;
+ assert_eq!(
+ "Projection: AVG(Int64(12))\n EmptyRelation",
Review Comment:
I don't see any coercion happening here. Maybe I am missing something
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
