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

---
 source/libs/planner/src/planOptimizer.c       | 24 +++++++++----------
 source/libs/planner/src/planPhysiCreater.c    |  1 -
 source/libs/planner/src/planSpliter.c         | 14 ++++++++++-
 source/libs/planner/src/planUtil.c            |  3 ++-
 source/libs/planner/test/planOptimizeTest.cpp |  2 ++
 5 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/source/libs/planner/src/planOptimizer.c b/source/libs/planner/src/planOptimizer.c
index b006ac2b0a..f107b2343c 100644
--- a/source/libs/planner/src/planOptimizer.c
+++ b/source/libs/planner/src/planOptimizer.c
@@ -997,10 +997,7 @@ static int32_t sortPriKeyOptGetScanNodesImpl(SLogicNode* pNode, bool* pNotOptimi
 
   switch (nodeType(pNode)) {
     case QUERY_NODE_LOGIC_PLAN_SCAN:
-      if (TSDB_SUPER_TABLE != ((SScanLogicNode*)pNode)->tableType) {
-        return nodesListMakeAppend(pScanNodes, (SNode*)pNode);
-      }
-      break;
+      return nodesListMakeAppend(pScanNodes, (SNode*)pNode);
     case QUERY_NODE_LOGIC_PLAN_JOIN:
       code =
           sortPriKeyOptGetScanNodesImpl((SLogicNode*)nodesListGetNode(pNode->pChildren, 0), pNotOptimize, pScanNodes);
@@ -1040,13 +1037,16 @@ static EOrder sortPriKeyOptGetPriKeyOrder(SSortLogicNode* pSort) {
 static int32_t sortPriKeyOptApply(SOptimizeContext* pCxt, SLogicSubplan* pLogicSubplan, SSortLogicNode* pSort,
                                   SNodeList* pScanNodes) {
   EOrder order = sortPriKeyOptGetPriKeyOrder(pSort);
-  if (ORDER_DESC == order) {
-    SNode* pScanNode = NULL;
-    FOREACH(pScanNode, pScanNodes) {
-      SScanLogicNode* pScan = (SScanLogicNode*)pScanNode;
-      if (pScan->scanSeq[0] > 0) {
-        TSWAP(pScan->scanSeq[0], pScan->scanSeq[1]);
-      }
+  SNode* pScanNode = NULL;
+  FOREACH(pScanNode, pScanNodes) {
+    SScanLogicNode* pScan = (SScanLogicNode*)pScanNode;
+    if (ORDER_DESC == order && pScan->scanSeq[0] > 0) {
+      TSWAP(pScan->scanSeq[0], pScan->scanSeq[1]);
+    }
+    if (TSDB_SUPER_TABLE == pScan->tableType) {
+      pScan->scanType = SCAN_TYPE_TABLE_MERGE;
+      pScan->node.resultDataOrder = DATA_ORDER_LEVEL_GLOBAL;
+      pScan->node.requireDataOrder = DATA_ORDER_LEVEL_GLOBAL;
     }
   }
 
@@ -2191,7 +2191,7 @@ static bool tagScanMayBeOptimized(SLogicNode* pNode) {
       !planOptNodeListHasTbname(pAgg->pGroupKeys)) {
     return false;
   }
-  
+
   SNode* pGroupKey = NULL;
   FOREACH(pGroupKey, pAgg->pGroupKeys) {
     SNode* pGroup = NULL;
diff --git a/source/libs/planner/src/planPhysiCreater.c b/source/libs/planner/src/planPhysiCreater.c
index 0a1f8bbd0b..fd359fce45 100644
--- a/source/libs/planner/src/planPhysiCreater.c
+++ b/source/libs/planner/src/planPhysiCreater.c
@@ -415,7 +415,6 @@ static int32_t createScanPhysiNodeFinalize(SPhysiPlanContext* pCxt, SSubplan* pS
                                            SScanPhysiNode* pScanPhysiNode, SPhysiNode** pPhyNode) {
   int32_t code = createScanCols(pCxt, pScanPhysiNode, pScanLogicNode->pScanCols);
   if (TSDB_CODE_SUCCESS == code) {
-    // Data block describe also needs to be set without scanning column, such as SELECT COUNT(*) FROM t
     code = addDataBlockSlots(pCxt, pScanPhysiNode->pScanCols, pScanPhysiNode->node.pOutputDataBlockDesc);
   }
 
diff --git a/source/libs/planner/src/planSpliter.c b/source/libs/planner/src/planSpliter.c
index 8586234b7e..10604fac19 100644
--- a/source/libs/planner/src/planSpliter.c
+++ b/source/libs/planner/src/planSpliter.c
@@ -913,13 +913,25 @@ static int32_t stbSplSplitScanNodeWithPartTags(SSplitContext* pCxt, SStableSplit
 }
 
 static SNode* stbSplFindPrimaryKeyFromScan(SScanLogicNode* pScan) {
+  bool   find = false;
   SNode* pCol = NULL;
   FOREACH(pCol, pScan->pScanCols) {
     if (PRIMARYKEY_TIMESTAMP_COL_ID == ((SColumnNode*)pCol)->colId) {
+      find = true;
+      break;
+    }
+  }
+  if (!find) {
+    return NULL;
+  }
+  SNode* pTarget = NULL;
+  FOREACH(pTarget, pScan->node.pTargets) {
+    if (nodesEqualNode(pTarget, pCol)) {
       return pCol;
     }
   }
-  return NULL;
+  nodesListStrictAppend(pScan->node.pTargets, nodesCloneNode(pCol));
+  return pCol;
 }
 
 static int32_t stbSplSplitMergeScanNode(SSplitContext* pCxt, SLogicSubplan* pSubplan, SScanLogicNode* pScan,
diff --git a/source/libs/planner/src/planUtil.c b/source/libs/planner/src/planUtil.c
index bfa6079cb1..7aab8a7ca3 100644
--- a/source/libs/planner/src/planUtil.c
+++ b/source/libs/planner/src/planUtil.c
@@ -124,7 +124,8 @@ int32_t replaceLogicNode(SLogicSubplan* pSubplan, SLogicNode* pOld, SLogicNode*
 }
 
 static int32_t adjustScanDataRequirement(SScanLogicNode* pScan, EDataOrderLevel requirement) {
-  if (SCAN_TYPE_TABLE != pScan->scanType && SCAN_TYPE_TABLE_MERGE != pScan->scanType) {
+  if ((SCAN_TYPE_TABLE != pScan->scanType && SCAN_TYPE_TABLE_MERGE != pScan->scanType) ||
+      DATA_ORDER_LEVEL_GLOBAL == pScan->node.requireDataOrder) {
     return TSDB_CODE_SUCCESS;
   }
   // The lowest sort level of scan output data is DATA_ORDER_LEVEL_IN_BLOCK
diff --git a/source/libs/planner/test/planOptimizeTest.cpp b/source/libs/planner/test/planOptimizeTest.cpp
index 770ac94e5b..058705403b 100644
--- a/source/libs/planner/test/planOptimizeTest.cpp
+++ b/source/libs/planner/test/planOptimizeTest.cpp
@@ -53,6 +53,8 @@ TEST_F(PlanOptimizeTest, sortPrimaryKey) {
 
   run("SELECT c1 FROM t1 ORDER BY ts");
 
+  run("SELECT c1 FROM st1 ORDER BY ts");
+
   run("SELECT c1 FROM t1 ORDER BY ts DESC");
 
   run("SELECT COUNT(*) FROM t1 INTERVAL(10S) ORDER BY _WSTART DESC");
-- 
GitLab