enh: ins_tables optimize

include/libs/nodes/plannodes.h

@@ -315,7 +315,6 @@ typedef struct SScanPhysiNode {
 
 typedef SScanPhysiNode STagScanPhysiNode;
 typedef SScanPhysiNode SBlockDistScanPhysiNode;
-typedef SScanPhysiNode STableCountScanPhysiNode;
 
 typedef struct SLastRowScanPhysiNode {
   SScanPhysiNode scan;
@@ -324,6 +323,8 @@ typedef struct SLastRowScanPhysiNode {
   bool           ignoreNull;
 } SLastRowScanPhysiNode;
 
+typedef SLastRowScanPhysiNode STableCountScanPhysiNode;
+
 typedef struct SSystemTableScanPhysiNode {
   SScanPhysiNode scan;
   SEpSet         mgmtEpSet;

source/libs/nodes/src/nodesCodeFuncs.c

@@ -4632,9 +4632,9 @@ static int32_t specificNodeToJson(const void* pObj, SJson* pJson) {
       return logicPlanToJson(pObj, pJson);
     case QUERY_NODE_PHYSICAL_PLAN_TAG_SCAN:
     case QUERY_NODE_PHYSICAL_PLAN_BLOCK_DIST_SCAN:
-    case QUERY_NODE_PHYSICAL_PLAN_TABLE_COUNT_SCAN:
       return physiScanNodeToJson(pObj, pJson);
     case QUERY_NODE_PHYSICAL_PLAN_LAST_ROW_SCAN:
+    case QUERY_NODE_PHYSICAL_PLAN_TABLE_COUNT_SCAN:
       return physiLastRowScanNodeToJson(pObj, pJson);
     case QUERY_NODE_PHYSICAL_PLAN_TABLE_SCAN:
     case QUERY_NODE_PHYSICAL_PLAN_TABLE_MERGE_SCAN:

source/libs/planner/src/planOptimizer.c

@@ -2461,17 +2461,51 @@ typedef struct STbCntScanOptInfo {
   SName           table;
 } STbCntScanOptInfo;
 
-static bool tbCntScanOptIsEligibleAgg(SAggLogicNode* pAgg) {
-  if (1 != LIST_LENGTH(pAgg->pAggFuncs) || NULL != pAgg->pGroupKeys) {
-    return false;
+static bool tbCntScanOptIsEligibleGroupKeys(SNodeList* pGroupKeys) {
+  if (NULL == pGroupKeys) {
+    return true;
+  }
+
+  SNode* pGroupKey = NULL;
+  FOREACH(pGroupKey, pGroupKeys) {
+    SNode* pKey = nodesListGetNode(((SGroupingSetNode*)pGroupKey)->pParameterList, 0);
+    if (QUERY_NODE_COLUMN != nodeType(pKey)) {
+      return false;
+    }
+    SColumnNode* pCol = (SColumnNode*)pKey;
+    if (0 != strcmp(pCol->colName, "db_name") && 0 != strcmp(pCol->colName, "stable_name")) {
+      return false;
+    }
   }
-  SFunctionNode* pFunc = (SFunctionNode*)nodesListGetNode(pAgg->pAggFuncs, 0);
-  if (FUNCTION_TYPE_COUNT != pFunc->funcType) {
+
+  return true;
+}
+
+static bool tbCntScanOptNotNullableExpr(SNode* pNode) {
+  if (QUERY_NODE_COLUMN != nodeType(pNode)) {
     return false;
   }
+  const char* pColName = ((SColumnNode*)pNode)->colName;
+  return 0 == strcmp(pColName, "*") || 0 == strcmp(pColName, "db_name") || 0 == strcmp(pColName, "stable_name") ||
+         0 == strcmp(pColName, "table_name");
+}
+
+static bool tbCntScanOptIsEligibleAggFuncs(SNodeList* pAggFuncs) {
+  SNode* pNode = NULL;
+  FOREACH(pNode, pAggFuncs) {
+    SFunctionNode* pFunc = (SFunctionNode*)nodesListGetNode(pAggFuncs, 0);
+    if (FUNCTION_TYPE_COUNT != pFunc->funcType ||
+        !tbCntScanOptNotNullableExpr(nodesListGetNode(pFunc->pParameterList, 0))) {
+      return false;
+    }
+  }
   return true;
 }
 
+static bool tbCntScanOptIsEligibleAgg(SAggLogicNode* pAgg) {
+  return tbCntScanOptIsEligibleGroupKeys(pAgg->pGroupKeys) && tbCntScanOptIsEligibleAggFuncs(pAgg->pAggFuncs);
+}
+
 static bool tbCntScanOptGetColValFromCond(SOperatorNode* pOper, SColumnNode** pCol, SValueNode** pVal) {
   if (OP_TYPE_EQUAL != pOper->opType) {
     return false;
@@ -2578,17 +2612,25 @@ static SNode* tbCntScanOptCreateTableCountFunc() {
   return (SNode*)pFunc;
 }
 
-static int32_t tbCntScanOptRewriteScan(STbCntScanOptInfo* pInfo, SScanLogicNode* pScan) {
-  pScan->scanType = SCAN_TYPE_TABLE_COUNT;
-  strcpy(pScan->tableName.dbname, pInfo->table.dbname);
-  strcpy(pScan->tableName.tname, pInfo->table.tname);
-  NODES_DESTORY_LIST(pScan->node.pTargets);
-  NODES_DESTORY_NODE(pScan->node.pConditions);
-  NODES_DESTORY_LIST(pScan->pScanCols);
-  NODES_DESTORY_LIST(pScan->pScanPseudoCols);
-  int32_t code = nodesListMakeStrictAppend(&pScan->pScanPseudoCols, tbCntScanOptCreateTableCountFunc());
+static int32_t tbCntScanOptRewriteScan(STbCntScanOptInfo* pInfo) {
+  pInfo->pScan->scanType = SCAN_TYPE_TABLE_COUNT;
+  strcpy(pInfo->pScan->tableName.dbname, pInfo->table.dbname);
+  strcpy(pInfo->pScan->tableName.tname, pInfo->table.tname);
+  NODES_DESTORY_LIST(pInfo->pScan->node.pTargets);
+  NODES_DESTORY_NODE(pInfo->pScan->node.pConditions);
+  NODES_DESTORY_LIST(pInfo->pScan->pScanCols);
+  NODES_DESTORY_LIST(pInfo->pScan->pScanPseudoCols);
+  int32_t code = nodesListMakeStrictAppend(&pInfo->pScan->pScanPseudoCols, tbCntScanOptCreateTableCountFunc());
   if (TSDB_CODE_SUCCESS == code) {
-    code = createColumnByRewriteExpr(nodesListGetNode(pScan->pScanPseudoCols, 0), &pScan->node.pTargets);
+    code = createColumnByRewriteExpr(nodesListGetNode(pInfo->pScan->pScanPseudoCols, 0), &pInfo->pScan->node.pTargets);
+  }
+  SNode* pGroupKey = NULL;
+  FOREACH(pGroupKey, pInfo->pAgg->pGroupKeys) {
+    code = nodesListMakeStrictAppend(
+        &pInfo->pScan->pGroupTags, nodesCloneNode(nodesListGetNode(((SGroupingSetNode*)pGroupKey)->pParameterList, 0)));
+    if (TSDB_CODE_SUCCESS != code) {
+      break;
+    }
   }
   return code;
 }
@@ -2621,6 +2663,7 @@ static int32_t tbCntScanOptRewriteAgg(SAggLogicNode* pAgg) {
     NODES_DESTORY_LIST(pAgg->pAggFuncs);
     code = nodesListMakeStrictAppend(&pAgg->pAggFuncs, pSum);
   }
+  NODES_DESTORY_LIST(pAgg->pGroupKeys);
   return code;
 }
 
@@ -2630,7 +2673,7 @@ static int32_t tableCountScanOptimize(SOptimizeContext* pCxt, SLogicSubplan* pLo
     return TSDB_CODE_SUCCESS;
   }
 
-  int32_t code = tbCntScanOptRewriteScan(&info, info.pScan);
+  int32_t code = tbCntScanOptRewriteScan(&info);
   if (TSDB_CODE_SUCCESS == code) {
     code = tbCntScanOptRewriteAgg(info.pAgg);
   }

source/libs/planner/src/planPhysiCreater.c

@@ -526,6 +526,24 @@ static int32_t createLastRowScanPhysiNode(SPhysiPlanContext* pCxt, SSubplan* pSu
   pScan->ignoreNull = pScanLogicNode->igLastNull;
   vgroupInfoToNodeAddr(pScanLogicNode->pVgroupList->vgroups, &pSubplan->execNode);
 
+  return createScanPhysiNodeFinalize(pCxt, pSubplan, pScanLogicNode, (SScanPhysiNode*)pScan, pPhyNode);
+}
+
+static int32_t createTableCountScanPhysiNode(SPhysiPlanContext* pCxt, SSubplan* pSubplan,
+                                             SScanLogicNode* pScanLogicNode, SPhysiNode** pPhyNode) {
+  STableCountScanPhysiNode* pScan = (STableCountScanPhysiNode*)makePhysiNode(pCxt, (SLogicNode*)pScanLogicNode,
+                                                                             QUERY_NODE_PHYSICAL_PLAN_TABLE_COUNT_SCAN);
+  if (NULL == pScan) {
+    return TSDB_CODE_OUT_OF_MEMORY;
+  }
+
+  pScan->pGroupTags = nodesCloneList(pScanLogicNode->pGroupTags);
+  if (NULL != pScanLogicNode->pGroupTags && NULL == pScan->pGroupTags) {
+    nodesDestroyNode((SNode*)pScan);
+    return TSDB_CODE_OUT_OF_MEMORY;
+  }
+
+  pScan->groupSort = pScanLogicNode->groupSort;
   vgroupInfoToNodeAddr(pScanLogicNode->pVgroupList->vgroups, &pSubplan->execNode);
 
   return createScanPhysiNodeFinalize(pCxt, pSubplan, pScanLogicNode, (SScanPhysiNode*)pScan, pPhyNode);
@@ -624,8 +642,9 @@ static int32_t createScanPhysiNode(SPhysiPlanContext* pCxt, SSubplan* pSubplan,
   switch (pScanLogicNode->scanType) {
     case SCAN_TYPE_TAG:
     case SCAN_TYPE_BLOCK_INFO:
-    case SCAN_TYPE_TABLE_COUNT:
       return createSimpleScanPhysiNode(pCxt, pSubplan, pScanLogicNode, pPhyNode);
+    case SCAN_TYPE_TABLE_COUNT:
+      return createTableCountScanPhysiNode(pCxt, pSubplan, pScanLogicNode, pPhyNode);
     case SCAN_TYPE_LAST_ROW:
       return createLastRowScanPhysiNode(pCxt, pSubplan, pScanLogicNode, pPhyNode);
     case SCAN_TYPE_TABLE: