fix: keep n rows if limit n

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

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

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