Extend a new AM method to do acquire sample rows.

src/backend/access/aocs/aocsam_handler.c

@@ -2447,6 +2447,7 @@ static TableAmRoutine ao_column_methods = {
 	.scan_bitmap_next_tuple = aoco_scan_bitmap_next_tuple,
 	.scan_sample_next_block = aoco_scan_sample_next_block,
 	.scan_sample_next_tuple = aoco_scan_sample_next_tuple,
+	.acquire_sample_rows = acquire_sample_rows,
 
 	.amoptions = ao_amoptions,
 	.swap_relation_files = aoco_swap_relation_files,

src/backend/access/appendonly/appendonlyam_handler.c

@@ -2391,6 +2391,7 @@ static const TableAmRoutine ao_row_methods = {
 	.scan_bitmap_next_tuple = appendonly_scan_bitmap_next_tuple,
 	.scan_sample_next_block = appendonly_scan_sample_next_block,
 	.scan_sample_next_tuple = appendonly_scan_sample_next_tuple,
+	.acquire_sample_rows = acquire_sample_rows,
 
 	.amoptions = ao_amoptions,
 	.swap_relation_files = appendonly_swap_relation_files,

src/backend/access/heap/heapam_handler.c

@@ -2642,6 +2642,7 @@ static const TableAmRoutine heapam_methods = {
 	.scan_bitmap_next_tuple = heapam_scan_bitmap_next_tuple,
 	.scan_sample_next_block = heapam_scan_sample_next_block,
 	.scan_sample_next_tuple = heapam_scan_sample_next_tuple,
+	.acquire_sample_rows = acquire_sample_rows,
 
 	.amoptions	= heapam_amoptions,
 

src/backend/commands/analyze.c

@@ -347,7 +347,15 @@ analyze_rel_internal(Oid relid, RangeVar *relation,
 		onerel->rd_rel->relkind == RELKIND_MATVIEW)
 	{
 		/* Regular table, so we'll use the regular row acquisition function */
-		acquirefunc = acquire_sample_rows;
+		if (onerel->rd_tableam)
+			acquirefunc = onerel->rd_tableam->acquire_sample_rows;
+
+		/*
+		 * If the TableAmRoutine's acquire_sample_rows if NULL, we use
+		 * acquire_sample_rows as default.
+		 */
+		if (acquirefunc == NULL)
+			acquirefunc = acquire_sample_rows;
 
 		/* Also get regular table's size */
 		relpages = AcquireNumberOfBlocks(onerel);
@@ -1908,7 +1916,7 @@ acquire_inherited_sample_rows(Relation onerel, int elevel,
 	 * Like in acquire_sample_rows(), if we're in the QD, fetch the sample
 	 * from segments.
 	 */
-	if (Gp_role == GP_ROLE_DISPATCH)
+	if (Gp_role == GP_ROLE_DISPATCH && ENABLE_DISPATCH())
 	{
 		return acquire_sample_rows_dispatcher(onerel,
 											  true, /* inherited stats */
@@ -1979,7 +1987,16 @@ acquire_inherited_sample_rows(Relation onerel, int elevel,
 			childrel->rd_rel->relkind == RELKIND_MATVIEW)
 		{
 			/* Regular table, so use the regular row acquisition function */
-			acquirefunc = acquire_sample_rows;
+			if (childrel->rd_tableam)
+				acquirefunc = childrel->rd_tableam->acquire_sample_rows;
+
+			/*
+			 * If the TableAmRoutine's acquire_sample_rows if NULL, we use
+			 * acquire_sample_rows as default.
+			 */
+			if (acquirefunc == NULL)
+				acquirefunc = acquire_sample_rows;
+
 			relpages = AcquireNumberOfBlocks(childrel);
 		}
 		else if (childrel->rd_rel->relkind == RELKIND_FOREIGN_TABLE)

src/backend/commands/vacuum.c

@@ -614,7 +614,7 @@ vacuum(List *relations, VacuumParams *params,
 	}
 
 	if ((params->options & VACOPT_VACUUM) && !IsAutoVacuumWorkerProcess() &&
-		(Gp_role != GP_ROLE_EXECUTE || (params->options & VACOPT_UPDATE_DATFROZENXID)))
+		(Gp_role != GP_ROLE_EXECUTE || (params->options & VACOPT_UPDATE_DATFROZENXID)) && ENABLE_DISPATCH())
 	{
 		/*
 		 * Update pg_database.datfrozenxid, and truncate pg_xact if possible.

src/include/access/tableam.h

@@ -887,6 +887,15 @@ typedef struct TableAmRoutine
 										   struct SampleScanState *scanstate,
 										   TupleTableSlot *slot);
 
+	/*
+	 * Like scan_sample_next_block and scan_sample_next_tuple, this callback
+	 * is also used to sample tuple rows. As different storage maybe need to
+	 * use different acquire sample rows process, we extend an interface to
+	 * achieve this requirement. 
+	 */
+	int			(*acquire_sample_rows) (Relation onerel, int elevel,
+										HeapTuple *rows, int targrows,
+										double *totalrows, double *totaldeadrows);
 	/*
 	 * This callback is used to parse reloptions for relation/matview/toast.
 	 */

src/include/miscadmin.h

@@ -602,6 +602,16 @@ extern int	get_hash_mem(void);
  * HOLD_DISPATCH() to disable dispatch temporarily, and calls RESUME_DISPATCH()
  * to resume the dispatch states to allow the outer utility to normally
  * dispatch the commands to the QEs.
+ *
+ * In Greenplum, in some situations, maybe we should not allow to dispatch
+ * from QD to QE. To achieve that, we should use dispatch_nest_level to judge
+ * whether dispatch or not. When dispatch_nest_level > 0, we not dispatch,
+ * only when dispatch_nest_level is 0, we can dispatch. However, sometimes
+ * we also need to dispatch when dispatch_nest_level > 0, to do that, we should
+ * save dispatch_nest_level first and set dispatch_nest_level to 0. After
+ * finishing dispatch, we should restore dispatch_nest_level to original value.
+ * SAVE_DISPATCH_LEVEL is used to save dispatch_nest_level and set it to 0.
+ * RESTORE_DISPATCH_LEVEL is used to restore the original dispatch_nest_level.
  */
 extern int dispatch_nest_level;
 extern void GpRecoveryFromError(void);
@@ -615,6 +625,13 @@ extern void GpRecoveryFromError(void);
 	Assert(dispatch_nest_level >  0); \
 	dispatch_nest_level--; \
 } while(0)
+#define SAVE_DISPATCH_LEVEL(saved_dispatch_nest_level)	do { \
+	saved_dispatch_nest_level = dispatch_nest_level; \
+	dispatch_nest_level = 0; \
+} while(0)
+#define RESTORE_DISPATCH_LEVEL(saved_dispatch_nest_level)	do { \
+	dispatch_nest_level = saved_dispatch_nest_level; \
+} while(0)
 #define CLEAR_DISPATCH()	do { \
 	dispatch_nest_level = 0; \
 } while(0)

src/include/utils/elog.h

@@ -19,6 +19,7 @@
 #include "c.h"
 #include <sys/time.h>
 #include <setjmp.h>
+#include "miscadmin.h" /* dispatch_nest_level */
 
 /* Error level codes */
 #define DEBUG5		10			/* Debugging messages, in categories of
@@ -369,6 +370,7 @@ extern PGDLLIMPORT ErrorContextCallback *error_context_stack;
 		ErrorContextCallback *_save_context_stack = error_context_stack; \
 		sigjmp_buf _local_sigjmp_buf; \
 		bool _do_rethrow = false; \
+		int _saved_dispatch_nest_level pg_attribute_unused() = dispatch_nest_level; \
 		if (sigsetjmp(_local_sigjmp_buf, 0) == 0) \
 		{ \
 			PG_exception_stack = &_local_sigjmp_buf
@@ -377,6 +379,7 @@ extern PGDLLIMPORT ErrorContextCallback *error_context_stack;
 		} \
 		else \
 		{ \
+            dispatch_nest_level = _saved_dispatch_nest_level; \
 			PG_exception_stack = _save_exception_stack; \
 			error_context_stack = _save_context_stack