From 1e8aaa273e89bdd6c1d47cdd8f82fe88d69f2302 Mon Sep 17 00:00:00 2001
From: kailixu <klxu@taosdata.com>
Date: Wed, 3 Apr 2024 16:32:59 +0800
Subject: [PATCH] chore: more logic for num of pks

---
 source/libs/nodes/src/nodesCloneFuncs.c | 1 +
 source/libs/nodes/src/nodesCodeFuncs.c  | 9 ++++++++-
 source/libs/nodes/src/nodesMsgFuncs.c   | 6 ++++++
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/source/libs/nodes/src/nodesCloneFuncs.c b/source/libs/nodes/src/nodesCloneFuncs.c
index 3e4b3cad31..e07fe07252 100644
--- a/source/libs/nodes/src/nodesCloneFuncs.c
+++ b/source/libs/nodes/src/nodesCloneFuncs.c
@@ -121,6 +121,7 @@ static int32_t columnNodeCopy(const SColumnNode* pSrc, SColumnNode* pDst) {
   COPY_SCALAR_FIELD(slotId);
   COPY_SCALAR_FIELD(tableHasPk);
   COPY_SCALAR_FIELD(isPk);
+  COPY_SCALAR_FIELD(numOfPKs);
   return TSDB_CODE_SUCCESS;
 }
 
diff --git a/source/libs/nodes/src/nodesCodeFuncs.c b/source/libs/nodes/src/nodesCodeFuncs.c
index f000d87fdc..cdf23e512e 100644
--- a/source/libs/nodes/src/nodesCodeFuncs.c
+++ b/source/libs/nodes/src/nodesCodeFuncs.c
@@ -3618,6 +3618,7 @@ static const char* jkColumnDataBlockId = "DataBlockId";
 static const char* jkColumnSlotId = "SlotId";
 static const char* jkColumnTableHasPk = "TableHasPk";
 static const char* jkColumnIsPk = "IsPk";
+static const char* jkColumnNumOfPKs = "NumOfPKs";
 
 static int32_t columnNodeToJson(const void* pObj, SJson* pJson) {
   const SColumnNode* pNode = (const SColumnNode*)pObj;
@@ -3662,6 +3663,9 @@ static int32_t columnNodeToJson(const void* pObj, SJson* pJson) {
   if (TSDB_CODE_SUCCESS == code) {
     code = tjsonAddBoolToObject(pJson, jkColumnIsPk, pNode->isPk);
   }  
+  if (TSDB_CODE_SUCCESS == code) {
+    code = tjsonAddIntegerToObject(pJson, jkColumnNumOfPKs, pNode->numOfPKs);
+  }
   return code;
 }
 
@@ -3707,7 +3711,10 @@ static int32_t jsonToColumnNode(const SJson* pJson, void* pObj) {
   }
   if (TSDB_CODE_SUCCESS == code) {
     code = tjsonGetBoolValue(pJson, jkColumnIsPk, &pNode->isPk);
-  }      
+  }    
+  if (TSDB_CODE_SUCCESS == code) {
+    code = tjsonGetSmallIntValue(pJson, jkColumnNumOfPKs, &pNode->numOfPKs);
+  }  
   return code;
 }
 
diff --git a/source/libs/nodes/src/nodesMsgFuncs.c b/source/libs/nodes/src/nodesMsgFuncs.c
index 0639a161bd..2cb2995a1d 100644
--- a/source/libs/nodes/src/nodesMsgFuncs.c
+++ b/source/libs/nodes/src/nodesMsgFuncs.c
@@ -731,6 +731,9 @@ static int32_t columnNodeInlineToMsg(const void* pObj, STlvEncoder* pEncoder) {
   if (TSDB_CODE_SUCCESS == code) {
     code = tlvEncodeValueBool(pEncoder, pNode->isPk);
   }  
+  if (TSDB_CODE_SUCCESS == code) {
+    code = tlvEncodeValueI16(pEncoder, pNode->numOfPKs);
+  }  
   return code;
 }
 
@@ -778,6 +781,9 @@ static int32_t msgToColumnNodeInline(STlvDecoder* pDecoder, void* pObj) {
   if (TSDB_CODE_SUCCESS == code) {
     code = tlvDecodeValueBool(pDecoder, &pNode->isPk);
   }  
+  if (TSDB_CODE_SUCCESS == code) {
+    code = tlvDecodeValueI16(pDecoder, &pNode->numOfPKs);
+  }  
   return code;
 }