Skip to content
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

Add logical plan support for aggregate expressions with filters (and upgrade to sqlparser 0.23) #3405

Merged
merged 8 commits into from
Sep 12, 2022
Merged

Add logical plan support for aggregate expressions with filters (and upgrade to sqlparser 0.23) #3405

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
1a1957d
Use sqlparser-0.23
andygrove Sep 8, 2022
a4ea9af
Add filter to aggregate expressions
andygrove Sep 8, 2022
d309a3c
clippy
andygrove Sep 10, 2022
7380307
implement protobuf serde
andygrove Sep 10, 2022
baf12a4
clippy
andygrove Sep 12, 2022
9461974
Merge remote-tracking branch 'apache/master' into sqlparser-0.23
andygrove Sep 12, 2022
00b7fdb
fix error message
andygrove Sep 12, 2022
b0d158a
Update datafusion/expr/src/expr.rs
andygrove Sep 12, 2022
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
7 changes: 6 additions & 1 deletion datafusion/core/src/physical_plan/planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,12 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result<String> {
args,
..
} => create_function_physical_name(&fun.to_string(), *distinct, args),
Expr::AggregateUDF { fun, args } => {
Expr::AggregateUDF { fun, args, filter } => {
if filter.is_some() {
return Err(DataFusionError::Execution(
"aggregate expression with filter is not supported".to_string(),
));
}
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_physical_name(e, false)?);
Expand Down
46 changes: 39 additions & 7 deletions datafusion/expr/src/expr.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,8 @@ pub enum Expr {
args: Vec<Expr>,
/// Whether this is a DISTINCT aggregation or not
distinct: bool,
/// Optional filter
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
/// Optional filter
/// Optional filter applied prior to aggregating

filter: Option<Box<Expr>>,
},
/// Represents the call of a window function with arguments.
WindowFunction {
Expand All @@ -251,6 +253,8 @@ pub enum Expr {
fun: Arc<AggregateUDF>,
/// List of expressions to feed to the functions as arguments
args: Vec<Expr>,
/// Optional filter applied prior to aggregating
filter: Option<Box<Expr>>,
},
/// Returns whether the list contains the expr value.
InList {
Expand Down Expand Up @@ -668,10 +672,26 @@ impl fmt::Debug for Expr {
fun,
distinct,
ref args,
filter,
..
} => fmt_function(f, &fun.to_string(), *distinct, args, true),
Expr::AggregateUDF { fun, ref args, .. } => {
fmt_function(f, &fun.name, false, args, false)
} => {
fmt_function(f, &fun.to_string(), *distinct, args, true)?;
if let Some(fe) = filter {
write!(f, " FILTER (WHERE {})", fe)?;
}
Ok(())
}
Expr::AggregateUDF {
fun,
ref args,
filter,
..
} => {
fmt_function(f, &fun.name, false, args, false)?;
if let Some(fe) = filter {
write!(f, " FILTER (WHERE {})", fe)?;
}
Ok(())
}
Expr::Between {
expr,
Expand Down Expand Up @@ -1010,14 +1030,26 @@ fn create_name(e: &Expr) -> Result<String> {
fun,
distinct,
args,
..
} => create_function_name(&fun.to_string(), *distinct, args),
Expr::AggregateUDF { fun, args } => {
filter,
} => {
let name = create_function_name(&fun.to_string(), *distinct, args)?;
if let Some(fe) = filter {
Ok(format!("{} FILTER (WHERE {})", name, fe))
} else {
Ok(name)
}
}
Expr::AggregateUDF { fun, args, filter } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e)?);
}
Ok(format!("{}({})", fun.name, names.join(",")))
let filter = if let Some(fe) = filter {
format!(" FILTER (WHERE {})", fe)
} else {
"".to_string()
};
Ok(format!("{}({}){}", fun.name, names.join(","), filter))
}
Expr::GroupingSet(grouping_set) => match grouping_set {
GroupingSet::Rollup(exprs) => {
Expand Down
10 changes: 10 additions & 0 deletions datafusion/expr/src/expr_fn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ pub fn min(expr: Expr) -> Expr {
fun: aggregate_function::AggregateFunction::Min,
distinct: false,
args: vec![expr],
filter: None,
}
}

Expand All @@ -75,6 +76,7 @@ pub fn max(expr: Expr) -> Expr {
fun: aggregate_function::AggregateFunction::Max,
distinct: false,
args: vec![expr],
filter: None,
}
}

Expand All @@ -84,6 +86,7 @@ pub fn sum(expr: Expr) -> Expr {
fun: aggregate_function::AggregateFunction::Sum,
distinct: false,
args: vec![expr],
filter: None,
}
}

