From c190fa5bd92c4c23c1699c645480e08df102864d Mon Sep 17 00:00:00 2001
From: wangmm0220 <wangmm0220@gmail.com>
Date: Mon, 27 Jun 2022 10:51:47 +0800
Subject: [PATCH] fix: keep n rows if limit n

---
 include/common/tdatablock.h             |  1 +
 source/common/src/tdatablock.c          | 28 +++++++++++++++++++++++++
 source/libs/executor/src/executorimpl.c |  5 +++--
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/include/common/tdatablock.h b/include/common/tdatablock.h
index 2a0d4e7ff6..60f517c7ba 100644
--- a/include/common/tdatablock.h
+++ b/include/common/tdatablock.h
@@ -222,6 +222,7 @@ void blockDataCleanup(SSDataBlock* pDataBlock);
 size_t blockDataGetCapacityInRow(const SSDataBlock* pBlock, size_t pageSize);
 
 int32_t blockDataTrimFirstNRows(SSDataBlock* pBlock, size_t n);
+int32_t blockDataKeepFirstNRows(SSDataBlock* pBlock, size_t n);
 
 int32_t      assignOneDataBlock(SSDataBlock* dst, const SSDataBlock* src);
 SSDataBlock* createOneDataBlock(const SSDataBlock* pDataBlock, bool copyData);
diff --git a/source/common/src/tdatablock.c b/source/common/src/tdatablock.c
index 593f8c5c0b..23bb14165d 100644
--- a/source/common/src/tdatablock.c
+++ b/source/common/src/tdatablock.c
@@ -1407,6 +1407,7 @@ static void doShiftBitmap(char* nullBitmap, size_t n, size_t total) {
 
 static void colDataTrimFirstNRows(SColumnInfoData* pColInfoData, size_t n, size_t total) {
   if (IS_VAR_DATA_TYPE(pColInfoData->info.type)) {
+    pColInfoData->varmeta.length -= pColInfoData->varmeta.offset[n];
     memmove(pColInfoData->varmeta.offset, &pColInfoData->varmeta.offset[n], (total - n) * sizeof(int32_t));
     memset(&pColInfoData->varmeta.offset[total - n], 0, n);
   } else {
@@ -1435,6 +1436,33 @@ int32_t blockDataTrimFirstNRows(SSDataBlock* pBlock, size_t n) {
   return TSDB_CODE_SUCCESS;
 }
 
+static void colDataKeepFirstNRows(SColumnInfoData* pColInfoData, size_t n, size_t total) {
+  if (IS_VAR_DATA_TYPE(pColInfoData->info.type)) {
+    pColInfoData->varmeta.length = pColInfoData->varmeta.offset[n] - pColInfoData->varmeta.offset[0];
+    memset(&pColInfoData->varmeta.offset[n], 0, total - n);
+  }
+}
+
+int32_t blockDataKeepFirstNRows(SSDataBlock* pBlock, size_t n) {
+  if (n == 0) {
+    blockDataCleanup(pBlock);
+    return TSDB_CODE_SUCCESS;
+  }
+
+  if (pBlock->info.rows <= n) {
+    return TSDB_CODE_SUCCESS;
+  } else {
+    size_t numOfCols = taosArrayGetSize(pBlock->pDataBlock);
+    for (int32_t i = 0; i < numOfCols; ++i) {
+      SColumnInfoData* pColInfoData = taosArrayGet(pBlock->pDataBlock, i);
+      colDataKeepFirstNRows(pColInfoData, n, pBlock->info.rows);
+    }
+
+    pBlock->info.rows = n;
+  }
+  return TSDB_CODE_SUCCESS;
+}
+
 int32_t tEncodeDataBlock(void** buf, const SSDataBlock* pBlock) {
   int64_t tbUid = pBlock->info.uid;
   int16_t numOfCols = taosArrayGetSize(pBlock->pDataBlock);
diff --git a/source/libs/executor/src/executorimpl.c b/source/libs/executor/src/executorimpl.c
index d680614799..897d607d20 100644
--- a/source/libs/executor/src/executorimpl.c
+++ b/source/libs/executor/src/executorimpl.c
@@ -3096,8 +3096,9 @@ static int32_t handleLimitOffset(SOperatorInfo* pOperator, SSDataBlock* pBlock)
   }
 
   // check for the limitation in each group
-  if (pProjectInfo->limit.limit > 0 && pProjectInfo->curOutput + pRes->info.rows >= pProjectInfo->limit.limit) {
-    pRes->info.rows = (int32_t)(pProjectInfo->limit.limit - pProjectInfo->curOutput);
+  if (pProjectInfo->limit.limit >= 0 && pProjectInfo->curOutput + pRes->info.rows >= pProjectInfo->limit.limit) {
+    int32_t keepRows = (int32_t)(pProjectInfo->limit.limit - pProjectInfo->curOutput);
+    blockDataKeepFirstNRows(pRes, keepRows);
     if (pProjectInfo->slimit.limit > 0 && pProjectInfo->slimit.limit <= pProjectInfo->curGroupOutput) {
       pOperator->status = OP_EXEC_DONE;
     }