From b4fa0d61b49221af5d4c6cd881cc864086c459fd Mon Sep 17 00:00:00 2001
From: Ganlin Zhao <ganlinzhao@hotmail.com>
Date: Mon, 1 Aug 2022 15:16:29 +0800
Subject: [PATCH] enh(query): add selectivity support for
 statecount/stateduration functions

TD-18050
---
 source/libs/function/src/builtins.c     |  4 ++--
 source/libs/function/src/builtinsimpl.c | 21 ++++++++++++++++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/source/libs/function/src/builtins.c b/source/libs/function/src/builtins.c
index e58e7475df..e16c6db7c6 100644
--- a/source/libs/function/src/builtins.c
+++ b/source/libs/function/src/builtins.c
@@ -2465,7 +2465,7 @@ const SBuiltinFuncDefinition funcMgtBuiltins[] = {
   {
     .name = "statecount",
     .type = FUNCTION_TYPE_STATE_COUNT,
-    .classification = FUNC_MGT_INDEFINITE_ROWS_FUNC | FUNC_MGT_TIMELINE_FUNC | FUNC_MGT_FORBID_STREAM_FUNC,
+    .classification = FUNC_MGT_INDEFINITE_ROWS_FUNC | FUNC_MGT_SELECT_FUNC | FUNC_MGT_TIMELINE_FUNC | FUNC_MGT_FORBID_STREAM_FUNC,
     .translateFunc = translateStateCount,
     .getEnvFunc   = getStateFuncEnv,
     .initFunc     = functionSetup,
@@ -2476,7 +2476,7 @@ const SBuiltinFuncDefinition funcMgtBuiltins[] = {
   {
     .name = "stateduration",
     .type = FUNCTION_TYPE_STATE_DURATION,
-    .classification = FUNC_MGT_INDEFINITE_ROWS_FUNC | FUNC_MGT_TIMELINE_FUNC | FUNC_MGT_IMPLICIT_TS_FUNC | FUNC_MGT_FORBID_STREAM_FUNC,
+    .classification = FUNC_MGT_INDEFINITE_ROWS_FUNC | FUNC_MGT_SELECT_FUNC | FUNC_MGT_TIMELINE_FUNC | FUNC_MGT_IMPLICIT_TS_FUNC | FUNC_MGT_FORBID_STREAM_FUNC,
     .translateFunc = translateStateDuration,
     .getEnvFunc   = getStateFuncEnv,
     .initFunc     = functionSetup,
diff --git a/source/libs/function/src/builtinsimpl.c b/source/libs/function/src/builtinsimpl.c
index 0767c2e5a2..94e7b77a48 100644
--- a/source/libs/function/src/builtinsimpl.c
+++ b/source/libs/function/src/builtinsimpl.c
@@ -4651,10 +4651,15 @@ int32_t stateCountFunction(SqlFunctionCtx* pCtx) {
     numOfElems++;
     if (colDataIsNull_f(pInputCol->nullbitmap, i)) {
       colDataAppendNULL(pOutput, i);
+      // handle selectivity
+      if (pCtx->subsidiaries.num > 0) {
+        appendSelectivityValue(pCtx, i, i);
+      }
       continue;
     }
 
-    bool    ret = checkStateOp(op, pInputCol, i, pCtx->param[2].param);
+    bool ret = checkStateOp(op, pInputCol, i, pCtx->param[2].param);
+
     int64_t output = -1;
     if (ret) {
       output = ++pInfo->count;
@@ -4662,6 +4667,11 @@ int32_t stateCountFunction(SqlFunctionCtx* pCtx) {
       pInfo->count = 0;
     }
     colDataAppend(pOutput, i, (char*)&output, false);
+
+    // handle selectivity
+    if (pCtx->subsidiaries.num > 0) {
+      appendSelectivityValue(pCtx, i, i);
+    }
   }
 
   return numOfElems;
@@ -4694,6 +4704,10 @@ int32_t stateDurationFunction(SqlFunctionCtx* pCtx) {
     numOfElems++;
     if (colDataIsNull_f(pInputCol->nullbitmap, i)) {
       colDataAppendNULL(pOutput, i);
+      // handle selectivity
+      if (pCtx->subsidiaries.num > 0) {
+        appendSelectivityValue(pCtx, i, i);
+      }
       continue;
     }
 
@@ -4710,6 +4724,11 @@ int32_t stateDurationFunction(SqlFunctionCtx* pCtx) {
       pInfo->durationStart = 0;
     }
     colDataAppend(pOutput, i, (char*)&output, false);
+
+    // handle selectivity
+    if (pCtx->subsidiaries.num > 0) {
+      appendSelectivityValue(pCtx, i, i);
+    }
   }
 
   return numOfElems;