Expand All @@ -93,6 +96,7 @@ pub fn avg(expr: Expr) -> Expr {
fun: aggregate_function::AggregateFunction::Avg,
distinct: false,
args: vec![expr],
filter: None,
}
}

Expand All @@ -102,6 +106,7 @@ pub fn count(expr: Expr) -> Expr {
fun: aggregate_function::AggregateFunction::Count,
distinct: false,
args: vec![expr],
filter: None,
}
}

Expand All @@ -111,6 +116,7 @@ pub fn count_distinct(expr: Expr) -> Expr {
fun: aggregate_function::AggregateFunction::Count,
distinct: true,
args: vec![expr],
filter: None,
}
}

Expand Down Expand Up @@ -163,6 +169,7 @@ pub fn approx_distinct(expr: Expr) -> Expr {
fun: aggregate_function::AggregateFunction::ApproxDistinct,
distinct: false,
args: vec![expr],
filter: None,
}
}

Expand All @@ -172,6 +179,7 @@ pub fn approx_median(expr: Expr) -> Expr {
fun: aggregate_function::AggregateFunction::ApproxMedian,
distinct: false,
args: vec![expr],
filter: None,
}
}

Expand All @@ -181,6 +189,7 @@ pub fn approx_percentile_cont(expr: Expr, percentile: Expr) -> Expr {
fun: aggregate_function::AggregateFunction::ApproxPercentileCont,
distinct: false,
args: vec![expr, percentile],
filter: None,
}
}

Expand All @@ -194,6 +203,7 @@ pub fn approx_percentile_cont_with_weight(
fun: aggregate_function::AggregateFunction::ApproxPercentileContWithWeight,
distinct: false,
args: vec![expr, weight_expr, percentile],
filter: None,
}
}

