From 9f529b7d3bb31080ff02e7b7650e4ef8689b8f39 Mon Sep 17 00:00:00 2001
From: Ganlin Zhao <ganlinzhao@hotmail.com>
Date: Mon, 4 Jul 2022 20:34:46 +0800
Subject: [PATCH] fix(query): avg function support unsigned data type

TD-16877
---
 source/libs/function/src/builtinsimpl.c | 70 +++++++++++++++++++++++--
 1 file changed, 67 insertions(+), 3 deletions(-)

diff --git a/source/libs/function/src/builtinsimpl.c b/source/libs/function/src/builtinsimpl.c
index 3f6afaf5fd..40fd7f00b2 100644
--- a/source/libs/function/src/builtinsimpl.c
+++ b/source/libs/function/src/builtinsimpl.c
@@ -719,8 +719,10 @@ int32_t avgFunction(SqlFunctionCtx* pCtx) {
     ASSERT(numOfElem >= 0);
 
     pAvgRes->count += numOfElem;
-    if (IS_INTEGER_TYPE(type)) {
+    if (IS_SIGNED_NUMERIC_TYPE(type)) {
       pAvgRes->sum.isum += pAgg->sum;
+    } else if (IS_UNSIGNED_NUMERIC_TYPE(type)) {
+      pAvgRes->sum.usum += pAgg->sum;
     } else if (IS_FLOAT_TYPE(type)) {
       pAvgRes->sum.dsum += GET_DOUBLE_VAL((const char*)&(pAgg->sum));
     }
@@ -784,6 +786,64 @@ int32_t avgFunction(SqlFunctionCtx* pCtx) {
         break;
       }
 
+      case TSDB_DATA_TYPE_UTINYINT: {
+        uint8_t* plist = (uint8_t*)pCol->pData;
+        for (int32_t i = start; i < numOfRows + pInput->startRowIndex; ++i) {
+          if (pCol->hasNull && colDataIsNull_f(pCol->nullbitmap, i)) {
+            continue;
+          }
+
+          numOfElem += 1;
+          pAvgRes->count += 1;
+          pAvgRes->sum.usum += plist[i];
+        }
+
+        break;
+      }
+
+      case TSDB_DATA_TYPE_USMALLINT: {
+        uint16_t* plist = (uint16_t*)pCol->pData;
+        for (int32_t i = start; i < numOfRows + pInput->startRowIndex; ++i) {
+          if (pCol->hasNull && colDataIsNull_f(pCol->nullbitmap, i)) {
+            continue;
+          }
+
+          numOfElem += 1;
+          pAvgRes->count += 1;
+          pAvgRes->sum.usum += plist[i];
+        }
+        break;
+      }
+
+      case TSDB_DATA_TYPE_UINT: {
+        uint32_t* plist = (uint32_t*)pCol->pData;
+        for (int32_t i = start; i < numOfRows + pInput->startRowIndex; ++i) {
+          if (pCol->hasNull && colDataIsNull_f(pCol->nullbitmap, i)) {
+            continue;
+          }
+
+          numOfElem += 1;
+          pAvgRes->count += 1;
+          pAvgRes->sum.usum += plist[i];
+        }
+
+        break;
+      }
+
+      case TSDB_DATA_TYPE_UBIGINT: {
+        uint64_t* plist = (uint64_t*)pCol->pData;
+        for (int32_t i = start; i < numOfRows + pInput->startRowIndex; ++i) {
+          if (pCol->hasNull && colDataIsNull_f(pCol->nullbitmap, i)) {
+            continue;
+          }
+
+          numOfElem += 1;
+          pAvgRes->count += 1;
+          pAvgRes->sum.usum += plist[i];
+        }
+        break;
+      }
+
       case TSDB_DATA_TYPE_FLOAT: {
         float* plist = (float*)pCol->pData;
         for (int32_t i = start; i < numOfRows + pInput->startRowIndex; ++i) {
@@ -825,8 +885,10 @@ _avg_over:
 
 static void avgTransferInfo(SAvgRes* pInput, SAvgRes* pOutput) {
   pOutput->type = pInput->type;
-  if (IS_INTEGER_TYPE(pOutput->type)) {
+  if (IS_SIGNED_NUMERIC_TYPE(pOutput->type)) {
     pOutput->sum.isum += pInput->sum.isum;
+  } else if (IS_UNSIGNED_NUMERIC_TYPE(pOutput->type)) {
+    pOutput->sum.usum += pInput->sum.usum;
   } else {
     pOutput->sum.dsum += pInput->sum.dsum;
   }
@@ -941,8 +1003,10 @@ int32_t avgFinalize(SqlFunctionCtx* pCtx, SSDataBlock* pBlock) {
   SAvgRes* pAvgRes = GET_ROWCELL_INTERBUF(GET_RES_INFO(pCtx));
   int32_t  type = pAvgRes->type;
 
-  if (IS_INTEGER_TYPE(type)) {
+  if (IS_SIGNED_NUMERIC_TYPE(type)) {
     pAvgRes->result = pAvgRes->sum.isum / ((double)pAvgRes->count);
+  } else if (IS_UNSIGNED_NUMERIC_TYPE(type)) {
+    pAvgRes->result = pAvgRes->sum.usum / ((double)pAvgRes->count);
   } else {
     pAvgRes->result = pAvgRes->sum.dsum / ((double)pAvgRes->count);
   }
-- 
GitLab