Merge pull request #13604 from taosdata/feat/split_enhance

enh(query): add spread function distributed splitting

include/libs/function/functionMgt.h

@@ -125,6 +125,8 @@ typedef enum EFunctionType {
   // distributed splitting functions
   FUNCTION_TYPE_APERCENTILE_PARTIAL,
   FUNCTION_TYPE_APERCENTILE_MERGE,
+  FUNCTION_TYPE_SPREAD_PARTIAL,
+  FUNCTION_TYPE_SPREAD_MERGE,
 
   // user defined funcion
   FUNCTION_TYPE_UDF = 10000

source/libs/function/inc/builtinsimpl.h

@@ -102,10 +102,13 @@ int32_t topFunction(SqlFunctionCtx *pCtx);
 int32_t bottomFunction(SqlFunctionCtx *pCtx);
 int32_t topBotFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock);
 
+int32_t getSpreadInfoSize();
 bool getSpreadFuncEnv(struct SFunctionNode* pFunc, SFuncExecEnv* pEnv);
 bool spreadFunctionSetup(SqlFunctionCtx *pCtx, SResultRowEntryInfo* pResultInfo);
 int32_t spreadFunction(SqlFunctionCtx* pCtx);
+int32_t spreadFunctionMerge(SqlFunctionCtx* pCtx);
 int32_t spreadFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock);
+int32_t spreadPartialFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock);
 
 bool getElapsedFuncEnv(struct SFunctionNode* pFunc, SFuncExecEnv* pEnv);
 bool elapsedFunctionSetup(SqlFunctionCtx *pCtx, SResultRowEntryInfo* pResultInfo);

source/libs/function/src/builtins.c

@@ -317,7 +317,6 @@ static int32_t translateApercentileMerge(SFunctionNode* pFunc, char* pErrBuf, in
   return translateApercentileImpl(pFunc, pErrBuf, len, false);
 }
 
-
 static int32_t translateTbnameColumn(SFunctionNode* pFunc, char* pErrBuf, int32_t len) {
   // pseudo column do not need to check parameters
   pFunc->node.resType =
@@ -378,6 +377,34 @@ static int32_t translateSpread(SFunctionNode* pFunc, char* pErrBuf, int32_t len)
   return TSDB_CODE_SUCCESS;
 }
 
