Use scan context for intermediate mem allocs too

src/backend/columnar/columnar_reader.c

@@ -81,6 +81,13 @@ struct ColumnarReadState
 
 	MemoryContext stripeReadContext;
 	int64 chunkGroupsFiltered;
+
+	/*
+	 * Memory context guaranteed to be not freed during scan so we can
+	 * safely use for any memory allocations regarding ColumnarReadState
+	 * itself.
+	 */
+	MemoryContext scanContext;
 };
 
 /* static function declarations */
@@ -159,7 +166,8 @@ static Datum ColumnDefaultValue(TupleConstr *tupleConstraints,
  */
 ColumnarReadState *
 ColumnarBeginRead(Relation relation, TupleDesc tupleDescriptor,
-				  List *projectedColumnList, List *whereClauseList)
+				  List *projectedColumnList, List *whereClauseList,
+				  MemoryContext scanContext)
 {
 	/*
 	 * We allocate all stripe specific data in the stripeReadContext, and reset
@@ -180,6 +188,7 @@ ColumnarBeginRead(Relation relation, TupleDesc tupleDescriptor,
 	readState->currentStripeMetadata = FindNextStripeByRowNumber(relation,
 																 COLUMNAR_INVALID_ROW_NUMBER,
 																 GetTransactionSnapshot());
+	readState->scanContext = scanContext;
 
 	return readState;
 }
@@ -432,11 +441,15 @@ HasUnreadStripe(ColumnarReadState *readState)
 void
 ColumnarRescan(ColumnarReadState *readState)
 {
+	MemoryContext oldContext = MemoryContextSwitchTo(readState->scanContext);
+
 	ColumnarResetRead(readState);
 	readState->currentStripeMetadata = FindNextStripeByRowNumber(readState->relation,
 																 COLUMNAR_INVALID_ROW_NUMBER,
 																 GetTransactionSnapshot());
 	readState->chunkGroupsFiltered = 0;
+
+	MemoryContextSwitchTo(oldContext);
 }
 
 
@@ -518,6 +531,8 @@ BeginStripeRead(StripeMetadata *stripeMetadata, Relation rel, TupleDesc tupleDes
 static void
 AdvanceStripeRead(ColumnarReadState *readState)
 {
+	MemoryContext oldContext = MemoryContextSwitchTo(readState->scanContext);
+
 	readState->chunkGroupsFiltered +=
 		readState->stripeReadState->chunkGroupsFiltered;
 
@@ -528,6 +543,8 @@ AdvanceStripeRead(ColumnarReadState *readState)
 																 GetTransactionSnapshot());
 	readState->stripeReadState = NULL;
 	MemoryContextReset(readState->stripeReadContext);
+
+	MemoryContextSwitchTo(oldContext);
 }
 
 

src/backend/columnar/columnar_tableam.c

@@ -259,7 +259,7 @@ init_columnar_read_state(Relation relation, TupleDesc tupdesc, Bitmapset *attr_n
 
 	List *neededColumnList = NeededColumnsList(tupdesc, attr_needed);
 	ColumnarReadState *readState = ColumnarBeginRead(relation, tupdesc, neededColumnList,
-													 scanQual);
+													 scanQual, scanContext);
 
 	MemoryContextSwitchTo(oldContext);
 

src/include/columnar/columnar.h

@@ -212,7 +212,8 @@ extern MemoryContext ColumnarWritePerTupleContext(ColumnarWriteState *state);
 extern ColumnarReadState * ColumnarBeginRead(Relation relation,
 											 TupleDesc tupleDescriptor,
 											 List *projectedColumnList,
-											 List *qualConditions);
+											 List *qualConditions,
+											 MemoryContext scanContext);
 extern bool ColumnarReadNextRow(ColumnarReadState *state, Datum *columnValues,
 								bool *columnNulls, uint64 *rowNumber);
 extern void ColumnarRescan(ColumnarReadState *readState);

src/test/regress/expected/columnar_alter.out

@@ -218,6 +218,21 @@ SELECT * FROM domain_test;
  1 | 2 | x
 (2 rows)
 
+-- similar to "add column c str_domain DEFAULT 'x'", both were getting
+-- stucked before fixing https://github.com/citusdata/citus/issues/5164
+BEGIN;
+  ALTER TABLE domain_test ADD COLUMN d INT DEFAULT random();
+ROLLBACK;
+BEGIN;
+  ALTER TABLE domain_test ADD COLUMN d SERIAL;
+  SELECT * FROM domain_test ORDER BY 1,2,3,4;
+ a | b | c | d
+---------------------------------------------------------------------
+ 1 | 2 | x | 1
+ 1 | 2 | x | 2
+(2 rows)
+
+ROLLBACK;
 set default_table_access_method TO 'columnar';
 CREATE TABLE has_volatile AS
 SELECT * FROM generate_series(1,10) id;

src/test/regress/sql/columnar_alter.sql

@@ -116,6 +116,16 @@ alter table domain_test add column c str_domain;
 alter table domain_test add column c str_domain DEFAULT 'x';
 SELECT * FROM domain_test;
 
+-- similar to "add column c str_domain DEFAULT 'x'", both were getting
+-- stucked before fixing https://github.com/citusdata/citus/issues/5164
+BEGIN;
+  ALTER TABLE domain_test ADD COLUMN d INT DEFAULT random();
+ROLLBACK;
+BEGIN;
+  ALTER TABLE domain_test ADD COLUMN d SERIAL;
+  SELECT * FROM domain_test ORDER BY 1,2,3,4;
+ROLLBACK;
+
 set default_table_access_method TO 'columnar';
 CREATE TABLE has_volatile AS
 SELECT * FROM generate_series(1,10) id;