update

Signed-off-by: Weizhen Wang <[email protected]>

Author: hawkingrei

Diff for: pkg/planner/core/casetest/tpch/BUILD.bazel

@@ -9,7 +9,7 @@ go_test(
     ],
     data = glob(["testdata/**"]),
     flaky = True,
-    shard_count = 3,
+    shard_count = 4,
     deps = [
         "//pkg/config",
         "//pkg/testkit",

Diff for: pkg/planner/core/casetest/tpch/testdata/tpch_suite_in.json

@@ -7,6 +7,12 @@
       "explain format='brief' SELECT /*+ SHUFFLE_JOIN(orders, lineitem) */ o.o_orderdate, SUM(l.l_extendedprice * (1 - l.l_discount)) AS revenue FROM orders AS o JOIN lineitem AS l ON o.o_orderkey = l.l_orderkey WHERE o.o_orderdate BETWEEN '1994-01-01' AND '1994-12-31' GROUP BY o.o_orderdate ORDER BY revenue DESC LIMIT 10;"
     ]
   },
+  {
+    "name": "TestQ9",
+    "cases": [
+      "explain format='brief' SELECT nation, o_year, SUM(amount) AS sum_profit FROM (SELECT n_name AS nation, EXTRACT(YEAR FROM o_orderdate) AS o_year, l_extendedprice * (1 - l_discount) - ps_supplycost * l_quantity AS amount FROM part, supplier, lineitem, partsupp, orders, nation WHERE s_suppkey = l_suppkey AND ps_suppkey = l_suppkey AND ps_partkey = l_partkey AND p_partkey = l_partkey AND o_orderkey = l_orderkey AND s_nationkey = n_nationkey AND p_name LIKE '%dim%') AS profit GROUP BY nation, o_year ORDER BY nation, o_year DESC;"
+    ]
+  },
   {
     "name": "TestQ13",
     "cases": [

Diff for: pkg/planner/core/casetest/tpch/testdata/tpch_suite_out.json

@@ -73,6 +73,64 @@
       }
     ]
   },
