xxq250
/
homework-jianmu
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

[TD-3980] support lua script

This commit is contained in:
yihaoDeng 2021-05-08 15:22:27 +08:00
parent 148a73fe50
commit 8155c9fb34
5 changed files with 79 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/query/inc/qScript.h
View File

@ -60,8 +60,9 @@ typedef struct ScriptCtx {
int taosLoadScriptInit(void *pInit);
void taosLoadScriptNormal(void *pInit, char *pInput, int16_t iType, int16_t iBytes, int32_t numOfRows,
int64_t *ptsList, char* pOutput, char *ptsOutput, int32_t *numOfOutput, int16_t oType, int16_t oBytes);
void taosLoadScriptFinalize(void *pInit, char *pOutput, int32_t *output);
int64_t *ptsList, int64_t key, char* pOutput, char *ptsOutput, int32_t *numOfOutput, int16_t oType, int16_t oBytes);
void taosLoadScriptFinalize(void *pInit, int64_t key, char *pOutput, int32_t *output);
void taosLoadScriptMerge(void *pCtx, char* data, int32_t numOfRows, char* dataOutput, int32_t* numOfOutput);
void taosLoadScriptDestroy(void *pInit);
typedef struct {

4
src/query/inc/qUdf.h
View File

@ -57,8 +57,8 @@ typedef struct SUdfInfo {
typedef int32_t (*scriptInitFunc)(void *pCtx);
typedef void (*scriptNormalFunc)(void *pCtx, char* data, int16_t iType, int16_t iBytes, int32_t numOfRows,
int64_t* ts, char* dataOutput, char* tsOutput, int32_t* numOfOutput, int16_t oType, int16_t oBytes);
typedef void (*scriptFinalizeFunc)(void *pCtx, char* dataOutput, int32_t* numOfOutput);
int64_t* ptList, int64_t key, char* dataOutput, char* tsOutput, int32_t* numOfOutput, int16_t oType, int16_t oBytes);
typedef void (*scriptFinalizeFunc)(void *pCtx, int64_t key, char* dataOutput, int32_t* numOfOutput);
typedef void (*scriptMergeFunc)(void *pCtx, char* data, int32_t numOfRows, char* dataOutput, int32_t* numOfOutput);
typedef void (*scriptDestroyFunc)(void* pCtx);

10
src/query/src/qExecutor.c
View File

@ -715,7 +715,7 @@ void doInvokeUdf(SUdfInfo* pUdfInfo, SQLFunctionCtx *pCtx, int32_t idx, int32_t
case TSDB_UDF_FUNC_NORMAL:
if (pUdfInfo->isScript) {
(*(scriptNormalFunc)pUdfInfo->funcs[TSDB_UDF_FUNC_NORMAL])(pUdfInfo->pScriptCtx,
(char *)pCtx->pInput + idx * pCtx->inputType, pCtx->inputType, pCtx->inputBytes, pCtx->size, pCtx->ptsList, pCtx->pOutput,
(char *)pCtx->pInput + idx * pCtx->inputType, pCtx->inputType, pCtx->inputBytes, pCtx->size, pCtx->ptsList, pCtx->startTs, pCtx->pOutput,
(char *)pCtx->ptsOutputBuf, &output, pCtx->outputType, pCtx->outputBytes);
} else {
SResultRowCellInfo *pResInfo = GET_RES_INFO(pCtx);
@ -739,7 +739,11 @@ void doInvokeUdf(SUdfInfo* pUdfInfo, SQLFunctionCtx *pCtx, int32_t idx, int32_t
break;
case TSDB_UDF_FUNC_MERGE:
if (pUdfInfo->isScript) {
(*(scriptMergeFunc)pUdfInfo->funcs[TSDB_UDF_FUNC_MERGE])(pUdfInfo->pScriptCtx, pCtx->pInput, pCtx->size, pCtx->pOutput, &output);
} else {
(*(udfMergeFunc)pUdfInfo->funcs[TSDB_UDF_FUNC_MERGE])(pCtx->pInput, pCtx->size, pCtx->pOutput, &output, &pUdfInfo->init);
}
// set the output value exist
pCtx->resultInfo->numOfRes = output;
@ -753,7 +757,7 @@ void doInvokeUdf(SUdfInfo* pUdfInfo, SQLFunctionCtx *pCtx, int32_t idx, int32_t
SResultRowCellInfo *pResInfo = GET_RES_INFO(pCtx);
void *interBuf = (void *)GET_ROWCELL_INTERBUF(pResInfo);
if (pUdfInfo->isScript) {
(*(scriptFinalizeFunc)pUdfInfo->funcs[TSDB_UDF_FUNC_FINALIZE])(pUdfInfo->pScriptCtx, pCtx->pOutput, &output);
(*(scriptFinalizeFunc)pUdfInfo->funcs[TSDB_UDF_FUNC_FINALIZE])(pUdfInfo->pScriptCtx, pCtx->startTs, pCtx->pOutput, &output);
} else {
(*(udfFinalizeFunc)pUdfInfo->funcs[TSDB_UDF_FUNC_FINALIZE])(pCtx->pOutput, interBuf, &output, &pUdfInfo->init);
}
@ -3279,6 +3283,7 @@ void finalizeQueryResult(SOperatorInfo* pOperator, SQLFunctionCtx* pCtx, SResult
setResultOutputBuf(pRuntimeEnv, buf, pCtx, numOfOutput, rowCellInfoOffset);
for (int32_t j = 0; j < numOfOutput; ++j) {
pCtx[j].startTs = buf->win.skey;
if (pCtx[j].functionId < 0) {
doInvokeUdf(pRuntimeEnv->pUdfInfo, &pCtx[j], 0, TSDB_UDF_FUNC_FINALIZE);
} else {
@ -6676,6 +6681,7 @@ int32_t initUdfInfo(SUdfInfo* pUdfInfo) {
if (pUdfInfo->funcType == TSDB_UDF_TYPE_AGGREGATE) {
pUdfInfo->funcs[TSDB_UDF_FUNC_FINALIZE] = taosLoadScriptFinalize;
pUdfInfo->funcs[TSDB_UDF_FUNC_MERGE] = taosLoadScriptMerge;
}
pUdfInfo->funcs[TSDB_UDF_FUNC_DESTROY] = taosLoadScriptDestroy;

40
src/query/src/qScript.c
View File

@ -103,7 +103,7 @@ int taosLoadScriptInit(void* pInit) {
return 0;
}
void taosLoadScriptNormal(void *pInit, char *pInput, int16_t iType, int16_t iBytes, int32_t numOfRows,
int64_t *ptsList, char* pOutput, char *ptsOutput, int32_t *numOfOutput, int16_t oType, int16_t oBytes) {
int64_t *ptsList, int64_t key, char* pOutput, char *ptsOutput, int32_t *numOfOutput, int16_t oType, int16_t oBytes) {
ScriptCtx* pCtx = pInit;
char funcName[MAX_FUNC_NAME] = {0};
sprintf(funcName, "%s_add", pCtx->funcName);
@ -124,8 +124,9 @@ void taosLoadScriptNormal(void *pInit, char *pInput, int16_t iType, int16_t iByt
if (lua_istable(lua, -1)) {
isGlobalState = true;
}
lua_pushnumber(lua, key);
// do call lua script
if (lua_pcall(lua, 2, 1, 0) != 0) {
if (lua_pcall(lua, 3, 1, 0) != 0) {
qError("SCRIPT ERROR: %s", lua_tostring(lua, -1));
lua_pop(lua, -1);
return;
@ -139,22 +140,51 @@ void taosLoadScriptNormal(void *pInit, char *pInput, int16_t iType, int16_t iByt
*numOfOutput = tNumOfOutput;
}
void taosLoadScriptMerge(void *pInit, char* data, int32_t numOfRows, char* pOutput, int32_t* numOfOutput) {
ScriptCtx *pCtx = pInit;
char funcName[MAX_FUNC_NAME] = {0};
sprintf(funcName, "%s_merge", pCtx->funcName);
lua_State* lua = pCtx->pEnv->lua_state;
lua_getglobal(lua, funcName);
if (!lua_isfunction(lua, -1)) {
qError("SCRIPT ERROR: %s", lua_tostring(lua, -1));
return;
}
lua_getglobal(lua, "global");
if (lua_pcall(lua, 1, 1, 0) != 0) {
qError("SCRIPT ERROR: %s", lua_tostring(lua, -1));
lua_pop(lua, -1);
return;
}
int tNumOfOutput = 0;
luaValueToTaosType(lua, pOutput, &tNumOfOutput, pCtx->resType, pCtx->resBytes);
*numOfOutput = tNumOfOutput;
}
//do not support agg now
void taosLoadScriptFinalize(void *pInit, char *pOutput, int32_t* numOfOutput) {
void taosLoadScriptFinalize(void *pInit,int64_t key, char *pOutput, int32_t* numOfOutput) {
ScriptCtx *pCtx = pInit;
char funcName[MAX_FUNC_NAME] = {0};
sprintf(funcName, "%s_finalize", pCtx->funcName);
lua_State* lua = pCtx->pEnv->lua_state;
lua_getglobal(lua, funcName);
if (!lua_isfunction(lua, -1)) {
qError("SCRIPT ERROR: %s", lua_tostring(lua, -1));
return;
}
lua_getglobal(lua, "global");
if (lua_pcall(lua, 1, 1, 0) != 0) {
lua_pushnumber(lua, key);
if (lua_pcall(lua, 2, 2, 0) != 0) {
qError("SCRIPT ERROR: %s", lua_tostring(lua, -1));
lua_pop(lua, -1);
return;
}
lua_setglobal(lua, "global");
int tNumOfOutput = 0;
luaValueToTaosType(lua, pOutput, &tNumOfOutput, pCtx->resType, pCtx->resBytes);
*numOfOutput = tNumOfOutput;
@ -256,6 +286,8 @@ void luaValueToTaosType(lua_State *lua, char *interBuf, int32_t *numOfOutput, in
}
break;
default:
setNull(interBuf, oType, oBytes);
sz = 1;
break;
}
lua_pop(lua,1); // pop ret value from script

36
tests/script/sh/demo.lua
View File

@ -1,21 +1,43 @@
funcName = "test"
global = {}
global["sum"] = 0.0
global["num"] = 0
function test_init()
return global
end
function test_add(rows, ans)
function test_add(rows, ans, key)
t = {}
t["sum"] = 0.0
t["num"] = 0
for i=1, #rows do
ans["sum"] = ans["sum"] + rows[i] * rows[i]
t["sum"] = t["sum"] + rows[i] * rows[i]
end
ans["num"] = ans["num"] + #rows
t["num"] = #rows
if (ans[key] ~= nil)
then
ans[key]["sum"] = ans[key]["sum"] + t["sum"]
ans[key]["num"] = ans[key]["num"] + t["num"]
else
ans[key] = t
end
return ans;
end
function test_finalize(ans)
return ans["sum"]/ans["num"];
function test_finalize(ans, key)
local ret = 0.0
if (ans[key] ~= nil)
then
ret = ans[key]["sum"]/ans[key]["num"]
ans[key]["sum"] = 0.0
ans[key]["num"] = 0
else
ret = inf
end
return ret, ans
end