Merge pull request #14218 from taosdata/feat/groupkey_function

enh(query): add _group_key function for partion by

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,

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

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
 

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;