+static int32_t translateSpreadImpl(SFunctionNode* pFunc, char* pErrBuf, int32_t len, bool isPartial) {
+  if (1 != LIST_LENGTH(pFunc->pParameterList)) {
+    return invaildFuncParaNumErrMsg(pErrBuf, len, pFunc->functionName);
+  }
+
+  uint8_t paraType = ((SExprNode*)nodesListGetNode(pFunc->pParameterList, 0))->resType.type;
+  if (isPartial) {
+    if (!IS_NUMERIC_TYPE(paraType) && TSDB_DATA_TYPE_TIMESTAMP != paraType) {
+      return invaildFuncParaTypeErrMsg(pErrBuf, len, pFunc->functionName);
+    }
+    pFunc->node.resType = (SDataType){.bytes = getSpreadInfoSize() + VARSTR_HEADER_SIZE, .type = TSDB_DATA_TYPE_BINARY};
+  } else {
+    if (TSDB_DATA_TYPE_BINARY != paraType) {
+      return invaildFuncParaTypeErrMsg(pErrBuf, len, pFunc->functionName);
+    }
+    pFunc->node.resType = (SDataType){.bytes = tDataTypes[TSDB_DATA_TYPE_DOUBLE].bytes, .type = TSDB_DATA_TYPE_DOUBLE};
+  }
+
+  return TSDB_CODE_SUCCESS;
+}
+
+static int32_t translateSpreadPartial(SFunctionNode* pFunc, char* pErrBuf, int32_t len) {
+  return translateSpreadImpl(pFunc, pErrBuf, len, true);
+}
+static int32_t translateSpreadMerge(SFunctionNode* pFunc, char* pErrBuf, int32_t len) {
+  return translateSpreadImpl(pFunc, pErrBuf, len, false);
+}
+
 static int32_t translateElapsed(SFunctionNode* pFunc, char* pErrBuf, int32_t len) {
   int32_t numOfParams = LIST_LENGTH(pFunc->pParameterList);
   if (1 != numOfParams && 2 != numOfParams) {
@@ -1277,6 +1304,30 @@ const SBuiltinFuncDefinition funcMgtBuiltins[] = {
     .getEnvFunc   = getSpreadFuncEnv,
     .initFunc     = spreadFunctionSetup,
     .processFunc  = spreadFunction,
+    .finalizeFunc = spreadFinalize,
+    .pPartialFunc = "_spread_partial",
+    .pMergeFunc   = "_spread_merge"
+  },
+  {
+    .name = "_spread_partial",
+    .type = FUNCTION_TYPE_SPREAD_PARTIAL,
+    .classification = FUNC_MGT_AGG_FUNC,
+    .translateFunc = translateSpreadPartial,
+    .dataRequiredFunc = statisDataRequired,
+    .getEnvFunc   = getSpreadFuncEnv,
+    .initFunc     = spreadFunctionSetup,
+    .processFunc  = spreadFunction,
+    .finalizeFunc = spreadPartialFinalize
+  },
+  {
+    .name = "_spread_merge",
+    .type = FUNCTION_TYPE_SPREAD_MERGE,
+    .classification = FUNC_MGT_AGG_FUNC,
+    .translateFunc = translateSpreadMerge,
+    .dataRequiredFunc = statisDataRequired,
+    .getEnvFunc   = getSpreadFuncEnv,
+    .initFunc     = spreadFunctionSetup,
+    .processFunc  = spreadFunctionMerge,
     .finalizeFunc = spreadFinalize
   },
   {

source/libs/function/src/builtinsimpl.c

@@ -2190,19 +2190,19 @@ int32_t apercentilePartialFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock) {
   int32_t bytesHist   = (int32_t)(sizeof(SAPercentileInfo) + sizeof(SHistogramInfo) + sizeof(SHistBin) * (MAX_HISTOGRAM_BIN + 1));
   int32_t bytesDigest = (int32_t)(sizeof(SAPercentileInfo) + TDIGEST_SIZE(COMPRESSION));
   int32_t resultBytes = TMAX(bytesHist, bytesDigest);
-  char *tmp = taosMemoryCalloc(resultBytes + VARSTR_HEADER_SIZE, sizeof(char));
+  char *res = taosMemoryCalloc(resultBytes + VARSTR_HEADER_SIZE, sizeof(char));
 
   if (pInfo->algo == APERCT_ALGO_TDIGEST) {
     if (pInfo->pTDigest->size > 0) {
-      memcpy(varDataVal(tmp), pInfo, resultBytes);
+      memcpy(varDataVal(res), pInfo, resultBytes);
-      varDataSetLen(tmp, resultBytes);
+      varDataSetLen(res, resultBytes);
     } else {
       return TSDB_CODE_SUCCESS;
     }
   } else {
     if (pInfo->pHisto->numOfElems > 0) {
-      memcpy(varDataVal(tmp), pInfo, resultBytes);
+      memcpy(varDataVal(res), pInfo, resultBytes);
-      varDataSetLen(tmp, resultBytes);
+      varDataSetLen(res, resultBytes);
     } else {
       return TSDB_CODE_SUCCESS;
     }
@@ -2211,9 +2211,9 @@ int32_t apercentilePartialFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock) {
   int32_t          slotId = pCtx->pExpr->base.resSchema.slotId;
   SColumnInfoData* pCol = taosArrayGet(pBlock->pDataBlock, slotId);
 
-  colDataAppend(pCol, pBlock->info.rows, tmp, false);
+  colDataAppend(pCol, pBlock->info.rows, res, false);
 
-  taosMemoryFree(tmp);
+  taosMemoryFree(res);
   return pResInfo->numOfRes;
 }
 
@@ -2868,6 +2868,10 @@ bool getSpreadFuncEnv(SFunctionNode* UNUSED_PARAM(pFunc), SFuncExecEnv* pEnv) {
   return true;
 }
 
+int32_t getSpreadInfoSize() {
+  return (int32_t)sizeof(SSpreadInfo);
+}
+
 bool spreadFunctionSetup(SqlFunctionCtx *pCtx, SResultRowEntryInfo* pResultInfo) {
   if (!functionSetup(pCtx, pResultInfo)) {
     return false;
@@ -2953,6 +2957,32 @@ _spread_over:
   return TSDB_CODE_SUCCESS;
 }
 
+int32_t spreadFunctionMerge(SqlFunctionCtx *pCtx) {
+  SInputColumnInfoData* pInput = &pCtx->input;
+  SColumnInfoData* pCol = pInput->pData[0];
+  ASSERT(pCol->info.type == TSDB_DATA_TYPE_BINARY);
+
+  SSpreadInfo* pInfo = GET_ROWCELL_INTERBUF(GET_RES_INFO(pCtx));
+  SSpreadInfo* pInputInfo;
+
+  int32_t start = pInput->startRowIndex;
+  char* data = colDataGetData(pCol, start);
+  pInputInfo = (SSpreadInfo *)varDataVal(data);
+
+  pInfo->hasResult = pInputInfo->hasResult;
+  if (pInputInfo->max > pInfo->max) {
+    pInfo->max = pInputInfo->max;
+  }
+
+  if (pInputInfo->min < pInfo->min) {
+    pInfo->min = pInputInfo->min;
+  }
+
+  SET_VAL(GET_RES_INFO(pCtx), 1, 1);
+
+  return TSDB_CODE_SUCCESS;
+}
+
 int32_t spreadFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock) {
   SSpreadInfo* pInfo = GET_ROWCELL_INTERBUF(GET_RES_INFO(pCtx));
   if (pInfo->hasResult == true) {
@@ -2961,6 +2991,24 @@ int32_t spreadFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock) {
   return functionFinalize(pCtx, pBlock);
 }
 
+int32_t spreadPartialFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock) {
+  SResultRowEntryInfo* pResInfo = GET_RES_INFO(pCtx);
+  SSpreadInfo* pInfo = GET_ROWCELL_INTERBUF(GET_RES_INFO(pCtx));
+  int32_t resultBytes = (int32_t)sizeof(SSpreadInfo);
+  char *res = taosMemoryCalloc(resultBytes + VARSTR_HEADER_SIZE, sizeof(char));
+
+  memcpy(varDataVal(res), pInfo, resultBytes);
+  varDataSetLen(res, resultBytes);
+
+  int32_t          slotId = pCtx->pExpr->base.resSchema.slotId;
+  SColumnInfoData* pCol = taosArrayGet(pBlock->pDataBlock, slotId);
+
+  colDataAppend(pCol, pBlock->info.rows, res, false);
+
+  taosMemoryFree(res);
+  return pResInfo->numOfRes;
+}
+
 bool getElapsedFuncEnv(SFunctionNode* UNUSED_PARAM(pFunc), SFuncExecEnv* pEnv) {
   pEnv->calcMemSize = sizeof(SElapsedInfo);
   return true;