From 272494d1c520af6ba1ba156636cf2350258cd7b8 Mon Sep 17 00:00:00 2001
From: Haojun Liao <hjliao@taosdata.com>
Date: Sat, 3 Apr 2021 23:45:56 +0800
Subject: [PATCH] [td-3662] <fix>: fix invalid write caused by top/bottom
 query.

---
 src/client/inc/tscUtil.h       |  2 ++
 src/client/src/tscLocalMerge.c | 11 +++++++++--
 src/client/src/tscUtil.c       | 35 ++++++++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/client/inc/tscUtil.h b/src/client/inc/tscUtil.h
index 9ef13cccbf..0eda49b1f4 100644
--- a/src/client/inc/tscUtil.h
+++ b/src/client/inc/tscUtil.h
@@ -129,6 +129,8 @@ bool tscIsPointInterpQuery(SQueryInfo* pQueryInfo);
 bool tscIsTWAQuery(SQueryInfo* pQueryInfo);
 bool tscIsSecondStageQuery(SQueryInfo* pQueryInfo);
 bool tscGroupbyColumn(SQueryInfo* pQueryInfo);
+bool tscIsTopbotQuery(SQueryInfo* pQueryInfo);
+int32_t tscGetTopbotQueryParam(SQueryInfo* pQueryInfo);
 
 bool tscNonOrderedProjectionQueryOnSTable(SQueryInfo *pQueryInfo, int32_t tableIndex);
 bool tscOrderedProjectionQueryOnSTable(SQueryInfo* pQueryInfo, int32_t tableIndex);
diff --git a/src/client/src/tscLocalMerge.c b/src/client/src/tscLocalMerge.c
index 8a301d1820..26e21170ac 100644
--- a/src/client/src/tscLocalMerge.c
+++ b/src/client/src/tscLocalMerge.c
@@ -338,12 +338,19 @@ void tscCreateLocalMerger(tExtMemBuffer **pMemBuffer, int32_t numOfBuffer, tOrde
   pReducer->resColModel->capacity = pReducer->nResultBufSize;
   pReducer->finalModel = pFFModel;
 
+  int32_t expandFactor = 1;
   if (finalmodel->rowSize > 0) {
-    pReducer->resColModel->capacity /= finalmodel->rowSize;
+    bool topBotQuery = tscIsTopbotQuery(pQueryInfo);
+    if (topBotQuery) {
+      expandFactor = tscGetTopbotQueryParam(pQueryInfo);
+      pReducer->resColModel->capacity /= (finalmodel->rowSize * expandFactor);
+    } else {
+      pReducer->resColModel->capacity /= finalmodel->rowSize;
+    }
   }
 
   assert(finalmodel->rowSize > 0 && finalmodel->rowSize <= pReducer->rowSize);
-  pReducer->pFinalRes = calloc(1, pReducer->rowSize * pReducer->resColModel->capacity);
+  pReducer->pFinalRes = calloc(1, pReducer->rowSize * pReducer->resColModel->capacity * expandFactor);
 
   if (pReducer->pTempBuffer == NULL || pReducer->discardData == NULL || pReducer->pResultBuf == NULL ||
       pReducer->pFinalRes == NULL || pReducer->prevRowOfInput == NULL) {
diff --git a/src/client/src/tscUtil.c b/src/client/src/tscUtil.c
index 386352bdc7..c6c0c9324e 100644
--- a/src/client/src/tscUtil.c
+++ b/src/client/src/tscUtil.c
@@ -271,6 +271,41 @@ bool tscIsTWAQuery(SQueryInfo* pQueryInfo) {
   return false;
 }
 
+bool tscIsTopbotQuery(SQueryInfo* pQueryInfo) {
+  size_t numOfExprs = tscSqlExprNumOfExprs(pQueryInfo);
+  for (int32_t i = 0; i < numOfExprs; ++i) {
+    SSqlExpr* pExpr = tscSqlExprGet(pQueryInfo, i);
+    if (pExpr == NULL) {
+      continue;
+    }
+
+    int32_t functionId = pExpr->functionId;
+    if (functionId == TSDB_FUNC_TOP || functionId == TSDB_FUNC_BOTTOM) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+int32_t tscGetTopbotQueryParam(SQueryInfo* pQueryInfo) {
+  size_t numOfExprs = tscSqlExprNumOfExprs(pQueryInfo);
+  for (int32_t i = 0; i < numOfExprs; ++i) {
+    SSqlExpr* pExpr = tscSqlExprGet(pQueryInfo, i);
+    if (pExpr == NULL) {
+      continue;
+    }
+
+    int32_t functionId = pExpr->functionId;
+    if (functionId == TSDB_FUNC_TOP || functionId == TSDB_FUNC_BOTTOM) {
+      return pExpr->param[0].i64;
+    }
+  }
+
+  return 0;
+}
+
+
 void tscClearInterpInfo(SQueryInfo* pQueryInfo) {
   if (!tscIsPointInterpQuery(pQueryInfo)) {
     return;
-- 
GitLab