+  {
+    "Name": "TestQ9",
+    "Cases": [
+      {
+        "SQL": "explain format='brief' SELECT nation, o_year, SUM(amount) AS sum_profit FROM (SELECT n_name AS nation, EXTRACT(YEAR FROM o_orderdate) AS o_year, l_extendedprice * (1 - l_discount) - ps_supplycost * l_quantity AS amount FROM part, supplier, lineitem, partsupp, orders, nation WHERE s_suppkey = l_suppkey AND ps_suppkey = l_suppkey AND ps_partkey = l_partkey AND p_partkey = l_partkey AND o_orderkey = l_orderkey AND s_nationkey = n_nationkey AND p_name LIKE '%dim%') AS profit GROUP BY nation, o_year ORDER BY nation, o_year DESC;",
+        "Result": [
+          "Sort 8000.00 root  test.nation.n_name, Column#52:desc",
+          "└─TableReader 8000.00 root  MppVersion: 3, data:ExchangeSender",
+          "  └─ExchangeSender 8000.00 mpp[tiflash]  ExchangeType: PassThrough",
+          "    └─Projection 8000.00 mpp[tiflash]  test.nation.n_name, Column#52, Column#54",
+          "      └─Projection 8000.00 mpp[tiflash]  Column#54, test.nation.n_name, Column#52",
+          "        └─HashAgg 8000.00 mpp[tiflash]  group by:Column#72, test.nation.n_name, funcs:sum(Column#73)->Column#54, funcs:firstrow(test.nation.n_name)->test.nation.n_name, funcs:firstrow(Column#72)->Column#52",
+          "          └─ExchangeReceiver 8000.00 mpp[tiflash]  ",
+          "            └─ExchangeSender 8000.00 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.nation.n_name, collate: utf8mb4_bin]",
+          "              └─HashAgg 8000.00 mpp[tiflash]  group by:Column#77, Column#78, funcs:sum(Column#76)->Column#73",
+          "                └─Projection 24414.06 mpp[tiflash]  minus(mul(test.lineitem.l_extendedprice, minus(1, test.lineitem.l_discount)), mul(test.partsupp.ps_supplycost, test.lineitem.l_quantity))->Column#76, test.nation.n_name->Column#77, extract(YEAR, test.orders.o_orderdate)->Column#78",
+          "                  └─Projection 24414.06 mpp[tiflash]  test.lineitem.l_quantity, test.lineitem.l_extendedprice, test.lineitem.l_discount, test.partsupp.ps_supplycost, test.orders.o_orderdate, test.nation.n_name",
+          "                    └─Projection 24414.06 mpp[tiflash]  test.lineitem.l_quantity, test.lineitem.l_extendedprice, test.lineitem.l_discount, test.partsupp.ps_supplycost, test.orders.o_orderdate, test.nation.n_name, test.supplier.s_nationkey",
+          "                      └─HashJoin 24414.06 mpp[tiflash]  inner join, equal:[eq(test.supplier.s_nationkey, test.nation.n_nationkey)]",
+          "                        ├─ExchangeReceiver(Build) 10000.00 mpp[tiflash]  ",
+          "                        │ └─ExchangeSender 10000.00 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.nation.n_nationkey, collate: binary]",
+          "                        │   └─TableFullScan 10000.00 mpp[tiflash] table:nation keep order:false, stats:pseudo",
+          "                        └─ExchangeReceiver(Probe) 19531.25 mpp[tiflash]  ",
+          "                          └─ExchangeSender 19531.25 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.supplier.s_nationkey, collate: binary]",
+          "                            └─Projection 19531.25 mpp[tiflash]  test.lineitem.l_quantity, test.lineitem.l_extendedprice, test.lineitem.l_discount, test.supplier.s_nationkey, test.partsupp.ps_supplycost, test.orders.o_orderdate, test.lineitem.l_orderkey",
+          "                              └─HashJoin 19531.25 mpp[tiflash]  inner join, equal:[eq(test.lineitem.l_orderkey, test.orders.o_orderkey)]",
+          "                                ├─ExchangeReceiver(Build) 10000.00 mpp[tiflash]  ",
+          "                                │ └─ExchangeSender 10000.00 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.orders.o_orderkey, collate: binary]",
+          "                                │   └─TableFullScan 10000.00 mpp[tiflash] table:orders keep order:false, stats:pseudo",
+          "                                └─ExchangeReceiver(Probe) 15625.00 mpp[tiflash]  ",
+          "                                  └─ExchangeSender 15625.00 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.lineitem.l_orderkey, collate: binary]",
+          "                                    └─Projection 15625.00 mpp[tiflash]  test.lineitem.l_orderkey, test.lineitem.l_quantity, test.lineitem.l_extendedprice, test.lineitem.l_discount, test.supplier.s_nationkey, test.partsupp.ps_supplycost, test.lineitem.l_suppkey, test.lineitem.l_partkey",
+          "                                      └─HashJoin 15625.00 mpp[tiflash]  inner join, equal:[eq(test.lineitem.l_suppkey, test.partsupp.ps_suppkey) eq(test.lineitem.l_partkey, test.partsupp.ps_partkey)]",
+          "                                        ├─ExchangeReceiver(Build) 10000.00 mpp[tiflash]  ",
+          "                                        │ └─ExchangeSender 10000.00 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.partsupp.ps_suppkey, collate: binary], [name: test.partsupp.ps_partkey, collate: binary]",
+          "                                        │   └─TableFullScan 10000.00 mpp[tiflash] table:partsupp keep order:false, stats:pseudo",
+          "                                        └─ExchangeReceiver(Probe) 12500.00 mpp[tiflash]  ",
+          "                                          └─ExchangeSender 12500.00 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.lineitem.l_suppkey, collate: binary], [name: test.lineitem.l_partkey, collate: binary]",
+          "                                            └─Projection 12500.00 mpp[tiflash]  test.lineitem.l_orderkey, test.lineitem.l_partkey, test.lineitem.l_suppkey, test.lineitem.l_quantity, test.lineitem.l_extendedprice, test.lineitem.l_discount, test.supplier.s_nationkey, test.supplier.s_suppkey",
+          "                                              └─HashJoin 12500.00 mpp[tiflash]  inner join, equal:[eq(test.lineitem.l_suppkey, test.supplier.s_suppkey)]",
+          "                                                ├─ExchangeReceiver(Build) 10000.00 mpp[tiflash]  ",
+          "                                                │ └─ExchangeSender 10000.00 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.lineitem.l_suppkey, collate: binary]",
+          "                                                │   └─Projection 10000.00 mpp[tiflash]  test.lineitem.l_orderkey, test.lineitem.l_partkey, test.lineitem.l_suppkey, test.lineitem.l_quantity, test.lineitem.l_extendedprice, test.lineitem.l_discount",
+          "                                                │     └─HashJoin 10000.00 mpp[tiflash]  inner join, equal:[eq(test.part.p_partkey, test.lineitem.l_partkey)]",
+          "                                                │       ├─ExchangeReceiver(Build) 8000.00 mpp[tiflash]  ",
+          "                                                │       │ └─ExchangeSender 8000.00 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.part.p_partkey, collate: binary]",
+          "                                                │       │   └─Selection 8000.00 mpp[tiflash]  like(test.part.p_name, \"%dim%\", 92)",
+          "                                                │       │     └─TableFullScan 10000.00 mpp[tiflash] table:part pushed down filter:empty, keep order:false, stats:pseudo",
+          "                                                │       └─ExchangeReceiver(Probe) 10000.00 mpp[tiflash]  ",
+          "                                                │         └─ExchangeSender 10000.00 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.lineitem.l_partkey, collate: binary]",
+          "                                                │           └─TableFullScan 10000.00 mpp[tiflash] table:lineitem keep order:false, stats:pseudo",
+          "                                                └─ExchangeReceiver(Probe) 10000.00 mpp[tiflash]  ",
+          "                                                  └─ExchangeSender 10000.00 mpp[tiflash]  ExchangeType: HashPartition, Compression: FAST, Hash Cols: [name: test.supplier.s_suppkey, collate: binary]",
+          "                                                    └─TableFullScan 10000.00 mpp[tiflash] table:supplier keep order:false, stats:pseudo"
+        ]
+      }
+    ]
+  },
   {
     "Name": "TestQ13",
     "Cases": [

Diff for: pkg/planner/core/casetest/tpch/tpch_test.go

@@ -178,39 +178,26 @@ CREATE TABLE orders (
 	testkit.SetTiFlashReplica(t, dom, "test", "part")
 	testkit.SetTiFlashReplica(t, dom, "test", "partsupp")
 	testkit.SetTiFlashReplica(t, dom, "test", "supplier")
-	tk.MustQuery(`explain select
-	nation,
-	o_year,
-	sum(amount) as sum_profit
-from
-	(
-		select
-			n_name as nation,
-			extract(year from o_orderdate) as o_year,
-			l_extendedprice * (1 - l_discount) - ps_supplycost * l_quantity as amount
-		from
-			part,
-			supplier,
-			lineitem,
-			partsupp,
-			orders,
-			nation
-		where
-			s_suppkey = l_suppkey
-			and ps_suppkey = l_suppkey
-			and ps_partkey = l_partkey
-			and p_partkey = l_partkey
-			and o_orderkey = l_orderkey
-			and s_nationkey = n_nationkey
-			and p_name like '%dim%'
-	) as profit
-group by
-	nation,
-	o_year
-order by
-	nation,
-	o_year desc;`).Check(testkit.Rows())
+	integrationSuiteData := GetTPCHSuiteData()
+	var (
+		input  []string
+		output []struct {
+			SQL    string
+			Result []string
+		}
+	)
+	integrationSuiteData.LoadTestCases(t, &input, &output)
+	for i := 0; i < len(input); i++ {
+		testdata.OnRecord(func() {
+			output[i].SQL = input[i]
+		})
+		testdata.OnRecord(func() {
+			output[i].Result = testdata.ConvertRowsToStrings(tk.MustQuery(input[i]).Rows())
+		})
+		tk.MustQuery(input[i]).Check(testkit.Rows(output[i].Result...))
+	}
 }
+
 func TestQ13(t *testing.T) {
 	store, dom := testkit.CreateMockStoreAndDomain(t)
 	tk := testkit.NewTestKit(t, store)

Diff for: pkg/planner/core/initialize.go

@@ -218,7 +218,7 @@ func (base basePhysicalAgg) Init(ctx base.PlanContext, stats *property.StatsInfo
 }
 
 func (base basePhysicalAgg) initForHash(ctx base.PlanContext, stats *property.StatsInfo, offset int, props ...*property.PhysicalProperty) *PhysicalHashAgg {
-	p := &PhysicalHashAgg{base, ""}
+	p := &PhysicalHashAgg{base, "", nil}
 	p.BasePhysicalPlan = physicalop.NewBasePhysicalPlan(ctx, plancodec.TypeHashAgg, p, offset)
 	p.SetChildrenReqProps(props)
 	p.SetStats(stats)

Diff for: pkg/planner/core/physical_plans.go

@@ -1503,6 +1503,12 @@ type PhysicalHashJoin struct {
 
 	// for runtime filter
 	runtimeFilterList []*RuntimeFilter `plan-cache-clone:"must-nil"` // plan with runtime filter is not cached
+
+	// backupHashProp is used to store the hash partitioning information before `NeedMPPExchangeByEquivalence`.
+	// because `NeedMPPExchangeByEquivalence` will change the hash partitioning information.
+	// Having different numbers of MPP partitions in the child nodes of a join can lead to invalid tasks,
+	// which ultimately result in incorrect plan.
+	backupHashProp *property.PhysicalProperty `plan-cache-clone:"shallow"`
 }
 
 // CanUseHashJoinV2 returns true if current join is supported by hash join v2
@@ -2221,6 +2227,12 @@ func (p *basePhysicalAgg) MemoryUsage() (sum int64) {
 type PhysicalHashAgg struct {
 	basePhysicalAgg
 	tiflashPreAggMode string
+
+	// backupHashProp is used to store the hash partitioning information before `NeedMPPExchangeByEquivalence`.
+	// because `NeedMPPExchangeByEquivalence` will change the hash partitioning information.
+	// Having different numbers of MPP partitions in the child nodes of a join can lead to invalid tasks,
+	// which ultimately result in incorrect plan.
+	backupHashProp *property.PhysicalProperty `plan-cache-clone:"shallow"`
 }
 
 func (p *PhysicalHashAgg) getPointer() *basePhysicalAgg {

Diff for: pkg/planner/core/task.go

@@ -431,17 +431,40 @@ func (p *PhysicalHashJoin) convertPartitionKeysIfNeed(lTask, rTask *MppTask) (_,
 	return lTask, rTask
 }
 
+func enforceExchangerIfValid(task *MppTask, prop *property.PhysicalProperty, expectedCols int) *MppTask {
+	if prop != nil && len(prop.MPPPartitionCols) == expectedCols {
+		return task.enforceExchangerImpl(prop)
+	}
+	return nil
+}
+
 func (p *PhysicalHashJoin) attach2TaskForMpp(tasks ...base.Task) base.Task {
 	rTask, rok := tasks[1].(*MppTask)
 	lTask, lok := tasks[0].(*MppTask)
 	if !lok || !rok {
 		return base.InvalidTask
 	}
 	if p.mppShuffleJoin {
-		// protection check is case of some bugs
-		if len(lTask.hashCols) != len(rTask.hashCols) || len(lTask.hashCols) == 0 {
+		if len(lTask.hashCols) == 0 || len(rTask.hashCols) == 0 {
+			// if the hash columns are empty, this is very likely a bug.
 			return base.InvalidTask
 		}
+		if len(lTask.hashCols) != len(rTask.hashCols) {
+			// if the hash columns are not the same, The most likely scenario is that
+			// they have undergone exchange optimization, removing some hash columns.
+			// In this case, we need to restore them.
+			var backupHashProp *property.PhysicalProperty
+			if len(lTask.hashCols) < len(rTask.hashCols) {
+				backupHashProp = lTask.getBackupHashProp()
+				lTask = enforceExchangerIfValid(lTask, backupHashProp, len(rTask.hashCols))
+			} else {
+				backupHashProp = rTask.getBackupHashProp()
+				rTask = enforceExchangerIfValid(rTask, backupHashProp, len(lTask.hashCols))
+			}
+			if lTask == nil || rTask == nil {
+				return base.InvalidTask
+			}
+		}
 		lTask, rTask = p.convertPartitionKeysIfNeed(lTask, rTask)
 	}
 	p.SetChildren(lTask.Plan(), rTask.Plan())
@@ -2680,6 +2703,19 @@ func tryExpandVirtualColumn(p base.PhysicalPlan) {
 	}
 }
 
+func setBackupHashProp(prop *property.PhysicalProperty, p base.PhysicalPlan) {
+	switch pp := p.(type) {
+	case *PhysicalHashJoin:
+		pp.backupHashProp = prop
+	case *PhysicalHashAgg:
+		pp.backupHashProp = prop
+	default:
+		if child := p.Children(); len(child) == 1 {
+			setBackupHashProp(prop, child[0])
+		}
+	}
+}
+
 func (t *MppTask) needEnforceExchanger(prop *property.PhysicalProperty, fd *funcdep.FDSet) bool {
 	switch prop.MPPPartitionTp {
 	case property.AnyType:
@@ -2694,7 +2730,11 @@ func (t *MppTask) needEnforceExchanger(prop *property.PhysicalProperty, fd *func
 		}
 		// for example, if already partitioned by hash(B,C), then same (A,B,C) must distribute on a same node.
 		if fd != nil && len(t.hashCols) != 0 {
-			return prop.NeedMPPExchangeByEquivalence(t.hashCols, fd)
+			if prop.NeedMPPExchangeByEquivalence(t.hashCols, fd) {
+				return true
+			}
+			setBackupHashProp(prop, t.p)
+			return false
 		}
 		if len(prop.MPPPartitionCols) != len(t.hashCols) {
 			return true
@@ -2708,6 +2748,26 @@ func (t *MppTask) needEnforceExchanger(prop *property.PhysicalProperty, fd *func
 	}
 }
 
+func (t *MppTask) getBackupHashProp() *property.PhysicalProperty {
+	return getBackupHashProp(t.p)
+}
+
+func getBackupHashProp(p base.PhysicalPlan) *property.PhysicalProperty {
+	if p == nil {
+		return nil
+	}
+	switch pp := p.(type) {
+	case *PhysicalHashJoin:
+		return pp.backupHashProp
+	case *PhysicalHashAgg:
+		return pp.backupHashProp
+	}
+	for _, child := range p.Children() {
+		return getBackupHashProp(child)
+	}
+	return nil
+}
+
 func (t *MppTask) enforceExchanger(prop *property.PhysicalProperty, fd *funcdep.FDSet) *MppTask {
 	if !t.needEnforceExchanger(prop, fd) {
 		return t