Merge pull request #5165 from citusdata/col/rewrite-mem-reset-bug

Use right mem cxt for read state when re-writing a columnar table

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

@@ -253,11 +253,15 @@ CreateColumnarScanMemoryContext(void)
  */
 static ColumnarReadState *
 init_columnar_read_state(Relation relation, TupleDesc tupdesc, Bitmapset *attr_needed,
-						 List *scanQual)
+						 List *scanQual, MemoryContext scanContext)
 {
+	MemoryContext oldContext = MemoryContextSwitchTo(scanContext);
+
 	List *neededColumnList = NeededColumnsList(tupdesc, attr_needed);
 	ColumnarReadState *readState = ColumnarBeginRead(relation, tupdesc, neededColumnList,
-													 scanQual);
+													 scanQual, scanContext);
+
+	MemoryContextSwitchTo(oldContext);
 
 	return readState;
 }
@@ -302,11 +306,10 @@ columnar_getnextslot(TableScanDesc sscan, ScanDirection direction, TupleTableSlo
 	 */
 	if (scan->cs_readState == NULL)
 	{
-		MemoryContext oldContext = MemoryContextSwitchTo(scan->scanContext);
 		scan->cs_readState =
 			init_columnar_read_state(scan->cs_base.rs_rd, slot->tts_tupleDescriptor,
-									 scan->attr_needed, scan->scanQual);
+									 scan->attr_needed, scan->scanQual,
-		MemoryContextSwitchTo(oldContext);
+									 scan->scanContext);
 	}
 
 	ExecClearTuple(slot);
@@ -490,8 +493,6 @@ columnar_index_fetch_tuple(struct IndexFetchTableData *sscan,
 	/* initialize read state for the first row */
 	if (scan->cs_readState == NULL)
 	{
-		MemoryContext oldContext = MemoryContextSwitchTo(scan->scanContext);
-
 		/* we need all columns */
 		int natts = columnarRelation->rd_att->natts;
 		Bitmapset *attr_needed = bms_add_range(NULL, 0, natts - 1);
@@ -501,8 +502,8 @@ columnar_index_fetch_tuple(struct IndexFetchTableData *sscan,
 
 		scan->cs_readState = init_columnar_read_state(columnarRelation,
 													  slot->tts_tupleDescriptor,
-													  attr_needed, scanQual);
+													  attr_needed, scanQual,
-		MemoryContextSwitchTo(oldContext);
+													  scan->scanContext);
 	}
 
 	uint64 rowNumber = tid_to_row_number(*tid);
@@ -795,10 +796,14 @@ columnar_relation_copy_for_cluster(Relation OldHeap, Relation NewHeap,
 	/* we need all columns */
 	int natts = OldHeap->rd_att->natts;
 	Bitmapset *attr_needed = bms_add_range(NULL, 0, natts - 1);
-	List *projectedColumnList = NeededColumnsList(sourceDesc, attr_needed);
+
-	ColumnarReadState *readState = ColumnarBeginRead(OldHeap, sourceDesc,
+	/* no quals for table rewrite */
-													 projectedColumnList,
+	List *scanQual = NIL;
-													 NULL);
+
+	MemoryContext scanContext = CreateColumnarScanMemoryContext();
+	ColumnarReadState *readState = init_columnar_read_state(OldHeap, sourceDesc,
+															attr_needed, scanQual,
+															scanContext);
 
 	Datum *values = palloc0(sourceDesc->natts * sizeof(Datum));
 	bool *nulls = palloc0(sourceDesc->natts * sizeof(bool));

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;