refactor TAOS_RES process

include/common/tdatablock.h

@@ -136,7 +136,8 @@ static FORCE_INLINE void colDataAppendInt8(SColumnInfoData* pColumnInfoData, uin
 }
 
 static FORCE_INLINE void colDataAppendInt16(SColumnInfoData* pColumnInfoData, uint32_t currentRow, int16_t* v) {
-  ASSERT(pColumnInfoData->info.type == TSDB_DATA_TYPE_SMALLINT || pColumnInfoData->info.type == TSDB_DATA_TYPE_USMALLINT);
+  ASSERT(pColumnInfoData->info.type == TSDB_DATA_TYPE_SMALLINT ||
+         pColumnInfoData->info.type == TSDB_DATA_TYPE_USMALLINT);
   char* p = pColumnInfoData->pData + pColumnInfoData->info.bytes * currentRow;
   *(int16_t*)p = *(int16_t*)v;
 }
@@ -210,6 +211,10 @@ SSDataBlock* createOneDataBlock(const SSDataBlock* pDataBlock);
 
 void blockDebugShowData(const SArray* dataBlocks);
 
+static FORCE_INLINE int32_t blockEstimateEncodeSize(const SSDataBlock* pBlock) {
+  return blockDataGetSerialMetaSize(pBlock) + (int32_t)ceil(blockDataGetSerialRowSize(pBlock) * pBlock->info.rows);
+}
+
 static FORCE_INLINE int32_t blockCompressColData(SColumnInfoData* pColRes, int32_t numOfRows, char* data,
                                                  int8_t compressed) {
   int32_t colSize = colDataGetLength(pColRes, numOfRows);

source/client/inc/clientInt.h

@@ -186,7 +186,25 @@ typedef struct SRequestSendRecvBody {
 
 #define ERROR_MSG_BUF_DEFAULT_SIZE 512
 
+enum {
+  RES_TYPE__QUERY = 1,
+  RES_TYPE__TMQ,
+};
+
+#define TD_RES_QUERY(res) (*(int8_t*)res == RES_TYPE__QUERY)
+#define TD_RES_TMQ(res)   (*(int8_t*)res == RES_TYPE__TMQ)
+
+struct tmq_message_t {
+  int8_t     resType;
+  SMqPollRsp msg;
+  char*      topic;
+  void*      vg;
+  SArray*    res;  // SArray<SReqResultInfo>
+  int32_t    resIter;
+};
+
 typedef struct SRequestObj {
+  int8_t               resType;  // query or tmq
   uint64_t             requestId;
   int32_t              type;  // request type
   STscObj*             pTscObj;
@@ -203,6 +221,25 @@ typedef struct SRequestObj {
   SRequestSendRecvBody body;
 } SRequestObj;
 
+static FORCE_INLINE SReqResultInfo* tmqGetCurResInfo(TAOS_RES* res) {
+  tmq_message_t* msg = (tmq_message_t*)res;
+  int32_t        resIter = msg->resIter == -1 ? 0 : msg->resIter;
+  return (SReqResultInfo*)taosArrayGet(msg->res, resIter);
+}
+
+static FORCE_INLINE SReqResultInfo* tmqGetNextResInfo(TAOS_RES* res) {
+  tmq_message_t* msg = (tmq_message_t*)res;
+  if (++msg->resIter < taosArrayGetSize(msg->res)) {
+    return (SReqResultInfo*)taosArrayGet(msg->res, msg->resIter);
+  }
+  return NULL;
+}
+
+static FORCE_INLINE SReqResultInfo* tscGetCurResInfo(TAOS_RES* res) {
+  if (TD_RES_QUERY(res)) return &(((SRequestObj*)res)->body.resInfo);
+  return tmqGetCurResInfo(res);
+}
+
 extern SAppInfo appInfo;
 extern int32_t  clientReqRefPool;
 extern int32_t  clientConnRefPool;
@@ -238,14 +275,17 @@ void initMsgHandleFp();
 TAOS* taos_connect_internal(const char* ip, const char* user, const char* pass, const char* auth, const char* db,
                             uint16_t port);
 
-void* doFetchRow(SRequestObj* pRequest, bool setupOneRowPtr);
+int32_t parseSql(SRequestObj* pRequest, bool topicQuery, SQuery** pQuery);
 
-int32_t setResultDataPtr(SReqResultInfo* pResultInfo, TAOS_FIELD* pFields, int32_t numOfCols, int32_t numOfRows);
+int32_t getPlan(SRequestObj* pRequest, SQuery* pQuery, SQueryPlan** pPlan, SArray* pNodeList);
 
 int32_t buildRequest(STscObj* pTscObj, const char* sql, int sqlLen, SRequestObj** pRequest);
 
-int32_t parseSql(SRequestObj* pRequest, bool topicQuery, SQuery** pQuery);
-int32_t getPlan(SRequestObj* pRequest, SQuery* pQuery, SQueryPlan** pPlan, SArray* pNodeList);
+void*   doFetchRow(SRequestObj* pRequest, bool setupOneRowPtr);
+void    doSetOneRowPtr(SReqResultInfo* pResultInfo);
+int32_t setResultDataPtr(SReqResultInfo* pResultInfo, TAOS_FIELD* pFields, int32_t numOfCols, int32_t numOfRows);
+void    setResSchemaInfo(SReqResultInfo* pResInfo, const SSchema* pSchema, int32_t numOfCols);
+int32_t setQueryResultFromRsp(SReqResultInfo* pResultInfo, const SRetrieveTableRsp* pRsp);
 
 // --- heartbeat
 // global, called by mgmt

source/client/src/clientEnv.c

@@ -149,6 +149,7 @@ void *createRequest(STscObj *pObj, __taos_async_fn_t fp, void *param, int32_t ty
     return NULL;
   }
 
+  pRequest->resType = RES_TYPE__QUERY;
   pRequest->pDb = getDbOfConnection(pObj);
   pRequest->requestId = generateRequestId();
   pRequest->metric.start = taosGetTimestampMs();

source/client/src/clientImpl.c

@@ -13,7 +13,6 @@
 static int32_t       initEpSetFromCfg(const char* firstEp, const char* secondEp, SCorEpSet* pEpSet);
 static SMsgSendInfo* buildConnectMsg(SRequestObj* pRequest);
 static void          destroySendMsgInfo(SMsgSendInfo* pMsgBody);
-static int32_t       setQueryResultFromRsp(SReqResultInfo* pResultInfo, const SRetrieveTableRsp* pRsp);
 
 static bool stringLengthCheck(const char* str, size_t maxsize) {
   if (str == NULL) {
@@ -42,7 +41,7 @@ static char* getClusterKey(const char* user, const char* auth, const char* ip, i
 
 static STscObj* taosConnectImpl(const char* user, const char* auth, const char* db, __taos_async_fn_t fp, void* param,
                                 SAppInstInfo* pAppInfo);
-static void     setResSchemaInfo(SReqResultInfo* pResInfo, const SSchema* pSchema, int32_t numOfCols);
+
 
 TAOS* taos_connect_internal(const char* ip, const char* user, const char* pass, const char* auth, const char* db,
                             uint16_t port) {
@@ -590,7 +589,7 @@ TAOS* taos_connect_l(const char* ip, int ipLen, const char* user, int userLen, c
   return taos_connect(ipStr, userStr, passStr, dbStr, port);
 }
 
-static void doSetOneRowPtr(SReqResultInfo* pResultInfo) {
+void doSetOneRowPtr(SReqResultInfo* pResultInfo) {
   for (int32_t i = 0; i < pResultInfo->numOfCols; ++i) {
     SResultColumn* pCol = &pResultInfo->pCol[i];
 

source/client/src/clientMain.c

@@ -133,8 +133,7 @@ int taos_field_count(TAOS_RES *res) {
     return 0;
   }
 
-  SRequestObj    *pRequest = (SRequestObj *)res;
-  SReqResultInfo *pResInfo = &pRequest->body.resInfo;
+  SReqResultInfo *pResInfo = tscGetCurResInfo(res);
   return pResInfo->numOfCols;
 }
 
@@ -145,7 +144,7 @@ TAOS_FIELD *taos_fetch_fields(TAOS_RES *res) {
     return NULL;
   }
 
-  SReqResultInfo *pResInfo = &(((SRequestObj *)res)->body.resInfo);
+  SReqResultInfo *pResInfo = tscGetCurResInfo(res);
   return pResInfo->userFields;
 }
 
@@ -162,6 +161,7 @@ TAOS_ROW taos_fetch_row(TAOS_RES *res) {
     return NULL;
   }
 
+  if (TD_RES_QUERY(res)) {
     SRequestObj *pRequest = (SRequestObj *)res;
     if (pRequest->type == TSDB_SQL_RETRIEVE_EMPTY_RESULT || pRequest->type == TSDB_SQL_INSERT ||
         pRequest->code != TSDB_CODE_SUCCESS || taos_num_fields(res) == 0) {
@@ -169,6 +169,28 @@ TAOS_ROW taos_fetch_row(TAOS_RES *res) {
     }
 
     return doFetchRow(pRequest, true);
+
+  } else if (TD_RES_TMQ(res)) {
+    tmq_message_t  *msg = ((tmq_message_t *)res);
+    SReqResultInfo *pResultInfo = taosArrayGet(msg->res, msg->resIter);
+
+    doSetOneRowPtr(pResultInfo);
+    pResultInfo->current += 1;
+
+    if (pResultInfo->row == NULL) {
+      msg->resIter++;
+      pResultInfo = taosArrayGet(msg->res, msg->resIter);
+      doSetOneRowPtr(pResultInfo);
+      pResultInfo->current += 1;
+    }
+
+    return pResultInfo->row;
+
+  } else {
+    // assert to avoid uninitialization error
+    ASSERT(0);
+  }
+  return NULL;
 }
 
 int taos_print_row(char *str, TAOS_ROW row, TAOS_FIELD *fields, int num_fields) {
@@ -260,12 +282,12 @@ int *taos_fetch_lengths(TAOS_RES *res) {
     return NULL;
   }
 
-  return ((SRequestObj *)res)->body.resInfo.length;
+  SReqResultInfo *pResInfo = tscGetCurResInfo(res);
+  return pResInfo->length;
 }
 
 TAOS_ROW *taos_result_block(TAOS_RES *res) {
-  SRequestObj* pRequest = (SRequestObj*) res;
-  if (pRequest == NULL) {
+  if (res == NULL) {
     terrno = TSDB_CODE_INVALID_PARA;
     return NULL;
   }
@@ -274,7 +296,8 @@ TAOS_ROW *taos_result_block(TAOS_RES *res) {
     return NULL;
   }
 
-  return &pRequest->body.resInfo.row;
+  SReqResultInfo *pResInfo = tscGetCurResInfo(res);
+  return &pResInfo->row;
 }
 
 // todo intergrate with tDataTypes
@@ -313,7 +336,7 @@ const char *taos_data_type(int type) {
 const char *taos_get_client_info() { return version; }
 
 int taos_affected_rows(TAOS_RES *res) {
-  if (res == NULL) {
+  if (res == NULL || TD_RES_TMQ(res)) {
     return 0;
   }
 
@@ -323,12 +346,17 @@ int taos_affected_rows(TAOS_RES *res) {
 }
 
 int taos_result_precision(TAOS_RES *res) {
-  SRequestObj* pRequest = (SRequestObj*) res;
-  if (pRequest == NULL) {
+  if (res == NULL) {
     return TSDB_TIME_PRECISION_MILLI;
   }
-
+  if (TD_RES_QUERY(res)) {
+    SRequestObj *pRequest = (SRequestObj *)res;
     return pRequest->body.resInfo.precision;
+  } else if (TD_RES_TMQ(res)) {
+    SReqResultInfo *info = tmqGetCurResInfo(res);
+    return info->precision;
+  }
+  return TSDB_TIME_PRECISION_MILLI;
 }
 
 int taos_select_db(TAOS *taos, const char *db) {
@@ -370,31 +398,29 @@ void taos_stop_query(TAOS_RES *res) {
 }
 
 bool taos_is_null(TAOS_RES *res, int32_t row, int32_t col) {
-  SRequestObj    *pRequestObj = res;
-  SReqResultInfo *pResultInfo = &pRequestObj->body.resInfo;
+  SReqResultInfo *pResultInfo = tscGetCurResInfo(res);
   if (col >= pResultInfo->numOfCols || col < 0 || row >= pResultInfo->numOfRows || row < 0) {
     return true;
   }
 
-  SResultColumn *pCol = &pRequestObj->body.resInfo.pCol[col];
+  SResultColumn *pCol = &pResultInfo->pCol[col];
   return colDataIsNull_f(pCol->nullbitmap, row);
 }
 
-bool taos_is_update_query(TAOS_RES *res) {
-  return taos_num_fields(res) == 0;
-}
+bool taos_is_update_query(TAOS_RES *res) { return taos_num_fields(res) == 0; }
 
 int taos_fetch_block(TAOS_RES *res, TAOS_ROW *rows) {
   int32_t numOfRows = 0;
-  /*int32_t code = */taos_fetch_block_s(res, &numOfRows, rows);
+  /*int32_t code = */ taos_fetch_block_s(res, &numOfRows, rows);
   return numOfRows;
 }
 
-int taos_fetch_block_s(TAOS_RES *res, int* numOfRows, TAOS_ROW *rows) {
-  SRequestObj *pRequest = (SRequestObj *)res;
-  if (pRequest == NULL) {
+int taos_fetch_block_s(TAOS_RES *res, int *numOfRows, TAOS_ROW *rows) {
+  if (res == NULL) {
     return 0;
   }
+  if (TD_RES_QUERY(res)) {
+    SRequestObj *pRequest = (SRequestObj *)res;
 
     (*rows) = NULL;
     (*numOfRows) = 0;
@@ -413,13 +439,26 @@ int taos_fetch_block_s(TAOS_RES *res, int* numOfRows, TAOS_ROW *rows) {
     (*rows) = pResultInfo->row;
     (*numOfRows) = pResultInfo->numOfRows;
     return pRequest->code;
+  } else if (TD_RES_TMQ(res)) {
+    SReqResultInfo *pResultInfo = tmqGetNextResInfo(res);
+    if (pResultInfo == NULL) return -1;
+
+    pResultInfo->current = pResultInfo->numOfRows;
+    (*rows) = pResultInfo->row;
+    (*numOfRows) = pResultInfo->numOfRows;
+    return 0;
+  } else {
+    ASSERT(0);
+    return -1;
+  }
 }
 
-int taos_fetch_raw_block(TAOS_RES *res, int* numOfRows, void** pData) {
-  SRequestObj *pRequest = (SRequestObj *)res;
-  if (pRequest == NULL) {
+int taos_fetch_raw_block(TAOS_RES *res, int *numOfRows, void **pData) {
+  if (res == NULL) {
     return 0;
   }
+  if (TD_RES_QUERY(res)) {
+    SRequestObj *pRequest = (SRequestObj *)res;
 
     if (pRequest->type == TSDB_SQL_RETRIEVE_EMPTY_RESULT || pRequest->type == TSDB_SQL_INSERT ||
         pRequest->code != TSDB_CODE_SUCCESS || taos_num_fields(res) == 0) {
@@ -432,28 +471,42 @@ int taos_fetch_raw_block(TAOS_RES *res, int* numOfRows, void** pData) {
 
     pResultInfo->current = pResultInfo->numOfRows;
     (*numOfRows) = pResultInfo->numOfRows;
-  (*pData) = (void*) pResultInfo->pData;
+    (*pData) = (void *)pResultInfo->pData;
 
     return 0;
+
+  } else if (TD_RES_TMQ(res)) {
+    SReqResultInfo *pResultInfo = tmqGetNextResInfo(res);
+    if (pResultInfo == NULL) return -1;
+
+    pResultInfo->current = pResultInfo->numOfRows;
+    (*numOfRows) = pResultInfo->numOfRows;
+    (*pData) = (void *)pResultInfo->pData;
+    return 0;
+
+  } else {
+    ASSERT(0);
+    return -1;
+  }
 }
 
 int *taos_get_column_data_offset(TAOS_RES *res, int columnIndex) {
-  SRequestObj *pRequest = (SRequestObj *)res;
-  if (pRequest == NULL) {
+  if (res == NULL) {
     return 0;
   }
 
-  int32_t numOfFields = taos_num_fields(pRequest);
+  int32_t numOfFields = taos_num_fields(res);
   if (columnIndex < 0 || columnIndex >= numOfFields || numOfFields == 0) {
     return 0;
   }
 
-  TAOS_FIELD* pField = &pRequest->body.resInfo.userFields[columnIndex];
+  SReqResultInfo *pResInfo = tscGetCurResInfo(res);
+  TAOS_FIELD     *pField = &pResInfo->userFields[columnIndex];
   if (!IS_VAR_DATA_TYPE(pField->type)) {
     return 0;
   }
 
-  return pRequest->body.resInfo.pCol[columnIndex].offset;
+  return pResInfo->pCol[columnIndex].offset;
 }
 
 int taos_validate_sql(TAOS *taos, const char *sql) { return true; }
@@ -483,7 +536,8 @@ void taos_fetch_rows_a(TAOS_RES *res, __taos_async_fn_t fp, void *param) {
   // TODO
 }
 
-TAOS_SUB *taos_subscribe(TAOS *taos, int restart, const char* topic, const char *sql, TAOS_SUBSCRIBE_CALLBACK fp, void *param, int interval) {
+TAOS_SUB *taos_subscribe(TAOS *taos, int restart, const char *topic, const char *sql, TAOS_SUBSCRIBE_CALLBACK fp,
+                         void *param, int interval) {
   // TODO
   return NULL;
 }
@@ -562,7 +616,7 @@ int taos_stmt_num_params(TAOS_STMT *stmt, int *nums) {
   return -1;
 }
 
-int taos_stmt_add_batch(TAOS_STMT* stmt) {
+int taos_stmt_add_batch(TAOS_STMT *stmt) {
   // TODO
   return -1;
 }
@@ -572,7 +626,7 @@ TAOS_RES *taos_stmt_use_result(TAOS_STMT *stmt) {
   return NULL;
 }
 
-int taos_stmt_bind_param_batch(TAOS_STMT* stmt, TAOS_MULTI_BIND* bind) {
+int taos_stmt_bind_param_batch(TAOS_STMT *stmt, TAOS_MULTI_BIND *bind) {
   // TODO
   return -1;
 }

source/client/src/tmq.c

@@ -17,27 +17,13 @@
 #include "clientLog.h"
 #include "parser.h"
 #include "planner.h"
-#include "scheduler.h"
 #include "tdatablock.h"
 #include "tdef.h"
 #include "tglobal.h"
 #include "tmsgtype.h"
-#include "tpagedbuf.h"
 #include "tqueue.h"
 #include "tref.h"
 
-typedef struct {
-  int32_t curBlock;
-  int32_t curRow;
-  void**  uData;
-} SMqRowIter;
-
-struct tmq_message_t {
-  SMqPollRsp msg;
-  void*      vg;
-  SMqRowIter iter;
-};
-
 struct tmq_list_t {
   SArray container;
 };
@@ -1566,30 +1552,4 @@ const char* tmq_err2str(tmq_resp_err_t err) {
   return "fail";
 }
 
-TAOS_ROW tmq_get_row(tmq_message_t* message) {
-  SMqPollRsp* rsp = &message->msg;
-  while (1) {
-    if (message->iter.curBlock < taosArrayGetSize(rsp->pBlockData)) {
-      SSDataBlock* pBlock = taosArrayGet(rsp->pBlockData, message->iter.curBlock);
-      if (message->iter.curRow < pBlock->info.rows) {
-        for (int i = 0; i < pBlock->info.numOfCols; i++) {
-          SColumnInfoData* pData = taosArrayGet(pBlock->pDataBlock, i);
-          if (colDataIsNull_s(pData, message->iter.curRow))
-            message->iter.uData[i] = NULL;
-          else {
-            message->iter.uData[i] = colDataGetData(pData, message->iter.curRow);
-          }
-        }
-        message->iter.curRow++;
-        return message->iter.uData;
-      } else {
-        message->iter.curBlock++;
-        message->iter.curRow = 0;
-        continue;
-      }
-    }
-    return NULL;
-  }
-}
-
 char* tmq_get_topic_name(tmq_message_t* message) { return "not implemented yet"; }

source/libs/executor/src/dataDispatcher.c

@@ -92,9 +92,7 @@ static bool allocBuf(SDataDispatchHandle* pDispatcher, const SInputData* pInput,
     return false;
   }
 
-  // NOTE: there are four bytes of an integer more than the required buffer space.
-  // struct size + data payload + length for each column + bitmap length
-  pBuf->allocSize = sizeof(SRetrieveTableRsp) + blockDataGetSerialMetaSize(pInput->pData) + blockDataGetSize(pInput->pData);
+  pBuf->allocSize = sizeof(SRetrieveTableRsp) + blockEstimateEncodeSize(pInput->pData);
 
   pBuf->pData = taosMemoryMalloc(pBuf->allocSize);
   if (pBuf->pData == NULL) {