(arrow-datafusion) branch main updated: Add test case for unnecessary hash when target is 1 (#8757)

[akurmustafa]

This is an automated email from the ASF dual-hosted git repository.

akurmustafa pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 05e3d456ea Add test case for unnecessary hash when target is 1 (#8757)
05e3d456ea is described below

commit 05e3d456ea5d8bf537b68787f67c8aa0af9e448a
Author: Mustafa Akur <106137913+mustafasr...@users.noreply.github.com>
AuthorDate: Fri Jan 5 15:20:40 2024 +0300

    Add test case for unnecessary hash when target is 1 (#8757)
---
 .../src/physical_optimizer/enforce_distribution.rs |  3 +-
 datafusion/sqllogictest/test_files/join.slt        | 32 ++++++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs 
b/datafusion/core/src/physical_optimizer/enforce_distribution.rs
index bf5aa7d022..1c86c4c320 100644
--- a/datafusion/core/src/physical_optimizer/enforce_distribution.rs
+++ b/datafusion/core/src/physical_optimizer/enforce_distribution.rs
@@ -927,9 +927,8 @@ fn add_hash_on_top(
     n_target: usize,
     repartition_beneficial_stats: bool,
 ) -> Result<DistributionContext> {
-    let partition_count = input.plan.output_partitioning().partition_count();
     // Early return if hash repartition is unnecessary
-    if n_target == partition_count && n_target == 1 {
+    if n_target == 1 {
         return Ok(input);
     }
 
diff --git a/datafusion/sqllogictest/test_files/join.slt 
b/datafusion/sqllogictest/test_files/join.slt
index c9dd7ca604..ca9b918ff3 100644
--- a/datafusion/sqllogictest/test_files/join.slt
+++ b/datafusion/sqllogictest/test_files/join.slt
@@ -626,6 +626,38 @@ Alice 100 Alice 1
 Alice 50 Alice 2
 Alice 100 Alice 2
 
+statement ok
+set datafusion.execution.target_partitions = 1;
+
+statement ok
+set datafusion.optimizer.repartition_joins = true;
+
+# make sure when target partition is 1, hash repartition is not added
+# to the final plan.
+query TT
+EXPLAIN SELECT *
+FROM t1,
+t1 as t2
+WHERE t1.a=t2.a;
+----
+logical_plan
+Inner Join: t1.a = t2.a
+--TableScan: t1 projection=[a, b]
+--SubqueryAlias: t2
+----TableScan: t1 projection=[a, b]
+physical_plan
+CoalesceBatchesExec: target_batch_size=8192
+--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, a@0)]
+----MemoryExec: partitions=1, partition_sizes=[1]
+----MemoryExec: partitions=1, partition_sizes=[1]
+
+# Reset the configs to old values
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+set datafusion.optimizer.repartition_joins = false;
+
 statement ok
 DROP TABLE t1;