Expand Down
5 changes: 4 additions & 1 deletion datafusion/expr/src/expr_rewriter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,10 +250,12 @@ impl ExprRewritable for Expr {
args,
fun,
distinct,
filter,
} => Expr::AggregateFunction {
args: rewrite_vec(args, rewriter)?,
fun,
distinct,
filter,
},
Expr::GroupingSet(grouping_set) => match grouping_set {
GroupingSet::Rollup(exprs) => {
Expand All @@ -271,9 +273,10 @@ impl ExprRewritable for Expr {
))
}
},
Expr::AggregateUDF { args, fun } => Expr::AggregateUDF {
Expr::AggregateUDF { args, fun, filter } => Expr::AggregateUDF {
args: rewrite_vec(args, rewriter)?,
fun,
filter,
},
Expr::InList {
expr,
Expand Down
1 change: 1 addition & 0 deletions datafusion/expr/src/udaf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ impl AggregateUDF {
Expr::AggregateUDF {
fun: Arc::new(self.clone()),
args,
filter: None,
}
}
}
6 changes: 5 additions & 1 deletion datafusion/optimizer/src/single_distinct_to_groupby.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,9 @@ fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
let new_aggr_exprs = aggr_expr
.iter()
.map(|aggr_expr| match aggr_expr {
Expr::AggregateFunction { fun, args, .. } => {
Expr::AggregateFunction {
fun, args, filter, ..
} => {
// is_single_distinct_agg ensure args.len=1
if group_fields_set.insert(args[0].name()?) {
inner_group_exprs
Expand All @@ -97,6 +99,7 @@ fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
fun: fun.clone(),
args: vec![col(SINGLE_DISTINCT_ALIAS)],
distinct: false, // intentional to remove distinct here
filter: filter.clone(),
})
}
_ => Ok(aggr_expr.clone()),
Expand Down Expand Up @@ -402,6 +405,7 @@ mod tests {
fun: AggregateFunction::Max,
distinct: true,
args: vec![col("b")],
filter: None,
},
],
)?
Expand Down
2 changes: 2 additions & 0 deletions datafusion/proto/proto/datafusion.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -504,11 +504,13 @@ message AggregateExprNode {
AggregateFunction aggr_function = 1;
repeated LogicalExprNode expr = 2;
bool distinct = 3;
LogicalExprNode filter = 4;
}

message AggregateUDFExprNode {
string fun_name = 1;
repeated LogicalExprNode args = 2;
LogicalExprNode filter = 3;
}

message ScalarUDFExprNode {
Expand Down
9 changes: 6 additions & 3 deletions datafusion/proto/src/from_proto.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -891,6 +891,7 @@ pub fn parse_expr(
.map(|e| parse_expr(e, registry))
.collect::<Result<Vec<_>, _>>()?,
distinct: expr.distinct,
filter: parse_optional_expr(&expr.filter, registry)?.map(Box::new),
})
}
ExprType::Alias(alias) => Ok(Expr::Alias(
Expand Down Expand Up @@ -1194,15 +1195,17 @@ pub fn parse_expr(
.collect::<Result<Vec<_>, Error>>()?,
})
}
ExprType::AggregateUdfExpr(protobuf::AggregateUdfExprNode { fun_name, args }) => {
let agg_fn = registry.udaf(fun_name.as_str())?;
ExprType::AggregateUdfExpr(pb) => {
let agg_fn = registry.udaf(pb.fun_name.as_str())?;

Ok(Expr::AggregateUDF {
fun: agg_fn,
args: args
args: pb
.args
.iter()
.map(|expr| parse_expr(expr, registry))
.collect::<Result<Vec<_>, Error>>()?,
filter: parse_optional_expr(&pb.filter, registry)?.map(Box::new),
})
}

Expand Down
4 changes: 4 additions & 0 deletions datafusion/proto/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1023,6 +1023,7 @@ mod roundtrip_tests {
fun: AggregateFunction::Count,
args: vec![col("bananas")],
distinct: false,
filter: None,
};
let ctx = SessionContext::new();
roundtrip_expr_test(test_expr, ctx);
Expand All @@ -1034,6 +1035,7 @@ mod roundtrip_tests {
fun: AggregateFunction::Count,
args: vec![col("bananas")],
distinct: true,
filter: None,
};
let ctx = SessionContext::new();
roundtrip_expr_test(test_expr, ctx);
Expand All @@ -1045,6 +1047,7 @@ mod roundtrip_tests {
fun: AggregateFunction::ApproxPercentileCont,
args: vec![col("bananas"), lit(0.42_f32)],
distinct: false,
filter: None,
};

let ctx = SessionContext::new();
Expand Down Expand Up @@ -1097,6 +1100,7 @@ mod roundtrip_tests {
let test_expr = Expr::AggregateUDF {
fun: Arc::new(dummy_agg.clone()),
args: vec![lit(1.0_f64)],
filter: Some(Box::new(lit(true))),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

};

let mut ctx = SessionContext::new();
Expand Down
35 changes: 23 additions & 12 deletions datafusion/proto/src/to_proto.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -585,6 +585,7 @@ impl TryFrom<&Expr> for protobuf::LogicalExprNode {
ref fun,
ref args,
ref distinct,
ref filter
} => {
let aggr_function = match fun {
AggregateFunction::ApproxDistinct => {
Expand Down Expand Up @@ -633,9 +634,13 @@ impl TryFrom<&Expr> for protobuf::LogicalExprNode {
.map(|v| v.try_into())
.collect::<Result<Vec<_>, _>>()?,
distinct: *distinct,
filter: match filter {
Some(e) => Some(Box::new(e.as_ref().try_into()?)),
None => None,
}
};
Self {
expr_type: Some(ExprType::AggregateExpr(aggregate_expr)),
expr_type: Some(ExprType::AggregateExpr(Box::new(aggregate_expr))),
}
}
Expr::ScalarVariable(_, _) => return Err(Error::General("Proto serialization error: Scalar Variable not supported".to_string())),
Expand Down Expand Up @@ -663,17 +668,23 @@ impl TryFrom<&Expr> for protobuf::LogicalExprNode {
.collect::<Result<Vec<_>, Error>>()?,
})),
},
Expr::AggregateUDF { fun, args } => Self {
expr_type: Some(ExprType::AggregateUdfExpr(
protobuf::AggregateUdfExprNode {
fun_name: fun.name.clone(),
args: args.iter().map(|expr| expr.try_into()).collect::<Result<
Vec<_>,
Error,
>>(
)?,
},
)),
Expr::AggregateUDF { fun, args, filter } => {
Self {
expr_type: Some(ExprType::AggregateUdfExpr(
Box::new(protobuf::AggregateUdfExprNode {
fun_name: fun.name.clone(),
args: args.iter().map(|expr| expr.try_into()).collect::<Result<
Vec<_>,
Error,
>>(
)?,
filter: match filter {
Some(e) => Some(Box::new(e.as_ref().try_into()?)),
None => None,
}
Comment on lines +681 to +684
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you can write this kind of thing somewhat more concisely via

Suggested change
filter: match filter {
Some(e) => Some(Box::new(e.as_ref().try_into()?)),
None => None,
}
filter: filter
.map(|filter| Box::new(e.as_ref().try_into()))
.transpose()?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I experimented with this locally but could not get this to work.

},
))),
}
},
Expr::Not(expr) => {
let expr = Box::new(protobuf::Not {
Expand Down
Loading