From 31944516837e5e80b67655839af59a679d4e7690 Mon Sep 17 00:00:00 2001
From: Japin Li <japinli@hotmail.com>
Date: Wed, 20 Mar 2024 23:30:33 +0800
Subject: [PATCH] Fix memory leak when executing vectorized quals (#242)

* Fix memory leak when executing vectorized quals

Currently, a VectorColumn structure is created when executing vectorized
quals for each tuple on the ExecutorState memory context.  However, it
will be freed only until execution finishes.  This commit changes the
memory context to a tuple memory context.

* Fix memory leak in ReadStripeNextVector()

In ReadStripeNextVector(), the memory of columnValueOffset is allocated
in ExecutorState memory content, it will be freed until execution is
finished, so call pfree() to explicitly release the memory to avoid
memory growing up.
---
 .../backend/columnar/columnar_customscan.c    |  2 +-
 .../src/backend/columnar/columnar_reader.c    |  2 ++
 .../vectorization/columnar_vector_execution.c | 21 +++++++++++++------
 .../vectorization/columnar_vector_execution.h |  5 +++--
 4 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/columnar/src/backend/columnar/columnar_customscan.c b/columnar/src/backend/columnar/columnar_customscan.c
index a58dbd1..7ab06a0 100644
--- a/columnar/src/backend/columnar/columnar_customscan.c
+++ b/columnar/src/backend/columnar/columnar_customscan.c
@@ -2380,7 +2380,7 @@ CustomExecScan(ColumnarScanState *columnarScanState,
 				bool *resultQual = 
 					ExecuteVectorizedQual(slot,
 										columnarScanState->vectorization.constructedVectorizedQualList,
-										AND_EXPR);
+										AND_EXPR, econtext);
 
 				memcpy(vectorSlot->keep, resultQual, COLUMNAR_VECTOR_COLUMN_SIZE);
 			}
diff --git a/columnar/src/backend/columnar/columnar_reader.c b/columnar/src/backend/columnar/columnar_reader.c
index 30448d5..fd63f55 100644
--- a/columnar/src/backend/columnar/columnar_reader.c
+++ b/columnar/src/backend/columnar/columnar_reader.c
@@ -2171,9 +2171,11 @@ ReadStripeNextVector(StripeReadState *stripeReadState, Datum *columnValues,
 		else
 			stripeReadState->currentRow += stripeReadState->chunkGroupReadState->rowCount;
 
+		pfree(columnValueOffset);
 		return true;
 	}
 
+	pfree(columnValueOffset);
 	return false;
 }
 
diff --git a/columnar/src/backend/columnar/vectorization/columnar_vector_execution.c b/columnar/src/backend/columnar/vectorization/columnar_vector_execution.c
index b3666da..367ae09 100644
--- a/columnar/src/backend/columnar/vectorization/columnar_vector_execution.c
+++ b/columnar/src/backend/columnar/vectorization/columnar_vector_execution.c
@@ -390,16 +390,21 @@ vectorizedOr(bool *left, bool *right, int dimension)
 }
 
 static bool *
-executeVectorizedExpr(VectorQual *vectorQual)
+executeVectorizedExpr(VectorQual *vectorQual, ExprContext *econtext)
 {
+	MemoryContext oldContext;
+
+	oldContext = MemoryContextSwitchTo(econtext->ecxt_per_tuple_memory);
 	VectorColumn *res = 
 		(VectorColumn *) vectorQual->u.expr.fmgrInfo->fn_addr(vectorQual->u.expr.fcInfo);
+	MemoryContextSwitchTo(oldContext);
 
 	return (bool *) res->value;
 }
 
 bool * 
-ExecuteVectorizedQual(TupleTableSlot *slot, List *vectorizedQualList, BoolExprType boolType)
+ExecuteVectorizedQual(TupleTableSlot *slot, List *vectorizedQualList,
+			BoolExprType boolType, ExprContext *econtext)
 {
 	VectorTupleTableSlot *vectorSlot = (VectorTupleTableSlot *) slot;
 	ListCell *lc;
@@ -416,18 +421,22 @@ ExecuteVectorizedQual(TupleTableSlot *slot, List *vectorizedQualList, BoolExprTy
 		{
 			case VECTOR_QUAL_EXPR:
 			{
-				qualResult = executeVectorizedExpr(vectorQual);
+				qualResult = executeVectorizedExpr(vectorQual, econtext);
 				break;
 			}
 			case VECTOR_QUAL_BOOL_EXPR:
 			{
 				if (vectorQual->u.boolExpr.boolExprType == AND_EXPR)
 				{
-					qualResult = ExecuteVectorizedQual(slot, vectorQual->u.boolExpr.vectorQualExprList, AND_EXPR);
+					qualResult = ExecuteVectorizedQual(slot,
+										vectorQual->u.boolExpr.vectorQualExprList,
+										AND_EXPR, econtext);
 				}
 				else if (vectorQual->u.boolExpr.boolExprType == OR_EXPR)
 				{
-					qualResult = ExecuteVectorizedQual(slot, vectorQual->u.boolExpr.vectorQualExprList, OR_EXPR);
+					qualResult = ExecuteVectorizedQual(slot,
+										vectorQual->u.boolExpr.vectorQualExprList,
+										OR_EXPR, econtext);
 				}
 				break;
 			}
@@ -448,4 +457,4 @@ ExecuteVectorizedQual(TupleTableSlot *slot, List *vectorizedQualList, BoolExprTy
 	}
 
 	return result;
-}
\ No newline at end of file
+}
diff --git a/columnar/src/include/columnar/vectorization/columnar_vector_execution.h b/columnar/src/include/columnar/vectorization/columnar_vector_execution.h
index f943719..72ca484 100644
--- a/columnar/src/include/columnar/vectorization/columnar_vector_execution.h
+++ b/columnar/src/include/columnar/vectorization/columnar_vector_execution.h
@@ -22,6 +22,7 @@ extern List * CreateVectorizedExprList(List *exprList);
 extern List * ConstructVectorizedQualList(TupleTableSlot *slot, List *vectorizedQual);
 extern bool * ExecuteVectorizedQual(TupleTableSlot *slot,
 									List *vectorizedQualList,
-									BoolExprType boolType);
+									BoolExprType boolType,
+									ExprContext *econtext);
 
-#endif
\ No newline at end of file
+#endif