Skip to content

Commit 2580bc5

Browse files
kumarUjjawalkumarUjjawal
committed
show lowered expressions in physical explain
1 parent 33e47c9 commit 2580bc5

14 files changed

Lines changed: 265 additions & 229 deletions

File tree

‎datafusion/core/src/physical_planner.rs‎

Lines changed: 62 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -2401,7 +2401,7 @@ type AggregateExprWithOptionalArgs = (
24012401
pub fn create_aggregate_expr_with_name_and_maybe_filter(
24022402
e: &Expr,
24032403
name: Option<String>,
2404-
human_display: String,
2404+
human_display: Option<String>,
24052405
logical_input_schema: &DFSchema,
24062406
physical_input_schema: &Schema,
24072407
execution_props: &ExecutionProps,
@@ -2445,16 +2445,17 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
24452445
execution_props,
24462446
)?;
24472447

2448-
let agg_expr =
2448+
let mut builder =
24492449
AggregateExprBuilder::new(func.to_owned(), physical_args.to_vec())
24502450
.order_by(order_bys.clone())
24512451
.schema(Arc::new(physical_input_schema.to_owned()))
24522452
.alias(name)
2453-
.human_display(human_display)
24542453
.with_ignore_nulls(ignore_nulls)
2455-
.with_distinct(*distinct)
2456-
.build()
2457-
.map(Arc::new)?;
2454+
.with_distinct(*distinct);
2455+
if let Some(human_display) = human_display {
2456+
builder = builder.human_display(human_display);
2457+
}
2458+
let agg_expr = builder.build().map(Arc::new)?;
24582459

24592460
(agg_expr, filter, order_bys)
24602461
};
@@ -2472,32 +2473,28 @@ pub fn create_aggregate_expr_and_maybe_filter(
24722473
physical_input_schema: &Schema,
24732474
execution_props: &ExecutionProps,
24742475
) -> Result<AggregateExprWithOptionalArgs> {
2475-
// Unpack (potentially nested) aliased logical expressions, e.g. "sum(col) as total"
2476-
// Some functions like `count_all()` create internal aliases,
2477-
// Unwrap all alias layers to get to the underlying aggregate function
2476+
// Unpack (potentially nested) aliased logical expressions, e.g. "sum(col) as total".
2477+
// Physical explain prefers the lowered aggregate form, so unwrap all alias
2478+
// layers to recover the underlying aggregate function and then re-attach
2479+
// only the visible output alias.
24782480
let (name, human_display, e) = match e {
24792481
Expr::Alias(alias) => {
24802482
let unaliased = e.clone().unalias_nested().data;
2481-
let human_display = if should_inline_aliased_aggregate_expr(alias) {
2482-
let human_display =
2483-
inherited_internal_aggregate_display(alias.expr.as_ref())
2484-
.unwrap_or_else(|| unaliased.human_display().to_string());
2485-
if human_display.is_empty() || human_display == alias.name {
2486-
alias.name.clone()
2487-
} else {
2488-
format!("{human_display} as {}", alias.name)
2489-
}
2490-
} else {
2483+
let human_display = unaliased.human_display().to_string();
2484+
let human_display = if human_display.is_empty() || human_display == alias.name
2485+
{
24912486
alias.name.clone()
2487+
} else {
2488+
format!("{human_display} as {}", alias.name)
24922489
};
2493-
(Some(alias.name.clone()), human_display, unaliased)
2490+
(Some(alias.name.clone()), Some(human_display), unaliased)
24942491
}
24952492
Expr::AggregateFunction(_) => (
24962493
Some(e.schema_name().to_string()),
2497-
e.human_display().to_string(),
2494+
Some(e.human_display().to_string()),
24982495
e.clone(),
24992496
),
2500-
_ => (None, String::default(), e.clone()),
2497+
_ => (None, None, e.clone()),
25012498
};
25022499

25032500
create_aggregate_expr_with_name_and_maybe_filter(
@@ -2510,25 +2507,6 @@ pub fn create_aggregate_expr_and_maybe_filter(
25102507
)
25112508
}
25122509

2513-
fn should_inline_aliased_aggregate_expr(alias: &Alias) -> bool {
2514-
!alias.is_internal
2515-
}
2516-
2517-
fn inherited_internal_aggregate_display(expr: &Expr) -> Option<String> {
2518-
let mut expr = expr;
2519-
while let Expr::Alias(alias) = expr {
2520-
let has_user_metadata = alias
2521-
.metadata
2522-
.as_ref()
2523-
.is_some_and(|metadata| !metadata.is_empty());
2524-
if alias.is_internal && !has_user_metadata {
2525-
return Some(alias.name.clone());
2526-
}
2527-
expr = alias.expr.as_ref();
2528-
}
2529-
None
2530-
}
2531-
25322510
impl DefaultPhysicalPlanner {
25332511
/// Handles capturing the various plans for EXPLAIN queries
25342512
///
@@ -3198,49 +3176,6 @@ mod tests {
31983176
Ok(displayable(physical_plan.as_ref()).indent(true).to_string())
31993177
}
32003178

3201-
#[test]
3202-
fn test_should_inline_aliased_aggregate_expr() {
3203-
let Expr::Alias(user_alias) = sum(col("column1")).alias("total rows") else {
3204-
panic!("expected alias");
3205-
};
3206-
assert!(should_inline_aliased_aggregate_expr(&user_alias));
3207-
3208-
let Expr::Alias(count_all_alias) = count_all() else {
3209-
panic!("expected alias");
3210-
};
3211-
assert!(!should_inline_aliased_aggregate_expr(&count_all_alias));
3212-
3213-
let Expr::Alias(internal_alias) =
3214-
sum(col("column1")).alias_internal("sum(random() + Int64(1))")
3215-
else {
3216-
panic!("expected alias");
3217-
};
3218-
assert!(!should_inline_aliased_aggregate_expr(&internal_alias));
3219-
}
3220-
3221-
#[test]
3222-
fn test_inherited_internal_aggregate_display() {
3223-
let expr = count_all().alias_internal("row_count").alias("total_rows");
3224-
let Expr::Alias(alias) = expr else {
3225-
panic!("expected alias");
3226-
};
3227-
assert_eq!(
3228-
inherited_internal_aggregate_display(alias.expr.as_ref()),
3229-
Some("row_count".to_string())
3230-
);
3231-
3232-
let expr = count_all()
3233-
.alias_internal("inner_count")
3234-
.alias_internal("outer_count");
3235-
assert_eq!(
3236-
inherited_internal_aggregate_display(&expr),
3237-
Some("outer_count".to_string())
3238-
);
3239-
3240-
let expr = sum(col("column1")).alias("total_rows");
3241-
assert_eq!(inherited_internal_aggregate_display(&expr), None);
3242-
}
3243-
32443179
#[derive(Debug, Default)]
32453180
struct NullAccumulator;
32463181

@@ -3925,7 +3860,30 @@ mod tests {
39253860
}
39263861

39273862
#[tokio::test]
3928-
async fn test_aggregate_explain_shows_aliased_expression() -> Result<()> {
3863+
async fn test_aggregate_explain_shows_quoted_user_alias() -> Result<()> {
3864+
let schema = Arc::new(Schema::new(vec![Field::new(
3865+
"column1",
3866+
DataType::Int64,
3867+
false,
3868+
)]));
3869+
3870+
let logical_plan = scan_empty(None, schema.as_ref(), None)?
3871+
.aggregate(
3872+
Vec::<Expr>::new(),
3873+
vec![sum(col("column1")).alias("total rows")],
3874+
)?
3875+
.build()?;
3876+
3877+
assert_contains!(
3878+
aggregate_explain(&logical_plan).await?,
3879+
"AggregateExec: mode=Single, gby=[], aggr=[sum(?table?.column1) as total rows]"
3880+
);
3881+
3882+
Ok(())
3883+
}
3884+
3885+
#[tokio::test]
3886+
async fn test_aggregate_explain_shows_aliased_filter_expression() -> Result<()> {
39293887
let schema = Arc::new(Schema::new(vec![
39303888
Field::new("column1", DataType::Int64, false),
39313889
Field::new("column2", DataType::Int64, false),
@@ -3952,23 +3910,29 @@ mod tests {
39523910
}
39533911

39543912
#[tokio::test]
3955-
async fn test_aggregate_explain_shows_quoted_user_alias() -> Result<()> {
3956-
let schema = Arc::new(Schema::new(vec![Field::new(
3957-
"column1",
3958-
DataType::Int64,
3959-
false,
3960-
)]));
3913+
async fn test_aggregate_explain_shows_aliased_respect_nulls() -> Result<()> {
3914+
let schema = Arc::new(Schema::new(vec![
3915+
Field::new("column1", DataType::Int64, true),
3916+
Field::new("column2", DataType::Int64, false),
3917+
]));
39613918

39623919
let logical_plan = scan_empty(None, schema.as_ref(), None)?
39633920
.aggregate(
39643921
Vec::<Expr>::new(),
3965-
vec![sum(col("column1")).alias("total rows")],
3922+
vec![
3923+
first_value_udaf()
3924+
.call(vec![col("column1")])
3925+
.order_by(vec![col("column2").sort(true, true)])
3926+
.null_treatment(NullTreatment::RespectNulls)
3927+
.build()?
3928+
.alias("agg"),
3929+
],
39663930
)?
39673931
.build()?;
39683932

39693933
assert_contains!(
39703934
aggregate_explain(&logical_plan).await?,
3971-
"AggregateExec: mode=Single, gby=[], aggr=[sum(?table?.column1) as total rows]"
3935+
"AggregateExec: mode=Single, gby=[], aggr=[first_value(?table?.column1) RESPECT NULLS ORDER BY [?table?.column2 ASC NULLS FIRST] as agg]"
39723936
);
39733937

39743938
Ok(())
@@ -3981,7 +3945,10 @@ mod tests {
39813945
.aggregate(Vec::<Expr>::new(), vec![count_all()])?
39823946
.build()?;
39833947

3984-
assert_contains!(aggregate_explain(&logical_plan).await?, "aggr=[count(*)]");
3948+
assert_contains!(
3949+
aggregate_explain(&logical_plan).await?,
3950+
"aggr=[count(1) as count(*)]"
3951+
);
39853952

39863953
Ok(())
39873954
}
@@ -3995,36 +3962,7 @@ mod tests {
39953962

39963963
assert_contains!(
39973964
aggregate_explain(&logical_plan).await?,
3998-
"aggr=[count(*) as total_rows]"
3999-
);
4000-
4001-
Ok(())
4002-
}
4003-
4004-
#[tokio::test]
4005-
async fn test_aggregate_explain_shows_aliased_respect_nulls() -> Result<()> {
4006-
let schema = Arc::new(Schema::new(vec![
4007-
Field::new("column1", DataType::Int64, true),
4008-
Field::new("column2", DataType::Int64, false),
4009-
]));
4010-
4011-
let logical_plan = scan_empty(None, schema.as_ref(), None)?
4012-
.aggregate(
4013-
Vec::<Expr>::new(),
4014-
vec![
4015-
first_value_udaf()
4016-
.call(vec![col("column1")])
4017-
.order_by(vec![col("column2").sort(true, true)])
4018-
.null_treatment(NullTreatment::RespectNulls)
4019-
.build()?
4020-
.alias("agg"),
4021-
],
4022-
)?
4023-
.build()?;
4024-
4025-
assert_contains!(
4026-
aggregate_explain(&logical_plan).await?,
4027-
"AggregateExec: mode=Single, gby=[], aggr=[first_value(?table?.column1) RESPECT NULLS ORDER BY [?table?.column2 ASC NULLS FIRST] as agg]"
3965+
"aggr=[count(1) as total_rows]"
40283966
);
40293967

40303968
Ok(())

‎datafusion/core/tests/dataframe/mod.rs‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3022,20 +3022,20 @@ async fn test_count_wildcard_on_sort() -> Result<()> {
30223022
assert_snapshot!(
30233023
pretty_format_batches(&df_results).unwrap(),
30243024
@r"
3025-
+---------------+----------------------------------------------------------------------------+
3026-
| plan_type | plan |
3027-
+---------------+----------------------------------------------------------------------------+
3028-
| logical_plan | Sort: count(*) AS count(*) ASC NULLS LAST |
3029-
| | Aggregate: groupBy=[[t1.b]], aggr=[[count(Int64(1)) AS count(*)]] |
3030-
| | TableScan: t1 projection=[b] |
3031-
| physical_plan | SortPreservingMergeExec: [count(*)@1 ASC NULLS LAST] |
3032-
| | SortExec: expr=[count(*)@1 ASC NULLS LAST], preserve_partitioning=[true] |
3033-
| | AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[count(*)] |
3034-
| | RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=1 |
3035-
| | AggregateExec: mode=Partial, gby=[b@0 as b], aggr=[count(*)] |
3036-
| | DataSourceExec: partitions=1, partition_sizes=[1] |
3037-
| | |
3038-
+---------------+----------------------------------------------------------------------------+
3025+
+---------------+---------------------------------------------------------------------------------------+
3026+
| plan_type | plan |
3027+
+---------------+---------------------------------------------------------------------------------------+
3028+
| logical_plan | Sort: count(*) AS count(*) ASC NULLS LAST |
3029+
| | Aggregate: groupBy=[[t1.b]], aggr=[[count(Int64(1)) AS count(*)]] |
3030+
| | TableScan: t1 projection=[b] |
3031+
| physical_plan | SortPreservingMergeExec: [count(*)@1 ASC NULLS LAST] |
3032+
| | SortExec: expr=[count(*)@1 ASC NULLS LAST], preserve_partitioning=[true] |
3033+
| | AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[count(1) as count(*)] |
3034+
| | RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=1 |
3035+
| | AggregateExec: mode=Partial, gby=[b@0 as b], aggr=[count(1) as count(*)] |
3036+
| | DataSourceExec: partitions=1, partition_sizes=[1] |
3037+
| | |
3038+
+---------------+---------------------------------------------------------------------------------------+
30393039
"
30403040
);
30413041
Ok(())
@@ -3500,9 +3500,9 @@ async fn test_count_wildcard_on_where_scalar_subquery() -> Result<()> {
35003500
| | HashJoinExec: mode=CollectLeft, join_type=Right, on=[(a@1, a@0)], projection=[a@3, b@4, count(*)@0, __always_true@2] |
35013501
| | CoalescePartitionsExec |
35023502
| | ProjectionExec: expr=[count(*)@1 as count(*), a@0 as a, true as __always_true] |
3503-
| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(*)] |
3503+
| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(1) as count(*)] |
35043504
| | RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1 |
3505-
| | AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(*)] |
3505+
| | AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(1) as count(*)] |
35063506
| | DataSourceExec: partitions=1, partition_sizes=[1] |
35073507
| | DataSourceExec: partitions=1, partition_sizes=[1] |
35083508
| | |

0 commit comments

Comments
 (0)