[Fix](bug) fix the divide zero in local shuffle: (#37948)

## Proposed changes cherry pick #37906
apache · Jul 16, 2024 · b6e5281 · b6e5281
1 parent 21c6b85
commit b6e5281
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 1 deletion.
diff --git a/be/src/pipeline/pipeline_x/pipeline_x_fragment_context.cpp b/be/src/pipeline/pipeline_x/pipeline_x_fragment_context.cpp
@@ -298,8 +298,12 @@ Status PipelineXFragmentContext::_plan_local_exchange(
             }
         }
 
+        // if 'num_buckets == 0' means the fragment is colocated by exchange node not the
+        // scan node. so here use `_num_instance` to replace the `num_buckets` to prevent dividing 0
+        // still keep colocate plan after local shuffle
         RETURN_IF_ERROR(_plan_local_exchange(
-                _pipelines[pip_idx]->operator_xs().front()->ignore_data_hash_distribution()
+                _pipelines[pip_idx]->operator_xs().front()->ignore_data_hash_distribution() ||
+                                num_buckets == 0
                         ? _num_instances
                         : num_buckets,
                 pip_idx, _pipelines[pip_idx], bucket_seq_to_instance_idx,

diff --git a/regression-test/data/query_p0/limit/sql/withGroupByUnion.out b/regression-test/data/query_p0/limit/sql/withGroupByUnion.out
@@ -0,0 +1,52 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !withGroupByUnion --
+0	ALGERIA
+1	ALGERIA
+1	ARGENTINA
+1	BRAZIL
+1	CANADA
+1	CHINA
+1	EGYPT
+1	ETHIOPIA
+1	FRANCE
+1	GERMANY
+1	INDIA
+1	INDONESIA
+1	IRAN
+1	IRAQ
+1	JAPAN
+1	JORDAN
+1	KENYA
+1	MOROCCO
+1	MOZAMBIQUE
+1	PERU
+1	ROMANIA
+1	RUSSIA
+1	SAUDI ARABIA
+1	UNITED KINGDOM
+1	UNITED STATES
+1	VIETNAM
+2	BRAZIL
+3	CANADA
+4	EGYPT
+5	ETHIOPIA
+6	FRANCE
+7	GERMANY
+8	INDIA
+9	INDONESIA
+10	IRAN
+11	IRAQ
+12	JAPAN
+13	JORDAN
+14	KENYA
+15	MOROCCO
+16	MOZAMBIQUE
+17	PERU
+18	CHINA
+19	ROMANIA
+20	SAUDI ARABIA
+21	VIETNAM
+22	RUSSIA
+23	UNITED KINGDOM
+24	UNITED STATES
+
diff --git a/regression-test/suites/query_p0/limit/sql/withGroupByUnion.sql b/regression-test/suites/query_p0/limit/sql/withGroupByUnion.sql
@@ -0,0 +1 @@
+select * from (select count(a.nationkey), a.name from tpch_tiny_nation a  join[shuffle] tpch_tiny_nation b on a.name = b.name group by a.name union select sum(c.nationkey), c.name from tpch_tiny_nation c  join[shuffle] tpch_tiny_nation d on c.name = d.name group by c.name) t order by 1,2 limit 50;
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		select * from (select count(a.nationkey), a.name from tpch_tiny_nation a join[shuffle] tpch_tiny_nation b on a.name = b.name group by a.name union select sum(c.nationkey), c.name from tpch_tiny_nation c join[shuffle] tpch_tiny_nation d on c.name = d.name group by c.name) t order by 1,2 limit 50;