diff --git a/source/libs/planner/src/planSpliter.c b/source/libs/planner/src/planSpliter.c
index 45fc45336c6f59d2ed12f024caf697e143f4fecd..8586234b7e28342b0e51890306a37c4b22e56487 100644
--- a/source/libs/planner/src/planSpliter.c
+++ b/source/libs/planner/src/planSpliter.c
@@ -977,12 +977,24 @@ static int32_t stbSplSplitJoinNode(SSplitContext* pCxt, SStableSplitInfo* pInfo)
   return code;
 }
 
+static int32_t stbSplCreateMergeKeysForPartitionNode(SLogicNode* pPart, SNodeList** pMergeKeys) {
+  SNode* pPrimaryKey =
+      nodesCloneNode(stbSplFindPrimaryKeyFromScan((SScanLogicNode*)nodesListGetNode(pPart->pChildren, 0)));
+  if (NULL == pPrimaryKey) {
+    return TSDB_CODE_OUT_OF_MEMORY;
+  }
+  int32_t code = nodesListAppend(pPart->pTargets, pPrimaryKey);
+  if (TSDB_CODE_SUCCESS == code) {
+    code = stbSplCreateMergeKeysByPrimaryKey(pPrimaryKey, pMergeKeys);
+  }
+  return code;
+}
+
 static int32_t stbSplSplitPartitionNode(SSplitContext* pCxt, SStableSplitInfo* pInfo) {
   int32_t    code = TSDB_CODE_SUCCESS;
   SNodeList* pMergeKeys = NULL;
   if (pInfo->pSplitNode->requireDataOrder >= DATA_ORDER_LEVEL_IN_GROUP) {
-    code = stbSplCreateMergeKeysByPrimaryKey(
-        stbSplFindPrimaryKeyFromScan((SScanLogicNode*)nodesListGetNode(pInfo->pSplitNode->pChildren, 0)), &pMergeKeys);
+    code = stbSplCreateMergeKeysForPartitionNode(pInfo->pSplitNode, &pMergeKeys);
   }
   if (TSDB_CODE_SUCCESS == code) {
     code = stbSplCreateMergeNode(pCxt, pInfo->pSubplan, pInfo->pSplitNode, pMergeKeys, pInfo->pSplitNode, true);
diff --git a/source/libs/planner/test/planPartByTest.cpp b/source/libs/planner/test/planPartByTest.cpp
index 256960a15ef70c6c5484b20a7022ff2e8f79abd9..48ab1e1ac2cc5cf9f8b836b9d7940df047019f6e 100644
--- a/source/libs/planner/test/planPartByTest.cpp
+++ b/source/libs/planner/test/planPartByTest.cpp
@@ -67,6 +67,8 @@ TEST_F(PlanPartitionByTest, withTimeLineFunc) {
   useDb("root", "test");
 
   run("SELECT TWA(c1) FROM st1 PARTITION BY c1");
+
+  run("SELECT MAVG(c1, 2) FROM st1 PARTITION BY c1");
 }
 
 TEST_F(PlanPartitionByTest, withSlimit) {