From 4113472e89c69c1bd55036b03016346305409351 Mon Sep 17 00:00:00 2001
From: Ganlin Zhao <ganlinzhao@hotmail.com>
Date: Tue, 14 Jun 2022 15:00:28 +0800
Subject: [PATCH] add avg function translate partial/merge

---
 include/libs/function/functionMgt.h     |  2 +
 source/libs/function/inc/builtinsimpl.h |  1 +
 source/libs/function/src/builtins.c     | 53 +++++++++++++++++++++++++
 source/libs/function/src/builtinsimpl.c |  2 +
 4 files changed, 58 insertions(+)

diff --git a/include/libs/function/functionMgt.h b/include/libs/function/functionMgt.h
index b5641f2d07..f324d8fed1 100644
--- a/include/libs/function/functionMgt.h
+++ b/include/libs/function/functionMgt.h
@@ -142,6 +142,8 @@ typedef enum EFunctionType {
   FUNCTION_TYPE_FIRST_MERGE,
   FUNCTION_TYPE_LAST_PARTIAL,
   FUNCTION_TYPE_LAST_MERGE,
+  FUNCTION_TYPE_AVG_PARTIAL,
+  FUNCTION_TYPE_AVG_MERGE,
 
   // user defined funcion
   FUNCTION_TYPE_UDF = 10000
diff --git a/source/libs/function/inc/builtinsimpl.h b/source/libs/function/inc/builtinsimpl.h
index 7db1396603..2f31d57907 100644
--- a/source/libs/function/inc/builtinsimpl.h
+++ b/source/libs/function/inc/builtinsimpl.h
@@ -58,6 +58,7 @@ int32_t avgFunction(SqlFunctionCtx* pCtx);
 int32_t avgFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock);
 int32_t avgInvertFunction(SqlFunctionCtx* pCtx);
 int32_t avgCombine(SqlFunctionCtx* pDestCtx, SqlFunctionCtx* pSourceCtx);
+int32_t getAvgInfoSize();
 
 bool getStddevFuncEnv(struct SFunctionNode* pFunc, SFuncExecEnv* pEnv);
 bool stddevFunctionSetup(SqlFunctionCtx *pCtx, SResultRowEntryInfo* pResultInfo);
diff --git a/source/libs/function/src/builtins.c b/source/libs/function/src/builtins.c
index 1c3f274ee9..4fd0db4819 100644
--- a/source/libs/function/src/builtins.c
+++ b/source/libs/function/src/builtins.c
@@ -155,6 +155,35 @@ static int32_t translateSum(SFunctionNode* pFunc, char* pErrBuf, int32_t len) {
   return TSDB_CODE_SUCCESS;
 }
 
+static int32_t translateAvgPartial(SFunctionNode* pFunc, char* pErrBuf, int32_t len) {
+  if (1 != LIST_LENGTH(pFunc->pParameterList)) {
+    return invaildFuncParaNumErrMsg(pErrBuf, len, pFunc->functionName);
+  }
+
+  uint8_t paraType = ((SExprNode*)nodesListGetNode(pFunc->pParameterList, 0))->resType.type;
+  if (!IS_NUMERIC_TYPE(paraType) && !IS_NULL_TYPE(paraType)) {
+    return invaildFuncParaTypeErrMsg(pErrBuf, len, pFunc->functionName);
+  }
+
+  pFunc->node.resType = (SDataType){.bytes = getAvgInfoSize() + VARSTR_HEADER_SIZE, .type = TSDB_DATA_TYPE_BINARY};
+  return TSDB_CODE_SUCCESS;
+}
+
+static int32_t translateAvgMerge(SFunctionNode* pFunc, char* pErrBuf, int32_t len) {
+  if (1 != LIST_LENGTH(pFunc->pParameterList)) {
+    return invaildFuncParaNumErrMsg(pErrBuf, len, pFunc->functionName);
+  }
+
+  uint8_t paraType = ((SExprNode*)nodesListGetNode(pFunc->pParameterList, 0))->resType.type;
+  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 translateWduration(SFunctionNode* pFunc, char* pErrBuf, int32_t len) {
   // pseudo column do not need to check parameters
   pFunc->node.resType = (SDataType){.bytes = sizeof(int64_t), .type = TSDB_DATA_TYPE_BIGINT};
@@ -1518,6 +1547,30 @@ const SBuiltinFuncDefinition funcMgtBuiltins[] = {
     .invertFunc   = avgInvertFunction,
     .combineFunc  = avgCombine,
   },
+  {
+    .name = "_avg_partial",
+    .type = FUNCTION_TYPE_AVG_PARTIAL,
+    .classification = FUNC_MGT_AGG_FUNC,
+    .translateFunc = translateInNumOutDou,
+    .getEnvFunc   = getAvgFuncEnv,
+    .initFunc     = avgFunctionSetup,
+    .processFunc  = avgFunction,
+    .finalizeFunc = avgFinalize,
+    .invertFunc   = avgInvertFunction,
+    .combineFunc  = avgCombine,
+  },
+  {
+    .name = "_avg_merge",
+    .type = FUNCTION_TYPE_AVG_MERGE,
+    .classification = FUNC_MGT_AGG_FUNC,
+    .translateFunc = translateInNumOutDou,
+    .getEnvFunc   = getAvgFuncEnv,
+    .initFunc     = avgFunctionSetup,
+    .processFunc  = avgFunction,
+    .finalizeFunc = avgFinalize,
+    .invertFunc   = avgInvertFunction,
+    .combineFunc  = avgCombine,
+  },
   {
     .name = "percentile",
     .type = FUNCTION_TYPE_PERCENTILE,
diff --git a/source/libs/function/src/builtinsimpl.c b/source/libs/function/src/builtinsimpl.c
index 3510a59254..ee47b1aa7f 100644
--- a/source/libs/function/src/builtinsimpl.c
+++ b/source/libs/function/src/builtinsimpl.c
@@ -624,6 +624,8 @@ bool getSumFuncEnv(SFunctionNode* UNUSED_PARAM(pFunc), SFuncExecEnv* pEnv) {
   return true;
 }
 
+int32_t getAvgInfoSize() { return (int32_t)sizeof(SAvgRes); }
+
 bool getAvgFuncEnv(SFunctionNode* UNUSED_PARAM(pFunc), SFuncExecEnv* pEnv) {
   pEnv->calcMemSize = sizeof(SAvgRes);
   return true;