fix: udf2 support int and double data type

This commit is contained in:
slzhou 2022-05-12 22:58:49 +08:00
parent 3967f79b38
commit 73d55e7c6b
1 changed files with 23 additions and 10 deletions

View File

@ -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;