From dbd71d247c627520cefc3bf1d914ec90195eb0d6 Mon Sep 17 00:00:00 2001
From: Xiaoyu Wang <xiaoyuwang@taosdata.com>
Date: Sat, 23 Jul 2022 20:46:39 +0800
Subject: [PATCH] feat: super table order by primary key optimization

---
 source/libs/planner/src/planSpliter.c      | 40 +++++++++++++++++++---
 source/libs/planner/test/planBasicTest.cpp |  7 ++--
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/source/libs/planner/src/planSpliter.c b/source/libs/planner/src/planSpliter.c
index 10604fac19..879c04743d 100644
--- a/source/libs/planner/src/planSpliter.c
+++ b/source/libs/planner/src/planSpliter.c
@@ -934,18 +934,48 @@ static SNode* stbSplFindPrimaryKeyFromScan(SScanLogicNode* pScan) {
   return pCol;
 }
 
+static int32_t stbSplCreateMergeScanNode(SScanLogicNode* pScan, SLogicNode** pOutputMergeScan,
+                                         SNodeList** pOutputMergeKeys) {
+  SNodeList* pChildren = pScan->node.pChildren;
+  pScan->node.pChildren = NULL;
+
+  int32_t         code = TSDB_CODE_SUCCESS;
+  SScanLogicNode* pMergeScan = (SScanLogicNode*)nodesCloneNode((SNode*)pScan);
+  if (NULL == pMergeScan) {
+    code = TSDB_CODE_OUT_OF_MEMORY;
+  }
+
+  SNodeList* pMergeKeys = NULL;
+  if (TSDB_CODE_SUCCESS == code) {
+    pMergeScan->scanType = SCAN_TYPE_TABLE_MERGE;
+    pMergeScan->node.pChildren = pChildren;
+    splSetParent((SLogicNode*)pMergeScan);
+    code = stbSplCreateMergeKeysByPrimaryKey(stbSplFindPrimaryKeyFromScan(pMergeScan), &pMergeKeys);
+  }
+
+  if (TSDB_CODE_SUCCESS == code) {
+    *pOutputMergeScan = (SLogicNode*)pMergeScan;
+    *pOutputMergeKeys = pMergeKeys;
+  } else {
+    nodesDestroyNode((SNode*)pMergeScan);
+    nodesDestroyList(pMergeKeys);
+  }
+
+  return code;
+}
+
 static int32_t stbSplSplitMergeScanNode(SSplitContext* pCxt, SLogicSubplan* pSubplan, SScanLogicNode* pScan,
                                         bool groupSort) {
-  SNodeList* pMergeKeys = NULL;
-  int32_t    code = stbSplCreateMergeKeysByPrimaryKey(stbSplFindPrimaryKeyFromScan(pScan), &pMergeKeys);
+  SLogicNode* pMergeScan = NULL;
+  SNodeList*  pMergeKeys = NULL;
+  int32_t     code = stbSplCreateMergeScanNode(pScan, &pMergeScan, &pMergeKeys);
   if (TSDB_CODE_SUCCESS == code) {
-    code = stbSplCreateMergeNode(pCxt, pSubplan, (SLogicNode*)pScan, pMergeKeys, (SLogicNode*)pScan, groupSort);
+    code = stbSplCreateMergeNode(pCxt, pSubplan, (SLogicNode*)pScan, pMergeKeys, pMergeScan, groupSort);
   }
   if (TSDB_CODE_SUCCESS == code) {
     code = nodesListMakeStrictAppend(&pSubplan->pChildren,
-                                     (SNode*)splCreateScanSubplan(pCxt, (SLogicNode*)pScan, SPLIT_FLAG_STABLE_SPLIT));
+                                     (SNode*)splCreateScanSubplan(pCxt, pMergeScan, SPLIT_FLAG_STABLE_SPLIT));
   }
-  pScan->scanType = SCAN_TYPE_TABLE_MERGE;
   ++(pCxt->groupId);
   return code;
 }
diff --git a/source/libs/planner/test/planBasicTest.cpp b/source/libs/planner/test/planBasicTest.cpp
index 8f9cd94c19..9cfae68d34 100644
--- a/source/libs/planner/test/planBasicTest.cpp
+++ b/source/libs/planner/test/planBasicTest.cpp
@@ -24,9 +24,10 @@ TEST_F(PlanBasicTest, selectClause) {
   useDb("root", "test");
 
   run("SELECT * FROM t1");
-  run("SELECT 1 FROM t1");
-  run("SELECT * FROM st1");
-  run("SELECT 1 FROM st1");
+
+  run("SELECT MAX(c1) c2, c2 FROM t1");
+
+  run("SELECT MAX(c1) c2, c2 FROM st1");
 }
 
 TEST_F(PlanBasicTest, whereClause) {