diff --git a/include/libs/function/functionMgt.h b/include/libs/function/functionMgt.h index 1ca51b3043..1ed78750d1 100644 --- a/include/libs/function/functionMgt.h +++ b/include/libs/function/functionMgt.h @@ -124,6 +124,7 @@ typedef enum EFunctionType { FUNCTION_TYPE_BLOCK_DIST, // block distribution aggregate function FUNCTION_TYPE_BLOCK_DIST_INFO, // block distribution pseudo column function FUNCTION_TYPE_TO_COLUMN, + FUNCTION_TYPE_GROUP_KEY, // distributed splitting functions FUNCTION_TYPE_APERCENTILE_PARTIAL = 4000, diff --git a/source/libs/function/inc/builtinsimpl.h b/source/libs/function/inc/builtinsimpl.h index 4d2f926c12..0820e884ce 100644 --- a/source/libs/function/inc/builtinsimpl.h +++ b/source/libs/function/inc/builtinsimpl.h @@ -202,6 +202,10 @@ bool blockDistSetup(SqlFunctionCtx *pCtx, SResultRowEntryInfo* pResultInfo); int32_t blockDistFunction(SqlFunctionCtx *pCtx); int32_t blockDistFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock); +bool getGroupKeyFuncEnv(SFunctionNode* pFunc, SFuncExecEnv* pEnv); +int32_t groupKeyFunction(SqlFunctionCtx* pCtx); +int32_t groupKeyFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock); + #ifdef __cplusplus } #endif diff --git a/source/libs/function/src/builtins.c b/source/libs/function/src/builtins.c index 6516a10795..6744ebe4b9 100644 --- a/source/libs/function/src/builtins.c +++ b/source/libs/function/src/builtins.c @@ -1515,6 +1515,16 @@ static bool getBlockDistFuncEnv(SFunctionNode* UNUSED_PARAM(pFunc), SFuncExecEnv return true; } +static int32_t translateGroupKey(SFunctionNode* pFunc, char* pErrBuf, int32_t len) { + if (1 != LIST_LENGTH(pFunc->pParameterList)) { + return TSDB_CODE_SUCCESS; + } + + SNode* pPara = nodesListGetNode(pFunc->pParameterList, 0); + pFunc->node.resType = ((SExprNode*)pPara)->resType; + return TSDB_CODE_SUCCESS; +} + // clang-format off const SBuiltinFuncDefinition funcMgtBuiltins[] = { { @@ -2499,7 +2509,17 @@ const SBuiltinFuncDefinition funcMgtBuiltins[] = { .type = FUNCTION_TYPE_BLOCK_DIST_INFO, .classification = FUNC_MGT_PSEUDO_COLUMN_FUNC | FUNC_MGT_SCAN_PC_FUNC, .translateFunc = translateBlockDistInfoFunc, - } + }, + { + .name = "_group_key", + .type = FUNCTION_TYPE_GROUP_KEY, + .classification = FUNC_MGT_AGG_FUNC, + .translateFunc = translateGroupKey, + .getEnvFunc = getGroupKeyFuncEnv, + .initFunc = functionSetup, + .processFunc = groupKeyFunction, + .finalizeFunc = groupKeyFinalize, + }, }; // clang-format on diff --git a/source/libs/function/src/builtinsimpl.c b/source/libs/function/src/builtinsimpl.c index 4ec6432fc8..61f3154811 100644 --- a/source/libs/function/src/builtinsimpl.c +++ b/source/libs/function/src/builtinsimpl.c @@ -262,6 +262,12 @@ typedef struct SRateInfo { int8_t hasResult; // flag to denote has value } SRateInfo; +typedef struct SGroupKeyInfo{ + bool hasResult; + char data[]; +} SGroupKeyInfo; + + #define SET_VAL(_info, numOfElem, res) \ do { \ if ((numOfElem) <= 0) { \ @@ -2402,6 +2408,12 @@ bool getSelectivityFuncEnv(SFunctionNode* pFunc, SFuncExecEnv* pEnv) { return true; } +bool getGroupKeyFuncEnv(SFunctionNode* pFunc, SFuncExecEnv* pEnv) { + SColumnNode* pNode = (SColumnNode*)nodesListGetNode(pFunc->pParameterList, 0); + pEnv->calcMemSize = sizeof(SGroupKeyInfo) + pNode->node.resType.bytes; + return true; +} + static FORCE_INLINE TSKEY getRowPTs(SColumnInfoData* pTsColInfo, int32_t rowIndex) { if (pTsColInfo == NULL) { return 0; @@ -5349,6 +5361,43 @@ int32_t irateFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock) { return pResInfo->numOfRes; } +int32_t groupKeyFunction(SqlFunctionCtx* pCtx) { + SResultRowEntryInfo* pResInfo = GET_RES_INFO(pCtx); + SGroupKeyInfo* pInfo = GET_ROWCELL_INTERBUF(pResInfo); + + SInputColumnInfoData* pInput = &pCtx->input; + SColumnInfoData* pInputCol = pInput->pData[0]; + + int32_t bytes = pInputCol->info.bytes; + + int32_t startIndex = pInput->startRowIndex; + if (colDataIsNull_s(pInputCol, startIndex)) { + pInfo->hasResult = false; + goto _group_key_over; + } + + pInfo->hasResult = true; + char* data = colDataGetData(pInputCol, startIndex); + memcpy(pInfo->data, data, bytes); + +_group_key_over: + + SET_VAL(pResInfo, 1, 1); + return TSDB_CODE_SUCCESS; +} + +int32_t groupKeyFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock) { + int32_t slotId = pCtx->pExpr->base.resSchema.slotId; + SColumnInfoData* pCol = taosArrayGet(pBlock->pDataBlock, slotId); + + SResultRowEntryInfo* pResInfo = GET_RES_INFO(pCtx); + + SGroupKeyInfo* pInfo = GET_ROWCELL_INTERBUF(pResInfo); + colDataAppend(pCol, pBlock->info.rows, pInfo->data, pInfo->hasResult ? false : true); + + return pResInfo->numOfRes; +} + int32_t interpFunction(SqlFunctionCtx* pCtx) { #if 0 int32_t fillType = (int32_t) pCtx->param[2].i64;