From 73d55e7c6b9fae35f2e596962d0667bb9213a273 Mon Sep 17 00:00:00 2001
From: slzhou <slzhou@taosdata.com>
Date: Thu, 12 May 2022 22:58:49 +0800
Subject: [PATCH] fix: udf2 support int and double data type

---
 source/libs/function/test/udf2.c | 33 ++++++++++++++++++++++----------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/source/libs/function/test/udf2.c b/source/libs/function/test/udf2.c
index b3b60f93a4..ba39b09f56 100644
--- a/source/libs/function/test/udf2.c
+++ b/source/libs/function/test/udf2.c
@@ -19,17 +19,18 @@ int32_t udf2_destroy() {
 
 int32_t udf2_start(SUdfInterBuf *buf) {
   *(int64_t*)(buf->buf) = 0;
-  buf->bufLen = sizeof(int64_t);
+  buf->bufLen = sizeof(double);
   buf->numOfResult = 0;
   return 0;
 }
 
 int32_t udf2(SUdfDataBlock* block, SUdfInterBuf *interBuf, SUdfInterBuf *newInterBuf) {
-  int64_t sumSquares = *(int64_t*)interBuf->buf;
+  double sumSquares = *(double*)interBuf->buf;
   int8_t numOutput = 0;
   for (int32_t i = 0; i < block->numOfCols; ++i) {
     SUdfColumn* col = block->udfCols[i];
-    if (col->colMeta.type != TSDB_DATA_TYPE_INT) {
+    if (!(col->colMeta.type == TSDB_DATA_TYPE_INT || 
+          col->colMeta.type == TSDB_DATA_TYPE_DOUBLE)) {
       return TSDB_CODE_UDF_INVALID_INPUT;
     }
   }
@@ -39,17 +40,29 @@ int32_t udf2(SUdfDataBlock* block, SUdfInterBuf *interBuf, SUdfInterBuf *newInte
       if (udfColDataIsNull(col, j)) {
         continue;
       }
-
-      char* cell = udfColDataGetData(col, j);
-      int32_t num = *(int32_t*)cell;
-      sumSquares += num * num;
+      switch (col->colMeta.type) {
+        case TSDB_DATA_TYPE_INT: {
+          char* cell = udfColDataGetData(col, j);
+          int32_t num = *(int32_t*)cell;
+          sumSquares += num * num;
+          break;
+        }
+        case TSDB_DATA_TYPE_DOUBLE: {
+          char* cell = udfColDataGetData(col, j);
+          double num = *(double*)cell;
+          sumSquares += num * num;
+          break;
+        }
+        default: 
+          break;
+      }
       numOutput = 1;
     }
   }
 
   if (numOutput == 1) {
-    *(int64_t*)(newInterBuf->buf) = sumSquares;
-    newInterBuf->bufLen = sizeof(int64_t);
+    *(double*)(newInterBuf->buf) = sumSquares;
+    newInterBuf->bufLen = sizeof(double);
   }
   newInterBuf->numOfResult = numOutput;
   return 0;
@@ -60,7 +73,7 @@ int32_t udf2_finish(SUdfInterBuf* buf, SUdfInterBuf *resultData) {
     resultData->numOfResult = 0;
     return 0;
   }
-  int64_t sumSquares = *(int64_t*)(buf->buf);
+  double sumSquares = *(double*)(buf->buf);
   *(double*)(resultData->buf) = sqrt(sumSquares);
   resultData->bufLen = sizeof(double);
   resultData->numOfResult = 1;