Fix aggregate type coercion bug #3710

alamb · 2022-10-04T18:58:10Z

~~DRAFT -- I am just pushing it up in case anyone else finds the debugging via logging helpful~~

Which issue does this PR close?

Closes #3704

Rationale for this change

There is a planning regression

What changes are included in this PR?

Fix bug (will comment inline)

Are there any user-facing changes?

alamb · 2022-10-04T20:39:48Z

datafusion/optimizer/src/optimizer.rs

@@ -178,16 +178,15 @@ impl Optimizer {
        F: FnMut(&LogicalPlan, &dyn OptimizerRule),
    {
        let mut new_plan = plan.clone();
-        debug!("Input logical plan:\n{}\n", plan.display_indent());
-        trace!("Full input logical plan:\n{:?}", plan);
+        log_plan("Optimizer input", plan);


This is a driveby cleanup to improve logging (specifically, also add trace! to log schema)

alamb · 2022-10-04T20:40:43Z

datafusion/core/tests/sql/subqueries.rs

@@ -147,8 +147,8 @@ order by s_acctbal desc, n_name, s_name, p_partkey;"#;
          Inner Join: #supplier.s_nationkey = #nation.n_nationkey
            Inner Join: #partsupp.ps_suppkey = #supplier.s_suppkey
              Inner Join: #part.p_partkey = #partsupp.ps_partkey
-                Filter: #part.p_size = Int32(15) AND #part.p_type LIKE Utf8("%BRASS")
-                  TableScan: part projection=[p_partkey, p_mfgr, p_type, p_size], partial_filters=[#part.p_size = Int32(15), #part.p_type LIKE Utf8("%BRASS")]
+                Filter: CAST(#part.p_size AS Int64) = Int64(15) AND #part.p_type LIKE Utf8("%BRASS")


to be honest, I am not sure why this has changed (aka the filters are no longer simplified). I will look into that in the morning

alamb · 2022-10-04T20:43:37Z

datafusion/optimizer/src/unwrap_cast_in_comparison.rs

@@ -97,7 +97,18 @@ fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
    let new_exprs = plan
        .expressions()
        .into_iter()
-        .map(|expr| expr.rewrite(&mut expr_rewriter))
+        .map(|expr| {
+            let original_name = expr.name()?;


I think the root cause issue is that the UnwrapCastInComparison can add potentially change the expression but it doesn't add an alias so the output name changes

Bigger picture there are at least three places that we have rediscovered this same problem when rewriting expressions --#3555 and https://github.com/apache/arrow-datafusion/blob/master/datafusion/optimizer/src/simplify_expressions.rs#L316 I will try and make a follow on PR to clean them all up. In particular, I think this is something from_plan could potentially handle

expr.name() here hides the casts that are added by this optimization rule, so expr.name() is the same as the original name (even though the expression is now different), and the alias does not get added.

I think this optimizer rule may actually change the name (e.g. from Int64(0) to Int32(0)) which i think is the root cause of the issue in this bug

andygrove · 2022-10-05T16:51:16Z

datafusion/optimizer/tests/integration-test.rs

+fn case_when_aggregate() -> Result<()> {
+    let sql = "SELECT col_utf8, SUM(CASE WHEN col_int32 > 0 THEN 1 ELSE 0 END) AS n FROM test GROUP BY col_utf8";
+    let plan = test_sql(sql)?;
+    let expected = "Projection: #test.col_utf8, #SUM(CASE WHEN test.col_int32 > Int64(0) THEN Int64(1) ELSE Int64(0) END) AS n\


Could you merge latest from master - we should not include # before column names now

in e3830a7

andygrove

LGTM. Thanks @alamb

andygrove · 2022-10-05T17:03:58Z

I tested this with the Dask SQL test suite and it fixed 5 queries 🚀 .. more than I expected

alamb · 2022-10-05T17:04:31Z

datafusion/optimizer/src/unwrap_cast_in_comparison.rs

        .collect::<Result<Vec<_>>>()?;

    from_plan(plan, new_exprs.as_slice(), new_inputs.as_slice())
 }

+fn name_for_alias(expr: &Expr) -> Result<String> {


I plan to make this easier on the eyes as a follow on PR

follow on #3727

alamb · 2022-10-05T17:04:48Z

datafusion/optimizer/tests/integration-test.rs

 #[test]
 fn case_when() -> Result<()> {
    let sql = "SELECT CASE WHEN col_int32 > 0 THEN 1 ELSE 0 END FROM test";
    let plan = test_sql(sql)?;
    let expected =
-        "Projection: CASE WHEN test.col_int32 > Int32(0) THEN Int64(1) ELSE Int64(0) END\
-    \n  TableScan: test projection=[col_int32]";
+        "Projection: CASE WHEN test.col_int32 > Int32(0) THEN Int64(1) ELSE Int64(0) END AS CASE WHEN test.col_int32 > Int64(0) THEN Int64(1) ELSE Int64(0) END\


FYI @andygrove the alias was added to this as well

andygrove · 2022-10-05T18:28:46Z

Thanks again @alamb. This was a huge improvement.

alamb · 2022-10-05T18:30:11Z

I am pretty excited where the datafusion optimizer is headign

ursabot · 2022-10-05T18:34:59Z

Benchmark runs are scheduled for baseline = 965133c and contender = 64669e9. 64669e9 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on ec2-t3-xlarge-us-east-2] ec2-t3-xlarge-us-east-2
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on test-mac-arm] test-mac-arm
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on ursa-i9-9960x] ursa-i9-9960x
[Skipped ⚠️ Benchmarking of arrow-datafusion-commits is not supported on ursa-thinkcentre-m75q] ursa-thinkcentre-m75q
Buildkite builds:
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java

liukun4515 · 2022-10-06T03:48:01Z

I think this fix is related to #3555
#3583

alamb marked this pull request as draft October 4, 2022 18:58

github-actions bot added the optimizer Optimizer rules label Oct 4, 2022

alamb added 2 commits October 4, 2022 16:30

Do not change output expr name in UnwrapCastInComparison

af3322a

Update

b486da7

alamb force-pushed the alamb/optimizer_error branch from 3edc1cd to b486da7 Compare October 4, 2022 20:38

github-actions bot added the core Core datafusion crate label Oct 4, 2022

alamb commented Oct 4, 2022

View reviewed changes

alamb changed the title ~~WIP fix aggregate type coercion bug~~ Fix aggregate type coercion bug Oct 4, 2022

alamb marked this pull request as ready for review October 4, 2022 20:40

alamb commented Oct 4, 2022

View reviewed changes

alamb mentioned this pull request Oct 4, 2022

Optimizer rule 'projection_push_down' failed due to unexpected error: Error during planning: Aggregate schema has wrong number of fields. Expected 3 got 8 #3704

Closed

andygrove mentioned this pull request Oct 5, 2022

Fix aggregate type coercion bug #3707

Closed

alamb added 3 commits October 5, 2022 11:59

Merge remote-tracking branch 'apache/master' into alamb/optimizer_error

e71aeed

Update test

edc2561

Fix regression

a9d6f8c

andygrove reviewed Oct 5, 2022

View reviewed changes

andygrove approved these changes Oct 5, 2022

View reviewed changes

alamb force-pushed the alamb/optimizer_error branch from f96e318 to a9d6f8c Compare October 5, 2022 16:57

github-actions bot removed the core Core datafusion crate label Oct 5, 2022

Update tests

e3830a7

alamb commented Oct 5, 2022

View reviewed changes

clippy

f6e8ffa

alamb mentioned this pull request Oct 5, 2022

Consolidate and better tests for expression re-rewriting / aliasing #3727

Merged

andygrove merged commit 64669e9 into apache:master Oct 5, 2022

alamb deleted the alamb/optimizer_error branch October 5, 2022 18:29

thinkharderdev approved these changes Oct 5, 2022

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix aggregate type coercion bug #3710

Fix aggregate type coercion bug #3710

alamb commented Oct 4, 2022 •

edited

alamb Oct 4, 2022

alamb Oct 4, 2022

alamb Oct 4, 2022

andygrove Oct 4, 2022

alamb Oct 4, 2022

andygrove Oct 5, 2022

alamb Oct 5, 2022

alamb Oct 5, 2022

andygrove left a comment •

edited

andygrove commented Oct 5, 2022

alamb Oct 5, 2022

alamb Oct 5, 2022

alamb Oct 5, 2022

andygrove commented Oct 5, 2022

alamb commented Oct 5, 2022

ursabot commented Oct 5, 2022

liukun4515 commented Oct 6, 2022

Fix aggregate type coercion bug #3710

Fix aggregate type coercion bug #3710

Conversation

alamb commented Oct 4, 2022 • edited

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andygrove left a comment • edited

Choose a reason for hiding this comment

andygrove commented Oct 5, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andygrove commented Oct 5, 2022

alamb commented Oct 5, 2022

ursabot commented Oct 5, 2022

liukun4515 commented Oct 6, 2022

alamb commented Oct 4, 2022 •

edited

andygrove left a comment •

edited