planner: open fully order prop push down for partition table's table scan and single index scan

executor/partition_table_test.go

@@ -630,9 +630,12 @@ func TestOrderByAndLimit(t *testing.T) {
 		require.True(t, tk.HasPlan(queryRangePartition, "TableReader"))
 		require.True(t, tk.HasPlan(queryHashPartition, "TableReader"))
 		require.True(t, tk.HasPlan(queryListPartition, "TableReader"))
-		require.True(t, tk.HasPlan(queryRangePartition, "Limit")) // check if order property is not pushed
+		require.True(t, tk.HasPlan(queryRangePartition, "Limit")) // check if order property is pushed
+		require.False(t, tk.HasPlan(queryRangePartition, "TopN")) // and is fully pushed
 		require.True(t, tk.HasPlan(queryHashPartition, "Limit"))
+		require.False(t, tk.HasPlan(queryHashPartition, "TopN"))
 		require.True(t, tk.HasPlan(queryListPartition, "Limit"))
+		require.False(t, tk.HasPlan(queryListPartition, "TopN"))
 		regularResult = tk.MustQuery(queryRegular).Rows()
 		tk.MustQuery(queryRangePartition).Check(regularResult)
 		tk.MustQuery(queryHashPartition).Check(regularResult)

planner/core/casetest/integration_test.go

@@ -948,6 +948,8 @@ func TestKeepOrderHint(t *testing.T) {
 	tk.MustExec("create table t(a int, b int, primary key(a));")
 	tk.MustExec("create table t1(a int, b int, index idx_a(a));")
 	tk.MustExec("create table th (a int, key(a)) partition by hash(a) partitions 4;")
+	tk.MustExec("create table thp (a int, primary key(a)) partition by hash(a) partitions 4;")
+	tk.MustExec("create table thh (a int, b int, key(a)) partition by hash(a) partitions 4;")
 	tk.MustExec("create definer='root'@'localhost' view v as select * from t1 where a<10 order by a limit 1;")
 	tk.MustExec("create definer='root'@'localhost' view v1 as select * from t where a<10 order by a limit 1;")
 
@@ -958,13 +960,14 @@ func TestKeepOrderHint(t *testing.T) {
 	err = tk.ExecToErr("explain select /*+ order_index(t, primary) */ * from t where a<10 limit 1;")
 	require.EqualError(t, err, "[planner:1815]Internal : Can't find a proper physical plan for this query")
 
-	// The partition table can not keep order
 	tk.MustExec("analyze table th;")
+	tk.MustExec("analyze table thp;")
+	tk.MustExec("analyze table thh;")
 	err = tk.ExecToErr("select a from th where a<1 order by a limit 1;")
 	require.NoError(t, err)
 
-	err = tk.ExecToErr("select /*+ order_index(th, a) */ a from th where a<1 order by a limit 1;")
-	require.EqualError(t, err, "[planner:1815]Internal : Can't find a proper physical plan for this query")
+	err = tk.ExecToErr("select /*+ order_index(thh, a) */ * from thh where a<1 order by a limit 1;")
+	require.NoError(t, err)
 
 	var input []string
 	var output []struct {

planner/core/casetest/testdata/integration_partition_suite_out.json

@@ -1289,35 +1289,35 @@
         "SQL": "explain format='brief' select a from trange use index (ia) where a > 10 and c = 10 order by a limit 10",
         "Plan": [
           "Projection 3.33 root  test.trange.a",
-          "└─TopN 3.33 root  test.trange.a, offset:0, count:10",
-          "  └─IndexLookUp 3.33 root partition:all ",
-          "    ├─IndexRangeScan(Build) 3333.33 cop[tikv] table:trange, index:ia(a) range:(10,+inf], keep order:false, stats:pseudo",
-          "    └─TopN(Probe) 3.33 cop[tikv]  test.trange.a, offset:0, count:10",
-          "      └─Selection 3.33 cop[tikv]  eq(test.trange.c, 10)",
+          "└─Limit 3.33 root  offset:0, count:10",
+          "  └─Projection 3.33 root  test.trange.a, test.trange.c",
+          "    └─IndexLookUp 3.33 root partition:all ",
+          "      ├─IndexRangeScan(Build) 3333.33 cop[tikv] table:trange, index:ia(a) range:(10,+inf], keep order:true, stats:pseudo",
+          "      └─Selection(Probe) 3.33 cop[tikv]  eq(test.trange.c, 10)",
           "        └─TableRowIDScan 3333.33 cop[tikv] table:trange keep order:false, stats:pseudo"
         ]
       },
       {
         "SQL": "explain format='brief' select a from tlist use index (ia) where a > 10 and c = 10 order by a limit 10",
         "Plan": [
           "Projection 3.33 root  test.tlist.a",
-          "└─TopN 3.33 root  test.tlist.a, offset:0, count:10",
-          "  └─IndexLookUp 3.33 root partition:all ",
-          "    ├─IndexRangeScan(Build) 3333.33 cop[tikv] table:tlist, index:ia(a) range:(10,+inf], keep order:false, stats:pseudo",
-          "    └─TopN(Probe) 3.33 cop[tikv]  test.tlist.a, offset:0, count:10",
-          "      └─Selection 3.33 cop[tikv]  eq(test.tlist.c, 10)",
+          "└─Limit 3.33 root  offset:0, count:10",
+          "  └─Projection 3.33 root  test.tlist.a, test.tlist.c",
+          "    └─IndexLookUp 3.33 root partition:all ",
+          "      ├─IndexRangeScan(Build) 3333.33 cop[tikv] table:tlist, index:ia(a) range:(10,+inf], keep order:true, stats:pseudo",
+          "      └─Selection(Probe) 3.33 cop[tikv]  eq(test.tlist.c, 10)",
           "        └─TableRowIDScan 3333.33 cop[tikv] table:tlist keep order:false, stats:pseudo"
         ]
       },
       {
         "SQL": "explain format='brief' select a from thash use index (ia) where a > 10 and c = 10 order by a limit 10",
         "Plan": [
           "Projection 3.33 root  test.thash.a",
-          "└─TopN 3.33 root  test.thash.a, offset:0, count:10",
-          "  └─IndexLookUp 3.33 root partition:all ",
-          "    ├─IndexRangeScan(Build) 3333.33 cop[tikv] table:thash, index:ia(a) range:(10,+inf], keep order:false, stats:pseudo",
-          "    └─TopN(Probe) 3.33 cop[tikv]  test.thash.a, offset:0, count:10",
-          "      └─Selection 3.33 cop[tikv]  eq(test.thash.c, 10)",
+          "└─Limit 3.33 root  offset:0, count:10",
+          "  └─Projection 3.33 root  test.thash.a, test.thash.c",
+          "    └─IndexLookUp 3.33 root partition:all ",
+          "      ├─IndexRangeScan(Build) 3333.33 cop[tikv] table:thash, index:ia(a) range:(10,+inf], keep order:true, stats:pseudo",
+          "      └─Selection(Probe) 3.33 cop[tikv]  eq(test.thash.c, 10)",
           "        └─TableRowIDScan 3333.33 cop[tikv] table:thash keep order:false, stats:pseudo"
         ]
       },

planner/core/casetest/testdata/integration_suite_in.json

@@ -704,7 +704,14 @@
       "explain WITH CTE AS (select /*+ order_index(t1, idx_a) */ * from t1 where a<10 order by a limit 1) SELECT * FROM CTE WHERE CTE.a <18 union select * from cte where cte.b > 1;",
       "explain WITH CTE AS (select /*+ order_index(t, primary) */ * from t where a<10 order by a limit 1) SELECT * FROM CTE WHERE CTE.a <18 union select * from cte where cte.b > 1;",
       "explain WITH CTE AS (select /*+ no_order_index(t1, idx_a) */ * from t1 where a<10 order by a limit 1) SELECT * FROM CTE WHERE CTE.a <18 union select * from cte where cte.b > 1;",
-      "explain WITH CTE AS (select /*+ no_order_index(t, primary) */ * from t where a<10 order by a limit 1) SELECT * FROM CTE WHERE CTE.a <18 union select * from cte where cte.b > 1;"
+      "explain WITH CTE AS (select /*+ no_order_index(t, primary) */ * from t where a<10 order by a limit 1) SELECT * FROM CTE WHERE CTE.a <18 union select * from cte where cte.b > 1;",
+
+      // Use hint for partition table
+      "explain select /*+ order_index(th, a) */ a from th where a<1 order by a limit 1",
+      "explain select /*+ no_order_index(th, a) */ a from th where a<1 order by a limit 1",
+      "explain select /*+ order_index(thp, primary) */ a from thp where a<1 order by a limit 1",
+      "explain select /*+ no_order_index(thp, primary) */ a from thp where a<1 order by a limit 1",
+      "explain select /*+ order_index(thh, a) */ * from thh where a<1 order by a limit 1"
     ]
   },
   {

planner/core/casetest/testdata/integration_suite_out.json

@@ -4772,6 +4772,58 @@
           "        └─TableRangeScan_25 3333.33 cop[tikv] table:t range:[-inf,10), keep order:false, stats:pseudo"
         ],
         "Warn": null
+      },
+      {
+        "SQL": "explain select /*+ order_index(th, a) */ a from th where a<1 order by a limit 1",
+        "Plan": [
+          "Limit_11 1.00 root  offset:0, count:1",
+          "└─IndexReader_15 1.00 root partition:all index:Limit_14",
+          "  └─Limit_14 1.00 cop[tikv]  offset:0, count:1",
+          "    └─IndexRangeScan_13 1.00 cop[tikv] table:th, index:a(a) range:[-inf,1), keep order:true, stats:pseudo"
+        ],
+        "Warn": null
+      },
+      {
+        "SQL": "explain select /*+ no_order_index(th, a) */ a from th where a<1 order by a limit 1",
+        "Plan": [
+          "TopN_8 1.00 root  test.th.a, offset:0, count:1",
+          "└─IndexReader_15 1.00 root partition:all index:TopN_14",
+          "  └─TopN_14 1.00 cop[tikv]  test.th.a, offset:0, count:1",
+          "    └─IndexRangeScan_13 3323.33 cop[tikv] table:th, index:a(a) range:[-inf,1), keep order:false, stats:pseudo"
+        ],
+        "Warn": null
+      },
+      {
+        "SQL": "explain select /*+ order_index(thp, primary) */ a from thp where a<1 order by a limit 1",
+        "Plan": [
+          "Limit_11 1.00 root  offset:0, count:1",
+          "└─TableReader_15 1.00 root partition:all data:Limit_14",
+          "  └─Limit_14 1.00 cop[tikv]  offset:0, count:1",
+          "    └─TableRangeScan_13 3333.33 cop[tikv] table:thp range:[-inf,1), keep order:true, stats:pseudo"
+        ],
+        "Warn": null
+      },
+      {
+        "SQL": "explain select /*+ no_order_index(thp, primary) */ a from thp where a<1 order by a limit 1",
+        "Plan": [
+          "TopN_8 1.00 root  test.thp.a, offset:0, count:1",
+          "└─TableReader_15 1.00 root partition:all data:TopN_14",
+          "  └─TopN_14 1.00 cop[tikv]  test.thp.a, offset:0, count:1",
+          "    └─TableRangeScan_13 3333.33 cop[tikv] table:thp range:[-inf,1), keep order:false, stats:pseudo"
+        ],
+        "Warn": null
+      },
+      {
+        "SQL": "explain select /*+ order_index(thh, a) */ * from thh where a<1 order by a limit 1",
+        "Plan": [
+          "Limit_12 1.00 root  offset:0, count:1",
+          "└─Projection_17 1.00 root  test.thh.a, test.thh.b",
+          "  └─IndexLookUp_16 1.00 root partition:all ",
+          "    ├─Limit_15(Build) 1.00 cop[tikv]  offset:0, count:1",
+          "    │ └─IndexRangeScan_13 1.00 cop[tikv] table:thh, index:a(a) range:[-inf,1), keep order:true, stats:pseudo",
+          "    └─TableRowIDScan_14(Probe) 1.00 cop[tikv] table:thh keep order:false, stats:pseudo"
+        ],
+        "Warn": null
       }
     ]
   },

