From 6cef3baa95578a056cfccdc6e2f8fb24a776dd28 Mon Sep 17 00:00:00 2001
From: Haojun Liao <hjliao@taosdata.com>
Date: Tue, 12 Jul 2022 15:17:39 +0800
Subject: [PATCH] fix(query): set correct number of result rows for
 taos_fetch_raw_block_a

---
 source/client/inc/clientInt.h  |  1 +
 source/client/src/clientMain.c | 20 +++++++++++++++-----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/source/client/inc/clientInt.h b/source/client/inc/clientInt.h
index 91f21f6e6a..641d235dd1 100644
--- a/source/client/inc/clientInt.h
+++ b/source/client/inc/clientInt.h
@@ -169,6 +169,7 @@ typedef struct SReqResultInfo {
   uint32_t       numOfRows;
   uint64_t       totalRows;
   uint32_t       current;
+  bool           localResultFetched;
   bool           completed;
   int32_t        precision;
   bool           convertUcs4;
diff --git a/source/client/src/clientMain.c b/source/client/src/clientMain.c
index 12de522cbc..d3fdfa7a39 100644
--- a/source/client/src/clientMain.c
+++ b/source/client/src/clientMain.c
@@ -852,23 +852,33 @@ void taos_fetch_rows_a(TAOS_RES *res, __taos_async_fn_t fp, void *param) {
   }
 
   // all data has returned to App already, no need to try again
-  if ((pResultInfo->pData == NULL || pResultInfo->current >= pResultInfo->numOfRows) && pResultInfo->completed) {
+  if (pResultInfo->completed && (pRequest->body.queryJob != 0)) {
     pResultInfo->numOfRows = 0;
     pRequest->body.fetchFp(param, pRequest, pResultInfo->numOfRows);
     return;
   }
 
   // it is a local executed query, no need to do async fetch
-  if (pResultInfo->current < pResultInfo->numOfRows && pRequest->body.queryJob == 0) {
-    pRequest->body.fetchFp(param, pRequest, pResultInfo->numOfRows);
+  if (pRequest->body.queryJob == 0) {
+    ASSERT(pResultInfo->completed && pResultInfo->numOfRows >= 0);
+    if (pResultInfo->localResultFetched) {
+      pResultInfo->numOfRows = 0;
+      pResultInfo->current = 0;
+      pRequest->body.fetchFp(param, pRequest, pResultInfo->numOfRows);
+    } else {
+      pResultInfo->localResultFetched = true;
+      pRequest->body.fetchFp(param, pRequest, pResultInfo->numOfRows);
+    }
     return;
   }
 
+
   SSchedulerReq req = {
     .syncReq = false,
     .fetchFp = fetchCallback,
     .cbParam = pRequest,
   };
+
   schedulerFetchRows(pRequest->body.queryJob, &req);
 }
 
@@ -880,10 +890,10 @@ void taos_fetch_raw_block_a(TAOS_RES *res, __taos_async_fn_t fp, void *param) {
   SReqResultInfo *pResultInfo = &pRequest->body.resInfo;
 
   // set the current block is all consumed
-  pResultInfo->current = pResultInfo->numOfRows;
   pResultInfo->convertUcs4 = false;
 
-  taos_fetch_rows_a(res, fp, param);
+  // it is a local executed query, no need to do async fetch
+  taos_fetch_rows_a(pRequest, fp, param);
 }
 
 const void *taos_get_raw_block(TAOS_RES *res) {
-- 
GitLab