From b12e65d144f4dbfbdc369903ce0c0190d84db98f Mon Sep 17 00:00:00 2001
From: Haojun Liao <hjliao@taosdata.com>
Date: Thu, 21 Mar 2024 09:37:46 +0800
Subject: [PATCH] fix(tsdb): set pk column index.

---
 include/common/tmsg.h                      | 1 +
 source/dnode/vnode/src/tsdb/tsdbRead2.c    | 4 ++++
 source/dnode/vnode/src/tsdb/tsdbReadUtil.h | 1 +
 source/libs/executor/src/executil.c        | 4 +++-
 source/libs/executor/src/sysscanoperator.c | 1 +
 5 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/include/common/tmsg.h b/include/common/tmsg.h
index abaa1b1854..80714dba67 100644
--- a/include/common/tmsg.h
+++ b/include/common/tmsg.h
@@ -1037,6 +1037,7 @@ typedef struct {
   uint8_t scale;
   int32_t bytes;
   int8_t  type;
+  uint8_t pk;
 } SColumnInfo;
 
 typedef struct STimeWindow {
diff --git a/source/dnode/vnode/src/tsdb/tsdbRead2.c b/source/dnode/vnode/src/tsdb/tsdbRead2.c
index 4960937297..4973fb5d02 100644
--- a/source/dnode/vnode/src/tsdb/tsdbRead2.c
+++ b/source/dnode/vnode/src/tsdb/tsdbRead2.c
@@ -129,6 +129,10 @@ static int32_t setColumnIdSlotList(SBlockLoadSuppInfo* pSupInfo, SColumnInfo* pC
     } else {
       pSupInfo->buildBuf[i] = NULL;
     }
+
+    if (pCols[i].pk) {
+      pSupInfo->pkSlotId = pCols[i].slotId;
+    }
   }
 
   return TSDB_CODE_SUCCESS;
diff --git a/source/dnode/vnode/src/tsdb/tsdbReadUtil.h b/source/dnode/vnode/src/tsdb/tsdbReadUtil.h
index 2e61f6dff9..274a9fa20a 100644
--- a/source/dnode/vnode/src/tsdb/tsdbReadUtil.h
+++ b/source/dnode/vnode/src/tsdb/tsdbReadUtil.h
@@ -159,6 +159,7 @@ typedef struct SBlockLoadSuppInfo {
   int16_t*            slotId;
   int32_t             numOfCols;
   char**              buildBuf;  // build string tmp buffer, todo remove it later after all string format being updated.
+  int16_t             pkSlotId;
   bool                smaValid;  // the sma on all queried columns are activated
 } SBlockLoadSuppInfo;
 
diff --git a/source/libs/executor/src/executil.c b/source/libs/executor/src/executil.c
index 5e11cd1115..beaa0684ad 100644
--- a/source/libs/executor/src/executil.c
+++ b/source/libs/executor/src/executil.c
@@ -370,7 +370,8 @@ static EDealRes getColumn(SNode** pNode, void* pContext) {
     pSColumnNode->slotId = pData->index++;
     SColumnInfo cInfo = {.colId = pSColumnNode->colId,
                          .type = pSColumnNode->node.resType.type,
-                         .bytes = pSColumnNode->node.resType.bytes};
+                         .bytes = pSColumnNode->node.resType.bytes,
+                         .pk = pSColumnNode->isPk};
 #if TAG_FILTER_DEBUG
     qDebug("tagfilter build column info, slotId:%d, colId:%d, type:%d", pSColumnNode->slotId, cInfo.colId, cInfo.type);
 #endif
@@ -1763,6 +1764,7 @@ int32_t initQueryTableDataCond(SQueryTableDataCond* pCond, const STableScanPhysi
     pCond->colList[j].type = pColNode->node.resType.type;
     pCond->colList[j].bytes = pColNode->node.resType.bytes;
     pCond->colList[j].colId = pColNode->colId;
+    pCond->colList[j].pk = pColNode->isPk;
 
     pCond->pSlotList[j] = pNode->slotId;
     j += 1;
diff --git a/source/libs/executor/src/sysscanoperator.c b/source/libs/executor/src/sysscanoperator.c
index 0c421bf354..1f41a0c7b3 100644
--- a/source/libs/executor/src/sysscanoperator.c
+++ b/source/libs/executor/src/sysscanoperator.c
@@ -2276,6 +2276,7 @@ static int32_t initTableblockDistQueryCond(uint64_t uid, SQueryTableDataCond* pC
   pCond->colList->colId = 1;
   pCond->colList->type = TSDB_DATA_TYPE_TIMESTAMP;
   pCond->colList->bytes = sizeof(TSKEY);
+  pCond->colList->pk = 0;
 
   pCond->pSlotList[0] = 0;