feat: refactor the plan implementation of unique function

source/libs/planner/src/planOptimizer.c

@@ -1034,12 +1034,30 @@ static SNodeList* partTagsGetPartKeys(SLogicNode* pNode) {
   }
 }
 
+static SNodeList* partTagsGetFuncs(SLogicNode* pNode) {
+  if (QUERY_NODE_LOGIC_PLAN_PARTITION == nodeType(pNode)) {
+    return NULL;
+  } else {
+    return ((SAggLogicNode*)pNode)->pAggFuncs;
+  }
+}
+
+static bool partTagsOptAreSupportedFuncs(SNodeList* pFuncs) {
+  SNode* pFunc = NULL;
+  FOREACH(pFunc, pFuncs) {
+    if (fmIsIndefiniteRowsFunc(((SFunctionNode*)pFunc)->funcId) && !fmIsSelectFunc(((SFunctionNode*)pFunc)->funcId)) {
+      return false;
+    }
+  }
+  return true;
+}
+
 static bool partTagsOptMayBeOptimized(SLogicNode* pNode) {
   if (!partTagsIsOptimizableNode(pNode)) {
     return false;
   }
 
-  return !partTagsOptHasCol(partTagsGetPartKeys(pNode));
+  return !partTagsOptHasCol(partTagsGetPartKeys(pNode)) && partTagsOptAreSupportedFuncs(partTagsGetFuncs(pNode));
 }
 
 static EDealRes partTagsOptRebuildTbanmeImpl(SNode** pNode, void* pContext) {
@@ -1065,13 +1083,13 @@ static int32_t partTagsOptRebuildTbanme(SNodeList* pPartKeys) {
   return code;
 }
 
-static SNode* partTagsCreateGroupKeyFunc(SNode* pNode) {
+static SNode* partTagsCreateWrapperFunc(const char* pFuncName, SNode* pNode) {
   SFunctionNode* pFunc = (SFunctionNode*)nodesMakeNode(QUERY_NODE_FUNCTION);
   if (NULL == pFunc) {
     return NULL;
   }
 
-  strcpy(pFunc->functionName, "_group_key");
+  strcpy(pFunc->functionName, pFuncName);
   if (QUERY_NODE_COLUMN == nodeType(pNode)) {
     SColumnNode* pCol = (SColumnNode*)pNode;
     snprintf(pFunc->node.aliasName, sizeof(pFunc->node.aliasName), "%s.%s", pCol->tableAlias, pCol->colName);
@@ -1091,15 +1109,31 @@ static SNode* partTagsCreateGroupKeyFunc(SNode* pNode) {
   return (SNode*)pFunc;
 }
 
-static int32_t partTagsRewriteGroupTagsToGroupKeyFuncs(SNodeList* pGroupTags, SNodeList* pAggFuncs) {
-  SNode* pNode = NULL;
+static bool partTagsHasIndefRowsSelectFunc(SNodeList* pFuncs) {
+  SNode* pFunc = NULL;
+  FOREACH(pFunc, pFuncs) {
+    if (fmIsIndefiniteRowsFunc(((SFunctionNode*)pFunc)->funcId)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+static int32_t partTagsRewriteGroupTagsToFuncs(SNodeList* pGroupTags, SNodeList* pAggFuncs) {
+  bool    hasIndefRowsSelectFunc = partTagsHasIndefRowsSelectFunc(pAggFuncs);
+  int32_t code = TSDB_CODE_SUCCESS;
+  SNode*  pNode = NULL;
   FOREACH(pNode, pGroupTags) {
-    int32_t code = nodesListStrictAppend(pAggFuncs, partTagsCreateGroupKeyFunc(pNode));
+    if (hasIndefRowsSelectFunc) {
+      code = nodesListStrictAppend(pAggFuncs, partTagsCreateWrapperFunc("_select_value", pNode));
+    } else {
+      code = nodesListStrictAppend(pAggFuncs, partTagsCreateWrapperFunc("_group_key", pNode));
+    }
     if (TSDB_CODE_SUCCESS != code) {
-      return code;
+      break;
     }
   }
-  return TSDB_CODE_SUCCESS;
+  return code;
 }
 
 static int32_t partTagsOptimize(SOptimizeContext* pCxt, SLogicSubplan* pLogicSubplan) {
@@ -1128,7 +1162,7 @@ static int32_t partTagsOptimize(SOptimizeContext* pCxt, SLogicSubplan* pLogicSub
       }
     }
     NODES_DESTORY_LIST(pAgg->pGroupKeys);
-    code = partTagsRewriteGroupTagsToGroupKeyFuncs(pScan->pGroupTags, pAgg->pAggFuncs);
+    code = partTagsRewriteGroupTagsToFuncs(pScan->pGroupTags, pAgg->pAggFuncs);
   }
   if (TSDB_CODE_SUCCESS == code) {
     code = partTagsOptRebuildTbanme(pScan->pGroupTags);

source/libs/planner/test/planBasicTest.cpp

@@ -99,6 +99,16 @@ TEST_F(PlanBasicTest, lastRowFunc) {
   run("SELECT LAST_ROW(c1) FROM st1");
 }
 
+TEST_F(PlanBasicTest, sampleFunc) {
+  useDb("root", "test");
+
+  run("SELECT SAMPLE(c1, 10) FROM t1");
+
+  run("SELECT SAMPLE(c1, 10) FROM st1");
+
+  run("SELECT SAMPLE(c1, 10) FROM st1 PARTITION BY TBNAME");
+}
+
 TEST_F(PlanBasicTest, withoutFrom) {
   useDb("root", "test");
 

tests/system-test/0-others/udfTest.py

@@ -301,13 +301,13 @@ class TDTestCase:
         tdSql.checkRows(1)
         tdSql.query("select ceil(num1) , min(num1) from tb;")
         tdSql.checkRows(1)
-        tdSql.error("select udf1(num1) , first(num1) from tb;")
+        tdSql.query("select udf1(num1) , first(num1) from tb;")
         
-        tdSql.error("select abs(num1) , first(num1) from tb;")
+        tdSql.query("select abs(num1) , first(num1) from tb;")
         
-        tdSql.error("select udf1(num1) , last(num1) from tb;")
+        tdSql.query("select udf1(num1) , last(num1) from tb;")
         
-        tdSql.error("select round(num1) , last(num1) from tb;")
+        tdSql.query("select round(num1) , last(num1) from tb;")
         
         tdSql.query("select udf1(num1) , top(num1,1) from tb;")
         tdSql.checkRows(1)
@@ -327,9 +327,9 @@ class TDTestCase:
         tdSql.checkRows(1)
         tdSql.query("select floor(c1) , min(c1) from stb1;")
         tdSql.checkRows(1)
-        tdSql.error("select udf1(c1) , first(c1) from stb1;")
+        tdSql.query("select udf1(c1) , first(c1) from stb1;")
         
-        tdSql.error("select udf1(c1) , last(c1) from stb1;")
+        tdSql.query("select udf1(c1) , last(c1) from stb1;")
         
         tdSql.query("select udf1(c1) , top(c1 ,1) from stb1;")
         tdSql.checkRows(1)

tests/system-test/0-others/udf_create.py

@@ -303,13 +303,13 @@ class TDTestCase:
         tdSql.checkRows(1)
         tdSql.query("select ceil(num1) , min(num1) from tb;")
         tdSql.checkRows(1)
-        tdSql.error("select udf1(num1) , first(num1) from tb;")
+        tdSql.query("select udf1(num1) , first(num1) from tb;")
         
-        tdSql.error("select abs(num1) , first(num1) from tb;")
+        tdSql.query("select abs(num1) , first(num1) from tb;")
         
-        tdSql.error("select udf1(num1) , last(num1) from tb;")
+        tdSql.query("select udf1(num1) , last(num1) from tb;")
         
-        tdSql.error("select round(num1) , last(num1) from tb;")
+        tdSql.query("select round(num1) , last(num1) from tb;")
         
         tdSql.query("select udf1(num1) , top(num1,1) from tb;")
         tdSql.checkRows(1)
@@ -329,9 +329,9 @@ class TDTestCase:
         tdSql.checkRows(1)
         tdSql.query("select floor(c1) , min(c1) from stb1;")
         tdSql.checkRows(1)
-        tdSql.error("select udf1(c1) , first(c1) from stb1;")
+        tdSql.query("select udf1(c1) , first(c1) from stb1;")
         
-        tdSql.error("select udf1(c1) , last(c1) from stb1;")
+        tdSql.query("select udf1(c1) , last(c1) from stb1;")
         
         tdSql.query("select udf1(c1) , top(c1 ,1) from stb1;")
         tdSql.checkRows(1)

tests/system-test/0-others/udf_restart_taosd.py

@@ -300,13 +300,13 @@ class TDTestCase:
         tdSql.checkRows(1)
         tdSql.query("select ceil(num1) , min(num1) from tb;")
         tdSql.checkRows(1)
-        tdSql.error("select udf1(num1) , first(num1) from tb;")
+        tdSql.query("select udf1(num1) , first(num1) from tb;")
         
-        tdSql.error("select abs(num1) , first(num1) from tb;")
+        tdSql.query("select abs(num1) , first(num1) from tb;")
         
-        tdSql.error("select udf1(num1) , last(num1) from tb;")
+        tdSql.query("select udf1(num1) , last(num1) from tb;")
         
-        tdSql.error("select round(num1) , last(num1) from tb;")
+        tdSql.query("select round(num1) , last(num1) from tb;")
         
         tdSql.query("select udf1(num1) , top(num1,1) from tb;")
         tdSql.checkRows(1)
@@ -326,9 +326,9 @@ class TDTestCase:
         tdSql.checkRows(1)
         tdSql.query("select floor(c1) , min(c1) from stb1;")
         tdSql.checkRows(1)
-        tdSql.error("select udf1(c1) , first(c1) from stb1;")
+        tdSql.query("select udf1(c1) , first(c1) from stb1;")
         
-        tdSql.error("select udf1(c1) , last(c1) from stb1;")
+        tdSql.query("select udf1(c1) , last(c1) from stb1;")
         
         tdSql.query("select udf1(c1) , top(c1 ,1) from stb1;")
         tdSql.checkRows(1)

taos-tools @ 28a49b44

-Subproject commit a875a057d1225d85c6323b9edaccc2b1a9641987
+Subproject commit 28a49b447f71c4f014ebbac858b7215b897d57fd