planner/core/casetest/testdata/ordered_result_mode_suite_out.json

@@ -440,30 +440,26 @@
     "Cases": [
       {
         "Plan": [
-          "Sort_6 2.00 root  test.thash.a",
-          "└─TableReader_9 2.00 root partition:p0,p1 data:TableRangeScan_8",
-          "  └─TableRangeScan_8 2.00 cop[tikv] table:thash range:[1,1], [200,200], keep order:false, stats:pseudo"
+          "TableReader_11 2.00 root partition:p0,p1 data:TableRangeScan_10",
+          "└─TableRangeScan_10 2.00 cop[tikv] table:thash range:[1,1], [200,200], keep order:true, stats:pseudo"
         ]
       },
       {
         "Plan": [
-          "Sort_6 100.00 root  test.thash.a",
-          "└─TableReader_9 100.00 root partition:all data:TableRangeScan_8",
-          "  └─TableRangeScan_8 100.00 cop[tikv] table:thash range:[50,150], keep order:false, stats:pseudo"
+          "TableReader_11 100.00 root partition:all data:TableRangeScan_10",
+          "└─TableRangeScan_10 100.00 cop[tikv] table:thash range:[50,150], keep order:true, stats:pseudo"
         ]
       },
       {
         "Plan": [
-          "Sort_6 2.00 root  test.trange.a",
-          "└─TableReader_9 2.00 root partition:p0,p2 data:TableRangeScan_8",
-          "  └─TableRangeScan_8 2.00 cop[tikv] table:trange range:[1,1], [200,200], keep order:false, stats:pseudo"
+          "TableReader_11 2.00 root partition:p0,p2 data:TableRangeScan_10",
+          "└─TableRangeScan_10 2.00 cop[tikv] table:trange range:[1,1], [200,200], keep order:true, stats:pseudo"
         ]
       },
       {
         "Plan": [
-          "Sort_6 100.00 root  test.trange.a",
-          "└─TableReader_9 100.00 root partition:p0,p1 data:TableRangeScan_8",
-          "  └─TableRangeScan_8 100.00 cop[tikv] table:trange range:[50,150], keep order:false, stats:pseudo"
+          "TableReader_11 100.00 root partition:p0,p1 data:TableRangeScan_10",
+          "└─TableRangeScan_10 100.00 cop[tikv] table:trange range:[50,150], keep order:true, stats:pseudo"
         ]
       }
     ]

planner/core/find_best_task.go

@@ -1520,10 +1520,26 @@ func (ds *DataSource) convertToIndexScan(prop *property.PhysicalProperty,
 	}
 	if candidate.isMatchProp {
 		cop.keepOrder = true
-		// IndexScan on partition table can't keep order.
 		if ds.tableInfo.GetPartitionInfo() != nil {
-			return invalidTask, nil
+			// Add sort items for index scan for merge-sort operation between partitions.
+			byItems := make([]*util.ByItems, 0, len(prop.SortItems))
+			for _, si := range prop.SortItems {
+				byItems = append(byItems, &util.ByItems{
+					Expr: si.Col,
+					Desc: si.Desc,
+				})
+			}
+			cop.indexPlan.(*PhysicalIndexScan).ByItems = byItems
+			if !is.Index.Global && cop.tablePlan != nil {
+				is.Columns = append(is.Columns, model.NewExtraPhysTblIDColInfo())
+				is.schema.Append(&expression.Column{
+					RetType:  types.NewFieldType(mysql.TypeLonglong),
+					UniqueID: ds.ctx.GetSessionVars().AllocPlanColumnID(),
+					ID:       model.ExtraPhysTblID,
+				})
+			}
 		}
+
 		if cop.tablePlan != nil && !ds.tableInfo.IsCommonHandle {
 			col, isNew := cop.tablePlan.(*PhysicalTableScan).appendExtraHandleCol(ds)
 			cop.extraHandleCol = col
@@ -2106,9 +2122,20 @@ func (ds *DataSource) convertToTableScan(prop *property.PhysicalProperty, candid
 	task = copTask
 	if candidate.isMatchProp {
 		copTask.keepOrder = true
-		// TableScan on partition table can't keep order.
 		if ds.tableInfo.GetPartitionInfo() != nil {
-			return invalidTask, nil
+			// TableScan on partition table on TiFlash can't keep order.
+			if ts.StoreType == kv.TiFlash {
+				return invalidTask, nil
+			}
+			// Add sort items for table scan for merge-sort operation between partitions.
+			byItems := make([]*util.ByItems, 0, len(prop.SortItems))
+			for _, si := range prop.SortItems {
+				byItems = append(byItems, &util.ByItems{
+					Expr: si.Col,
+					Desc: si.Desc,
+				})
+			}
+			ts.ByItems = byItems
 		}
 	}
 	ts.addPushedDownSelection(copTask, ds.stats.ScaleByExpectCnt(prop.ExpectedCnt))