This is an automated email from the ASF dual-hosted git repository.
yjhjstz pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/cloudberry.git
The following commit(s) were added to refs/heads/main by this push:
new 78153b4d53a Feature: dispatch T_CustomScanState in CBDB parallel
walkers (#1855)
78153b4d53a is described below
commit 78153b4d53a77c42cc9c6c858045688b6e735f7b
Author: roseduan <[email protected]>
AuthorDate: Sun Aug 9 17:34:13 2026 +0800
Feature: dispatch T_CustomScanState in CBDB parallel walkers (#1855)
Three GP-side walkers used during parallel setup —
EstimateGpParallelDSMEntrySize, InitializeGpParallelWorkers,
InitializeGpParallelDSMEntry — cased on other scan/join states but
skipped T_CustomScanState. A parallel_aware CustomScan therefore
silently under-sized its DSM and failed to attach in workers, with
no diagnostic.
Add T_CustomScanState arms that dispatch to
ExecCustomScan{Estimate,InitializeDSM,InitializeWorker}, gated on
parallel_aware to match the other node types in these switches.
Also handle T_CustomScanState in planstate_walk_kids. Upstream
planstate_tree_walker walks only css->custom_ps for a
CustomScanState; the CBDB walker previously fell through to
default (lefttree/righttree), skipping custom_ps children under
stateful walkers (cdbexplain_*, getMotionState, ...). The new
case walks custom_ps first, then lefttree/righttree if set, with
an Assert that a CustomScanState never populates both — otherwise
the child would be walked twice.
---
src/backend/executor/execParallel.c | 12 +
src/backend/executor/execProcnode.c | 22 ++
src/test/modules/Makefile | 2 +-
src/test/modules/parallel_customscan/Makefile | 29 ++
.../expected/parallel_customscan.out | 212 +++++++++++
.../parallel_customscan--1.0.sql | 15 +
.../parallel_customscan/parallel_customscan.c | 396 +++++++++++++++++++++
.../parallel_customscan.control | 5 +
.../sql/parallel_customscan.sql | 76 ++++
9 files changed, 768 insertions(+), 1 deletion(-)
diff --git a/src/backend/executor/execParallel.c
b/src/backend/executor/execParallel.c
index 16c4abe9b76..8d1a72cced7 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -1551,6 +1551,10 @@ EstimateGpParallelDSMEntrySize(PlanState *planstate,
ParallelContext *pctx)
case T_SortState:
ExecSortEstimate((SortState *) planstate, pctx);
break;
+ case T_CustomScanState:
+ if (planstate->plan->parallel_aware)
+ ExecCustomScanEstimate((CustomScanState *)
planstate, pctx);
+ break;
default:
break;
@@ -1604,6 +1608,10 @@ InitializeGpParallelWorkers(PlanState *planstate,
ParallelWorkerContext *pwcxt)
if (planstate->plan->parallel_aware)
ExecHashJoinInitializeWorker((HashJoinState *)
planstate, pwcxt);
break;
+ case T_CustomScanState:
+ if (planstate->plan->parallel_aware)
+ ExecCustomScanInitializeWorker((CustomScanState
*) planstate, pwcxt);
+ break;
default:
break;
}
@@ -1662,6 +1670,10 @@ InitializeGpParallelDSMEntry(PlanState *planstate,
ParallelContext *pctx)
case T_SortState:
ExecSortInitializeDSM((SortState *) planstate, pctx);
break;
+ case T_CustomScanState:
+ if (planstate->plan->parallel_aware)
+ ExecCustomScanInitializeDSM((CustomScanState *)
planstate, pctx);
+ break;
default:
break;
}
diff --git a/src/backend/executor/execProcnode.c
b/src/backend/executor/execProcnode.c
index 3f98f99f267..954ba935de1 100644
--- a/src/backend/executor/execProcnode.c
+++ b/src/backend/executor/execProcnode.c
@@ -1295,6 +1295,28 @@ planstate_walk_kids(PlanState *planstate,
Assert(!planstate->lefttree && !planstate->righttree);
break;
+ case T_CustomScanState:
+ {
+ CustomScanState *css = (CustomScanState *)
planstate;
+ ListCell *lc;
+
+ Assert(!(css->custom_ps != NIL &&
+ (planstate->lefttree ||
planstate->righttree)));
+
+ v = CdbVisit_Walk;
+ foreach(lc, css->custom_ps)
+ {
+ v =
planstate_walk_node_extended((PlanState *) lfirst(lc), walker, context, flags);
+ if (v != CdbVisit_Walk)
+ break;
+ }
+ if (v == CdbVisit_Walk && planstate->lefttree)
+ v =
planstate_walk_node_extended(planstate->lefttree, walker, context, flags);
+ if (v == CdbVisit_Walk && planstate->righttree)
+ v =
planstate_walk_node_extended(planstate->righttree, walker, context, flags);
+ break;
+ }
+
default:
/* Left subtree */
v = planstate_walk_node_extended(planstate->lefttree,
walker, context, flags);
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 44f9ac5fc7d..77af107b55e 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -43,7 +43,7 @@ SUBDIRS = \
# special check for DML on system relations in GPDB
# GPDB subdirs
-SUBDIRS += test_planner
+SUBDIRS += test_planner parallel_customscan
ifeq ($(with_ssl),openssl)
SUBDIRS += ssl_passphrase_callback
else
diff --git a/src/test/modules/parallel_customscan/Makefile
b/src/test/modules/parallel_customscan/Makefile
new file mode 100644
index 00000000000..3e4297d638a
--- /dev/null
+++ b/src/test/modules/parallel_customscan/Makefile
@@ -0,0 +1,29 @@
+# src/test/modules/parallel_customscan/Makefile
+
+MODULE_big = parallel_customscan
+OBJS = \
+ $(WIN32RES) \
+ parallel_customscan.o
+
+EXTENSION = parallel_customscan
+DATA = parallel_customscan--1.0.sql
+PGFILEDESC = "parallel_customscan - exercise parallel CustomScan dispatch"
+
+REGRESS = parallel_customscan
+
+# Run against an existing cluster (gpdemo). The cluster must have
+# 'parallel_customscan' in shared_preload_libraries so segment backends
+# and parallel workers have the CustomScan methods registered:
+# gpconfig -c shared_preload_libraries -v "'time_series,parallel_customscan'"
+# gpstop -ra
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/parallel_customscan
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git
a/src/test/modules/parallel_customscan/expected/parallel_customscan.out
b/src/test/modules/parallel_customscan/expected/parallel_customscan.out
new file mode 100644
index 00000000000..867b5c60e9a
--- /dev/null
+++ b/src/test/modules/parallel_customscan/expected/parallel_customscan.out
@@ -0,0 +1,212 @@
+-- start_matchsubs
+-- m/\(actual rows=[^)]*\)/
+-- s/\(actual rows=[^)]*\)/(actual rows=...)/
+-- end_matchsubs
+-- start_matchignore
+-- m/^\s*Buckets: \d+ Batches: \d+ Memory Usage: \d+kB/
+-- end_matchignore
+CREATE EXTENSION parallel_customscan;
+-- Set up data BEFORE enabling our hook so ANALYZE doesn't traverse it.
+CREATE TABLE pcs_t (a int)
+ WITH (parallel_workers = 2)
+ DISTRIBUTED BY (a);
+INSERT INTO pcs_t SELECT generate_series(1, 1000);
+ANALYZE pcs_t;
+SET optimizer = off;
+SET enable_parallel = on;
+SET enable_seqscan = off;
+SET min_parallel_table_scan_size = 0;
+SET parallel_setup_cost = 0;
+SET parallel_tuple_cost = 0;
+SET max_parallel_workers_per_gather = 2;
+-- (1) Enabling the hook replaces the Parallel Seq Scan with a Custom Scan.
+EXPLAIN (COSTS OFF) SELECT count(*) FROM pcs_t;
+ QUERY PLAN
+------------------------------------------------
+ Finalize Aggregate
+ -> Gather Motion 6:1 (slice1; segments: 6)
+ -> Partial Aggregate
+ -> Parallel Seq Scan on pcs_t
+ Optimizer: Postgres query optimizer
+(5 rows)
+
+SET parallel_customscan.enabled = on;
+EXPLAIN (COSTS OFF) SELECT count(*) FROM pcs_t;
+ QUERY PLAN
+-------------------------------------------------------------
+ Finalize Aggregate
+ -> Gather Motion 6:1 (slice1; segments: 6)
+ -> Partial Aggregate
+ -> Parallel Custom Scan (ParallelCustomScan)
+ -> Parallel Seq Scan on pcs_t
+ Optimizer: Postgres query optimizer
+(6 rows)
+
+-- (2) Correctness: results match only if every parallel worker runs our scan.
+SELECT count(*) FROM pcs_t;
+ count
+-------
+ 1000
+(1 row)
+
+SELECT sum(a) FROM pcs_t;
+ sum
+--------
+ 500500
+(1 row)
+
+-- (3) Scan-level qualifier and projection through the custom scan.
+EXPLAIN (COSTS OFF) SELECT a FROM pcs_t WHERE a > 990;
+ QUERY PLAN
+-------------------------------------------------
+ Gather Motion 6:1 (slice1; segments: 6)
+ -> Parallel Custom Scan (ParallelCustomScan)
+ -> Parallel Seq Scan on pcs_t
+ Filter: (a > 990)
+ Optimizer: Postgres query optimizer
+(5 rows)
+
+SELECT a FROM pcs_t WHERE a > 990 ORDER BY a;
+ a
+------
+ 991
+ 992
+ 993
+ 994
+ 995
+ 996
+ 997
+ 998
+ 999
+ 1000
+(10 rows)
+
+SELECT count(*) FROM pcs_t WHERE a % 2 = 0;
+ count
+-------
+ 500
+(1 row)
+
+-- (4) Join with both inputs scanned by the custom scan, plus a scan-level
qual.
+CREATE TABLE pcs_t2 (a int)
+ WITH (parallel_workers = 2)
+ DISTRIBUTED BY (a);
+INSERT INTO pcs_t2 SELECT generate_series(1, 500);
+ANALYZE pcs_t2;
+EXPLAIN (COSTS OFF)
+ SELECT count(*) FROM pcs_t x JOIN pcs_t2 y ON x.a = y.a WHERE x.a <= 100;
+ QUERY PLAN
+-------------------------------------------------------------------------
+ Finalize Aggregate
+ -> Gather Motion 6:1 (slice1; segments: 6)
+ -> Partial Aggregate
+ -> Parallel Hash Join
+ Hash Cond: (x.a = y.a)
+ -> Parallel Custom Scan (ParallelCustomScan)
+ -> Parallel Seq Scan on pcs_t x
+ Filter: (a <= 100)
+ -> Parallel Hash
+ -> Parallel Custom Scan (ParallelCustomScan)
+ -> Parallel Seq Scan on pcs_t2 y
+ Filter: (a <= 100)
+ Optimizer: Postgres query optimizer
+(13 rows)
+
+SELECT count(*) FROM pcs_t x JOIN pcs_t2 y ON x.a = y.a WHERE x.a <= 100;
+ count
+-------
+ 100
+(1 row)
+
+-- (5) Empty relation: the custom scan must handle an immediate end-of-scan.
+CREATE TABLE pcs_empty (a int)
+ WITH (parallel_workers = 2)
+ DISTRIBUTED BY (a);
+ANALYZE pcs_empty;
+EXPLAIN (COSTS OFF) SELECT count(*) FROM pcs_empty;
+ QUERY PLAN
+------------------------------------------------
+ Aggregate
+ -> Gather Motion 3:1 (slice1; segments: 3)
+ -> Custom Scan (ParallelCustomScan)
+ -> Seq Scan on pcs_empty
+ Optimizer: Postgres query optimizer
+(5 rows)
+
+SELECT count(*) FROM pcs_empty;
+ count
+-------
+ 0
+(1 row)
+
+SELECT * FROM pcs_empty;
+ a
+---
+(0 rows)
+
+-- (6) Serial path: with no workers, a non-parallel Custom Scan is used.
+SET max_parallel_workers_per_gather = 0;
+EXPLAIN (COSTS OFF) SELECT count(*) FROM pcs_t;
+ QUERY PLAN
+----------------------------------------------------
+ Finalize Aggregate
+ -> Gather Motion 3:1 (slice1; segments: 3)
+ -> Partial Aggregate
+ -> Custom Scan (ParallelCustomScan)
+ -> Seq Scan on pcs_t
+ Optimizer: Postgres query optimizer
+(6 rows)
+
+SELECT count(*) FROM pcs_t;
+ count
+-------
+ 1000
+(1 row)
+
+SELECT a FROM pcs_t WHERE a > 995 ORDER BY a;
+ a
+------
+ 996
+ 997
+ 998
+ 999
+ 1000
+(5 rows)
+
+SET max_parallel_workers_per_gather = 2;
+-- (7) EXPLAIN ANALYZE: exercises planstate_walk_kids' custom_ps recursion.
+EXPLAIN (ANALYZE, TIMING OFF, COSTS OFF, SUMMARY OFF)
+ SELECT count(*) FROM pcs_t;
+ QUERY PLAN
+---------------------------------------------------------------------------------------
+ Finalize Aggregate (actual rows=1 loops=1)
+ -> Gather Motion 6:1 (slice1; segments: 6) (actual rows=6 loops=1)
+ -> Partial Aggregate (actual rows=1 loops=1)
+ -> Parallel Custom Scan (ParallelCustomScan) (actual rows=340
loops=1)
+ -> Parallel Seq Scan on pcs_t (actual rows=340 loops=1)
+ Optimizer: Postgres query optimizer
+(6 rows)
+
+EXPLAIN (ANALYZE, TIMING OFF, COSTS OFF, SUMMARY OFF)
+ SELECT count(*) FROM pcs_t x JOIN pcs_t2 y ON x.a = y.a;
+ QUERY PLAN
+-------------------------------------------------------------------------------------------------
+ Finalize Aggregate (actual rows=1 loops=1)
+ -> Gather Motion 6:1 (slice1; segments: 6) (actual rows=6 loops=1)
+ -> Partial Aggregate (actual rows=1 loops=1)
+ -> Parallel Hash Join (actual rows=0 loops=1)
+ Hash Cond: (x.a = y.a)
+ -> Parallel Custom Scan (ParallelCustomScan) (actual
rows=340 loops=1)
+ -> Parallel Seq Scan on pcs_t x (actual rows=340
loops=1)
+ -> Parallel Hash (actual rows=0 loops=1)
+ Buckets: 524288 Batches: 1 Memory Usage: 4128kB
+ -> Parallel Custom Scan (ParallelCustomScan)
(actual rows=0 loops=1)
+ -> Parallel Seq Scan on pcs_t2 y (actual
rows=0 loops=1)
+ Optimizer: Postgres query optimizer
+(12 rows)
+
+-- cleanup
+DROP TABLE pcs_t2;
+DROP TABLE pcs_empty;
+DROP TABLE pcs_t;
+DROP EXTENSION parallel_customscan;
diff --git a/src/test/modules/parallel_customscan/parallel_customscan--1.0.sql
b/src/test/modules/parallel_customscan/parallel_customscan--1.0.sql
new file mode 100644
index 00000000000..1304f8af9a0
--- /dev/null
+++ b/src/test/modules/parallel_customscan/parallel_customscan--1.0.sql
@@ -0,0 +1,15 @@
+/* src/test/modules/parallel_customscan/parallel_customscan--1.0.sql */
+
+-- Complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION parallel_customscan" to load this file. \quit
+
+CREATE FUNCTION pcs_get_hook_calls(
+ OUT estimate_calls bigint,
+ OUT init_dsm_calls bigint,
+ OUT reinit_dsm_calls bigint,
+ OUT init_worker_calls bigint,
+ OUT shutdown_calls bigint
+)
+RETURNS record
+AS 'MODULE_PATHNAME', 'pcs_get_hook_calls'
+LANGUAGE C STRICT;
diff --git a/src/test/modules/parallel_customscan/parallel_customscan.c
b/src/test/modules/parallel_customscan/parallel_customscan.c
new file mode 100644
index 00000000000..9f5ae66b8b8
--- /dev/null
+++ b/src/test/modules/parallel_customscan/parallel_customscan.c
@@ -0,0 +1,396 @@
+#include "postgres.h"
+
+#include "access/relscan.h"
+#include "access/tableam.h"
+#include "catalog/pg_type.h"
+#include "executor/executor.h"
+#include "fmgr.h"
+#include "funcapi.h"
+#include "miscadmin.h"
+#include "nodes/extensible.h"
+#include "nodes/makefuncs.h"
+#include "optimizer/optimizer.h"
+#include "cdb/cdbpathlocus.h"
+#include "optimizer/pathnode.h"
+#include "optimizer/paths.h"
+#include "optimizer/restrictinfo.h"
+#include "port/atomics.h"
+#include "storage/shm_toc.h"
+#include "utils/builtins.h"
+#include "utils/guc.h"
+#include "utils/rel.h"
+
+PG_MODULE_MAGIC;
+
+void _PG_init(void);
+void _PG_fini(void);
+
+PG_FUNCTION_INFO_V1(pcs_get_hook_calls);
+
+/* GUC */
+static bool pcs_enabled = false;
+
+/*
+ * Per-process hook-call counters. Reset in EstimateDSMCustomScan (which is
+ * called once per parallel-mode invocation in the leader). The cross-worker
+ * InitializeWorkerCustomScan count is aggregated into the DSM atomic and
+ * harvested back into the leader-local counter in ShutdownCustomScan.
+ */
+static int64 pcs_n_estimate = 0;
+static int64 pcs_n_init_dsm = 0;
+static int64 pcs_n_reinit_dsm = 0;
+static int64 pcs_n_init_worker = 0;
+static int64 pcs_n_shutdown = 0;
+
+/*
+ * DSM header. The child SeqScan owns the parallel table scan descriptor, so
+ * the wrapper only needs a tiny shared area to count InitializeWorker calls
+ * across workers.
+ */
+typedef struct PcsDSM
+{
+ pg_atomic_uint32 init_worker_calls;
+} PcsDSM;
+
+typedef struct PcsState
+{
+ CustomScanState csstate;
+ PcsDSM *dsm; /* set in InitDSM /
InitWorker */
+} PcsState;
+
+static set_rel_pathlist_hook_type prev_pathlist_hook = NULL;
+
+static Plan *pcs_plan_path(PlannerInfo *root, RelOptInfo *rel,
+ CustomPath *best_path, List
*tlist,
+ List *clauses, List
*custom_plans);
+static Node *pcs_create_state(CustomScan *cscan);
+static void pcs_begin(CustomScanState *node, EState *estate, int eflags);
+static TupleTableSlot *pcs_exec(CustomScanState *node);
+static void pcs_end(CustomScanState *node);
+static void pcs_rescan(CustomScanState *node);
+static Size pcs_estimate_dsm(CustomScanState *node, ParallelContext *pcxt);
+static void pcs_init_dsm(CustomScanState *node, ParallelContext *pcxt,
+ void *coord);
+static void pcs_reinit_dsm(CustomScanState *node, ParallelContext *pcxt,
+ void *coord);
+static void pcs_init_worker(CustomScanState *node, shm_toc *toc, void *coord);
+static void pcs_shutdown(CustomScanState *node);
+
+static const CustomPathMethods pcs_path_methods =
+{
+ .CustomName = "ParallelCustomScan",
+ .PlanCustomPath = pcs_plan_path,
+};
+
+static const CustomScanMethods pcs_scan_methods =
+{
+ .CustomName = "ParallelCustomScan",
+ .CreateCustomScanState = pcs_create_state,
+};
+
+static const CustomExecMethods pcs_exec_methods =
+{
+ .CustomName = "ParallelCustomScan",
+ .BeginCustomScan = pcs_begin,
+ .ExecCustomScan = pcs_exec,
+ .EndCustomScan = pcs_end,
+ .ReScanCustomScan = pcs_rescan,
+ .EstimateDSMCustomScan = pcs_estimate_dsm,
+ .InitializeDSMCustomScan = pcs_init_dsm,
+ .ReInitializeDSMCustomScan = pcs_reinit_dsm,
+ .InitializeWorkerCustomScan = pcs_init_worker,
+ .ShutdownCustomScan = pcs_shutdown,
+};
+
+static Plan *
+pcs_plan_path(PlannerInfo *root, RelOptInfo *rel, CustomPath *best_path,
+ List *tlist, List *clauses, List *custom_plans)
+{
+ CustomScan *cs = makeNode(CustomScan);
+
+ cs->scan.plan.targetlist = tlist;
+ /*
+ * The child SeqScan does the filtering (create_scan_plan attaches the
+ * base restrictions to it), so the wrapper itself carries no qual.
+ */
+ cs->scan.plan.qual = NIL;
+ cs->scan.plan.parallel_aware = best_path->path.parallel_aware;
+ cs->scan.plan.parallel_safe = best_path->path.parallel_safe;
+ /*
+ * The wrapper delegates scanning to its child, so it is not itself a
+ * base-relation scan: scanrelid = 0. custom_scan_tlist must then
+ * describe the scan tuple we hand upward; we mirror the child's output
+ * targetlist, and set_customscan_references() rewrites our own
+ * targetlist to reference it via INDEX_VAR. This keeps projection
+ * correct for multi-column relations -- e.g. the catalog scans
+ * (pg_class) that ANALYZE issues internally, which a fixed
+ * base-relation descriptor would mis-deform.
+ */
+ cs->scan.scanrelid = 0;
+ cs->flags = best_path->flags;
+ cs->custom_plans = custom_plans;
+ cs->custom_exprs = NIL;
+ cs->custom_private = NIL;
+ cs->custom_scan_tlist =
+ copyObject(((Plan *) linitial(custom_plans))->targetlist);
+ cs->methods = &pcs_scan_methods;
+
+ return (Plan *) cs;
+}
+
+static Node *
+pcs_create_state(CustomScan *cscan)
+{
+ PcsState *st = (PcsState *) newNode(sizeof(PcsState),
T_CustomScanState);
+
+ st->csstate.methods = &pcs_exec_methods;
+ return (Node *) st;
+}
+
+static void
+pcs_begin(CustomScanState *node, EState *estate, int eflags)
+{
+ CustomScan *cscan = (CustomScan *) node->ss.ps.plan;
+ Plan *childplan = (Plan *) linitial(cscan->custom_plans);
+
+ /*
+ * ExecInitCustomScan has already opened the scan relation and set up
our
+ * scan/result slots and projection. All we add is the child plan
state,
+ * which becomes our sole custom_ps entry. That child performs the
actual
+ * (parallel) heap scan and is the node the MPP planstate walkers
recurse
+ * into via the T_CustomScanState arm.
+ */
+ node->custom_ps = list_make1(ExecInitNode(childplan, estate, eflags));
+}
+
+static TupleTableSlot *
+pcs_child_next(CustomScanState *node)
+{
+ PlanState *child = (PlanState *) linitial(node->custom_ps);
+ TupleTableSlot *childslot = ExecProcNode(child);
+
+ if (TupIsNull(childslot))
+ return NULL;
+
+ /*
+ * Move the child's tuple into our scan slot. That slot's descriptor
was
+ * built from custom_scan_tlist, which is an exact copy of the child's
+ * output targetlist, so the columns line up positionally and ExecScan's
+ * projection (compiled against the same slot) reads correct values.
+ */
+ ExecCopySlot(node->ss.ss_ScanTupleSlot, childslot);
+ return node->ss.ss_ScanTupleSlot;
+}
+
+static bool
+pcs_recheck(CustomScanState *node, TupleTableSlot *slot)
+{
+ return true;
+}
+
+static TupleTableSlot *
+pcs_exec(CustomScanState *node)
+{
+ return ExecScan(&node->ss,
+ (ExecScanAccessMtd) pcs_child_next,
+ (ExecScanRecheckMtd) pcs_recheck);
+}
+
+static void
+pcs_end(CustomScanState *node)
+{
+ ExecEndNode((PlanState *) linitial(node->custom_ps));
+}
+
+static void
+pcs_rescan(CustomScanState *node)
+{
+ ExecReScan((PlanState *) linitial(node->custom_ps));
+}
+
+static Size
+pcs_estimate_dsm(CustomScanState *node, ParallelContext *pcxt)
+{
+ /* Reset per-process counters at the start of a parallel run. */
+ pcs_n_estimate++;
+ pcs_n_init_dsm = 0;
+ pcs_n_reinit_dsm = 0;
+ pcs_n_init_worker = 0;
+ pcs_n_shutdown = 0;
+
+ return sizeof(PcsDSM);
+}
+
+static void
+pcs_init_dsm(CustomScanState *node, ParallelContext *pcxt, void *coord)
+{
+ PcsState *st = (PcsState *) node;
+ PcsDSM *dsm = (PcsDSM *) coord;
+
+ pg_atomic_init_u32(&dsm->init_worker_calls, 0);
+ st->dsm = dsm;
+ pcs_n_init_dsm++;
+}
+
+static void
+pcs_reinit_dsm(CustomScanState *node, ParallelContext *pcxt, void *coord)
+{
+ PcsState *st = (PcsState *) node;
+ PcsDSM *dsm = (PcsDSM *) coord;
+
+ pg_atomic_write_u32(&dsm->init_worker_calls, 0);
+ st->dsm = dsm;
+ pcs_n_reinit_dsm++;
+}
+
+static void
+pcs_init_worker(CustomScanState *node, shm_toc *toc, void *coord)
+{
+ PcsState *st = (PcsState *) node;
+ PcsDSM *dsm = (PcsDSM *) coord;
+
+ st->dsm = dsm;
+
+ /* This runs in a worker; the leader reads the aggregate in Shutdown. */
+ pg_atomic_fetch_add_u32(&dsm->init_worker_calls, 1);
+}
+
+static void
+pcs_shutdown(CustomScanState *node)
+{
+ PcsState *st = (PcsState *) node;
+
+ if (st->dsm != NULL)
+ pcs_n_init_worker =
pg_atomic_read_u32(&st->dsm->init_worker_calls);
+ pcs_n_shutdown++;
+}
+
+static void
+pcs_set_rel_pathlist(PlannerInfo *root, RelOptInfo *rel, Index rti,
+ RangeTblEntry *rte)
+{
+ CustomPath *cp;
+ CustomPath *pp;
+ double rows;
+ CdbPathLocus inherited_locus;
+
+ if (prev_pathlist_hook)
+ prev_pathlist_hook(root, rel, rti, rte);
+
+ if (!pcs_enabled)
+ return;
+ if (rte->rtekind != RTE_RELATION || rte->relkind != RELKIND_RELATION)
+ return;
+
+ /*
+ * In Cloudberry every Path must carry a CdbPathLocus. Inherit it from
+ * the first existing pathlist entry (set up by set_plain_rel_pathlist)
+ * so the planner can dispatch our CustomScan the same way a regular
+ * SeqScan would be dispatched.
+ */
+ if (rel->pathlist == NIL)
+ return;
+ inherited_locus = ((Path *) linitial(rel->pathlist))->locus;
+
+ rows = rel->tuples > 0 ? rel->tuples : 1.0;
+
+ cp = makeNode(CustomPath);
+ cp->path.pathtype = T_CustomScan;
+ cp->path.parent = rel;
+ cp->path.pathtarget = rel->reltarget;
+ cp->path.param_info = NULL;
+ cp->path.parallel_aware = false;
+ cp->path.parallel_safe = true;
+ cp->path.parallel_workers = 0;
+ cp->path.pathkeys = NIL;
+ cp->path.rows = rows;
+ cp->path.startup_cost = 0;
+ cp->path.total_cost = rows * 0.001; /* cheaper than seqscan to win
*/
+ cp->path.locus = inherited_locus;
+ cp->flags = 0;
+ /*
+ * A freshly built SeqScan path is the wrapper's child. It must be a
new
+ * node (not one already in rel->pathlist): add_path() below may pfree a
+ * dominated seqscan path, which would leave a dangling child pointer.
+ */
+ cp->custom_paths = list_make1(create_seqscan_path(root, rel, NULL, 0));
+ cp->custom_private = NIL;
+ cp->methods = &pcs_path_methods;
+ add_path(rel, &cp->path, root);
+
+ if (rel->consider_parallel)
+ {
+ CdbPathLocus parallel_locus = cdbpathlocus_from_baserel(root,
rel, 2);
+
+ pp = makeNode(CustomPath);
+ pp->path.pathtype = T_CustomScan;
+ pp->path.parent = rel;
+ pp->path.pathtarget = rel->reltarget;
+ pp->path.param_info = NULL;
+ pp->path.parallel_aware = true;
+ pp->path.parallel_safe = true;
+ pp->path.parallel_workers = parallel_locus.parallel_workers;
+ pp->path.pathkeys = NIL;
+ pp->path.rows = rows / Max(2, parallel_locus.parallel_workers);
+ pp->path.startup_cost = 0;
+ pp->path.total_cost = pp->path.rows * 0.0001; /* much cheaper
*/
+ pp->path.locus = parallel_locus;
+ pp->flags = 0;
+ /* Parallel-aware child so it splits blocks across workers. */
+ pp->custom_paths =
+ list_make1(create_seqscan_path(root, rel, NULL,
+
parallel_locus.parallel_workers));
+ pp->custom_private = NIL;
+ pp->methods = &pcs_path_methods;
+ add_partial_path(rel, &pp->path);
+ }
+}
+
+Datum
+pcs_get_hook_calls(PG_FUNCTION_ARGS)
+{
+ TupleDesc tupdesc;
+ HeapTuple tuple;
+ Datum values[5];
+ bool nulls[5] = {false, false, false, false, false};
+
+ if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
+ ereport(ERROR,
+ (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+ errmsg("function returning record called in
context that cannot accept type record")));
+
+ tupdesc = BlessTupleDesc(tupdesc);
+
+ values[0] = Int64GetDatum(pcs_n_estimate);
+ values[1] = Int64GetDatum(pcs_n_init_dsm);
+ values[2] = Int64GetDatum(pcs_n_reinit_dsm);
+ values[3] = Int64GetDatum(pcs_n_init_worker);
+ values[4] = Int64GetDatum(pcs_n_shutdown);
+
+ tuple = heap_form_tuple(tupdesc, values, nulls);
+ PG_RETURN_DATUM(HeapTupleGetDatum(tuple));
+}
+
+void
+_PG_init(void)
+{
+ DefineCustomBoolVariable("parallel_customscan.enabled",
+ "Replace seqscan paths
with TestParallelCustomScan.",
+ NULL,
+ &pcs_enabled,
+ false,
+ PGC_USERSET,
+ 0,
+ NULL, NULL, NULL);
+
+ RegisterCustomScanMethods(&pcs_scan_methods);
+
+ prev_pathlist_hook = set_rel_pathlist_hook;
+ set_rel_pathlist_hook = pcs_set_rel_pathlist;
+}
+
+void
+_PG_fini(void)
+{
+ set_rel_pathlist_hook = prev_pathlist_hook;
+}
diff --git a/src/test/modules/parallel_customscan/parallel_customscan.control
b/src/test/modules/parallel_customscan/parallel_customscan.control
new file mode 100644
index 00000000000..d9657db2036
--- /dev/null
+++ b/src/test/modules/parallel_customscan/parallel_customscan.control
@@ -0,0 +1,5 @@
+# parallel_customscan extension
+comment = 'Exercise parallel-aware CustomScan dispatch (Cloudberry MPP
walkers)'
+default_version = '1.0'
+module_pathname = '$libdir/parallel_customscan'
+relocatable = true
diff --git a/src/test/modules/parallel_customscan/sql/parallel_customscan.sql
b/src/test/modules/parallel_customscan/sql/parallel_customscan.sql
new file mode 100644
index 00000000000..08821ccb8e0
--- /dev/null
+++ b/src/test/modules/parallel_customscan/sql/parallel_customscan.sql
@@ -0,0 +1,76 @@
+-- start_matchsubs
+-- m/\(actual rows=[^)]*\)/
+-- s/\(actual rows=[^)]*\)/(actual rows=...)/
+-- end_matchsubs
+-- start_matchignore
+-- m/^\s*Buckets: \d+ Batches: \d+ Memory Usage: \d+kB/
+-- end_matchignore
+
+CREATE EXTENSION parallel_customscan;
+
+-- Set up data BEFORE enabling our hook so ANALYZE doesn't traverse it.
+CREATE TABLE pcs_t (a int)
+ WITH (parallel_workers = 2)
+ DISTRIBUTED BY (a);
+INSERT INTO pcs_t SELECT generate_series(1, 1000);
+ANALYZE pcs_t;
+
+SET optimizer = off;
+SET enable_parallel = on;
+SET enable_seqscan = off;
+SET min_parallel_table_scan_size = 0;
+SET parallel_setup_cost = 0;
+SET parallel_tuple_cost = 0;
+SET max_parallel_workers_per_gather = 2;
+
+-- (1) Enabling the hook replaces the Parallel Seq Scan with a Custom Scan.
+EXPLAIN (COSTS OFF) SELECT count(*) FROM pcs_t;
+SET parallel_customscan.enabled = on;
+EXPLAIN (COSTS OFF) SELECT count(*) FROM pcs_t;
+
+-- (2) Correctness: results match only if every parallel worker runs our scan.
+SELECT count(*) FROM pcs_t;
+SELECT sum(a) FROM pcs_t;
+
+-- (3) Scan-level qualifier and projection through the custom scan.
+EXPLAIN (COSTS OFF) SELECT a FROM pcs_t WHERE a > 990;
+SELECT a FROM pcs_t WHERE a > 990 ORDER BY a;
+SELECT count(*) FROM pcs_t WHERE a % 2 = 0;
+
+-- (4) Join with both inputs scanned by the custom scan, plus a scan-level
qual.
+CREATE TABLE pcs_t2 (a int)
+ WITH (parallel_workers = 2)
+ DISTRIBUTED BY (a);
+INSERT INTO pcs_t2 SELECT generate_series(1, 500);
+ANALYZE pcs_t2;
+EXPLAIN (COSTS OFF)
+ SELECT count(*) FROM pcs_t x JOIN pcs_t2 y ON x.a = y.a WHERE x.a <= 100;
+SELECT count(*) FROM pcs_t x JOIN pcs_t2 y ON x.a = y.a WHERE x.a <= 100;
+
+-- (5) Empty relation: the custom scan must handle an immediate end-of-scan.
+CREATE TABLE pcs_empty (a int)
+ WITH (parallel_workers = 2)
+ DISTRIBUTED BY (a);
+ANALYZE pcs_empty;
+EXPLAIN (COSTS OFF) SELECT count(*) FROM pcs_empty;
+SELECT count(*) FROM pcs_empty;
+SELECT * FROM pcs_empty;
+
+-- (6) Serial path: with no workers, a non-parallel Custom Scan is used.
+SET max_parallel_workers_per_gather = 0;
+EXPLAIN (COSTS OFF) SELECT count(*) FROM pcs_t;
+SELECT count(*) FROM pcs_t;
+SELECT a FROM pcs_t WHERE a > 995 ORDER BY a;
+SET max_parallel_workers_per_gather = 2;
+
+-- (7) EXPLAIN ANALYZE: exercises planstate_walk_kids' custom_ps recursion.
+EXPLAIN (ANALYZE, TIMING OFF, COSTS OFF, SUMMARY OFF)
+ SELECT count(*) FROM pcs_t;
+EXPLAIN (ANALYZE, TIMING OFF, COSTS OFF, SUMMARY OFF)
+ SELECT count(*) FROM pcs_t x JOIN pcs_t2 y ON x.a = y.a;
+
+-- cleanup
+DROP TABLE pcs_t2;
+DROP TABLE pcs_empty;
+DROP TABLE pcs_t;
+DROP EXTENSION parallel_customscan;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
