From 95697bcddbb33b01e5b8e8f014ad953e9df572ef Mon Sep 17 00:00:00 2001
From: Haojun Liao <hjliao@taosdata.com>
Date: Mon, 8 Apr 2024 18:04:53 +0800
Subject: [PATCH] fix(tsdb): prepare pk buf

---
 source/common/src/tdatablock.c | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/source/common/src/tdatablock.c b/source/common/src/tdatablock.c
index e748cce643..69b2a2e6a3 100644
--- a/source/common/src/tdatablock.c
+++ b/source/common/src/tdatablock.c
@@ -535,8 +535,8 @@ int32_t blockDataUpdatePkRange(SSDataBlock* pDataBlock, int32_t pkColumnIndex, b
 
   if (asc) {
     if (IS_NUMERIC_TYPE(pColInfoData->info.type)) {
-      pDataBlock->info.pks[0].val = *(int32_t*) skey;
-      pDataBlock->info.pks[1].val = *(int32_t*) ekey;
+      GET_TYPED_DATA(pDataBlock->info.pks[0].val, int64_t, pColInfoData->info.type, skey);
+      GET_TYPED_DATA(pDataBlock->info.pks[1].val, int64_t, pColInfoData->info.type, ekey);
     } else {  // todo refactor
       memcpy(pDataBlock->info.pks[0].pData, varDataVal(skey), varDataLen(skey));
       pDataBlock->info.pks[0].nData = varDataLen(skey);
@@ -546,8 +546,8 @@ int32_t blockDataUpdatePkRange(SSDataBlock* pDataBlock, int32_t pkColumnIndex, b
     }
   } else {
     if (IS_NUMERIC_TYPE(pColInfoData->info.type)) {
-      pDataBlock->info.pks[0].val = *(int32_t*) ekey;
-      pDataBlock->info.pks[1].val = *(int32_t*) skey;
+      GET_TYPED_DATA(pDataBlock->info.pks[0].val, int64_t, pColInfoData->info.type, ekey);
+      GET_TYPED_DATA(pDataBlock->info.pks[1].val, int64_t, pColInfoData->info.type, skey);
     } else {  // todo refactor
       memcpy(pDataBlock->info.pks[0].pData, varDataVal(ekey), varDataLen(ekey));
       pDataBlock->info.pks[0].nData = varDataLen(ekey);
@@ -1491,6 +1491,18 @@ SSDataBlock* createOneDataBlock(const SSDataBlock* pDataBlock, bool copyData) {
     blockDataAppendColInfo(pBlock, &colInfo);
   }
 
+  // prepare the pk buffer if necessary
+  if (IS_VAR_DATA_TYPE(pDataBlock->info.pks[0].type)) {
+    SValue* pVal = &pBlock->info.pks[0];
+
+    pVal->type = pDataBlock->info.pks[0].type;
+    pVal->pData = taosMemoryCalloc(1, pDataBlock->info.pks[0].nData);
+
+    pVal = &pBlock->info.pks[1];
+    pVal->type = pDataBlock->info.pks[1].type;
+    pVal->pData = taosMemoryCalloc(1, pDataBlock->info.pks[1].nData);
+  }
+
   if (copyData) {
     int32_t code = blockDataEnsureCapacity(pBlock, pDataBlock->info.rows);
     if (code != TSDB_CODE_SUCCESS) {