Honor forceDistRandom in subquery pull-up and equality

src/backend/nodes/equalfuncs.c

@@ -3177,6 +3177,7 @@ _equalRangeTblEntry(const RangeTblEntry *a, const RangeTblEntry *b)
 	COMPARE_BITMAPSET_FIELD(updatedCols);
 	COMPARE_BITMAPSET_FIELD(extraUpdatedCols);
 	COMPARE_NODE_FIELD(securityQuals);
+	COMPARE_SCALAR_FIELD(forceDistRandom);
 
 	return true;
 }

src/backend/optimizer/prep/prepjointree.c

@@ -934,14 +934,16 @@ pull_up_subqueries_recurse(PlannerInfo *root, Node *jtnode,
 
 		/*
 		 * Alternatively, is it a simple UNION ALL subquery?  If so, flatten
-		 * into an "append relation".
+		 * into an "append relation".  As above, keep gp_dist_random subqueries
+		 * intact so set_subquery_pathlist can enforce segment execution.
 		 *
 		 * It's safe to do this regardless of whether this query is itself an
 		 * appendrel member.  (If you're thinking we should try to flatten the
 		 * two levels of appendrel together, you're right; but we handle that
 		 * in set_append_rel_pathlist, not here.)
 		 */
 		if (rte->rtekind == RTE_SUBQUERY &&
+			!rte->forceDistRandom &&
 			is_simple_union_all(rte->subquery))
 			return pull_up_simple_union_all(root, jtnode, rte);
 

src/test/regress/expected/gpdist.out

@@ -771,3 +771,99 @@ SELECT * FROM gp_dist_random('"gp.dist.random.schema".gp_dist_random_table_with_
 
 DROP SCHEMA "gp.dist.random.schema" CASCADE;
 NOTICE:  drop cascades to table "gp.dist.random.schema".gp_dist_random_table_with_schema
+-- gp_dist_random must preserve segment execution when a view is a UNION ALL.
+BEGIN;
+SET LOCAL optimizer = off;
+CREATE TEMP VIEW gdr_union AS
+  SELECT gp_execution_segment() AS seg, 1 AS branch FROM gp_id
+  UNION ALL
+  SELECT gp_execution_segment(), 2 FROM gp_id;
+SELECT branch,
+       bool_and(seg >= 0) AS on_segments,
+       count(*) = (SELECT count(*) FROM gp_segment_configuration
+                    WHERE role = 'p' AND content >= 0)
+         AND count(*) = count(DISTINCT seg) AS once_per_segment
+  FROM gp_dist_random('gdr_union')
+ GROUP BY branch
+ ORDER BY branch;
+ branch | on_segments | once_per_segment 
+--------+-------------+------------------
+      1 | t           | t
+      2 | t           | t
+(2 rows)
+
+-- Ordinary access still executes on the coordinator.
+SELECT count(*) = 2 AS two_rows, bool_and(seg = -1) AS on_coordinator
+  FROM gdr_union;
+ two_rows | on_coordinator 
+----------+----------------
+ t        | t
+(1 row)
+
+-- A nested view with a filter must also execute on the segments.
+CREATE TEMP VIEW gdr_union_filtered AS
+  SELECT * FROM gdr_union WHERE branch = 2;
+SELECT bool_and(seg >= 0) AS on_segments,
+       count(*) = (SELECT count(*) FROM gp_segment_configuration
+                    WHERE role = 'p' AND content >= 0)
+         AND count(*) = count(DISTINCT seg) AS once_per_segment
+  FROM gp_dist_random('gdr_union_filtered');
+ on_segments | once_per_segment 
+-------------+------------------
+ t           | t
+(1 row)
+
+-- An impossible filter must remain empty.
+SELECT count(*) FROM gp_dist_random('gdr_union') WHERE branch = 3;
+ count 
+-------
+     0
+(1 row)
+
+ROLLBACK;
+-- RTE equality must distinguish ordinary scans from gp_dist_random scans.
+-- ONLY makes the inheritance flag identical, leaving forceDistRandom different.
+BEGIN;
+SET LOCAL optimizer = off;
+SELECT EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0)
+         AS on_coordinator,
+       EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0)
+         AS on_segments;
+ on_coordinator | on_segments 
+----------------+-------------
+ f              | t
+(1 row)
+
+-- Both OR orders must retain the distributed subquery and return one row.
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0)
+    OR EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0);
+ result 
+--------
+      1
+(1 row)
+
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0)
+    OR EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0);
+ result 
+--------
+      1
+(1 row)
+
+-- Both AND orders must retain the coordinator subquery and return no rows.
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0)
+   AND EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0);
+ result 
+--------
+(0 rows)
+
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0)
+   AND EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0);
+ result 
+--------
+(0 rows)
+
+ROLLBACK;

src/test/regress/expected/gpdist_optimizer.out

@@ -775,3 +775,99 @@ SELECT * FROM gp_dist_random('"gp.dist.random.schema".gp_dist_random_table_with_
 
 DROP SCHEMA "gp.dist.random.schema" CASCADE;
 NOTICE:  drop cascades to table "gp.dist.random.schema".gp_dist_random_table_with_schema
+-- gp_dist_random must preserve segment execution when a view is a UNION ALL.
+BEGIN;
+SET LOCAL optimizer = off;
+CREATE TEMP VIEW gdr_union AS
+  SELECT gp_execution_segment() AS seg, 1 AS branch FROM gp_id
+  UNION ALL
+  SELECT gp_execution_segment(), 2 FROM gp_id;
+SELECT branch,
+       bool_and(seg >= 0) AS on_segments,
+       count(*) = (SELECT count(*) FROM gp_segment_configuration
+                    WHERE role = 'p' AND content >= 0)
+         AND count(*) = count(DISTINCT seg) AS once_per_segment
+  FROM gp_dist_random('gdr_union')
+ GROUP BY branch
+ ORDER BY branch;
+ branch | on_segments | once_per_segment 
+--------+-------------+------------------
+      1 | t           | t
+      2 | t           | t
+(2 rows)
+
+-- Ordinary access still executes on the coordinator.
+SELECT count(*) = 2 AS two_rows, bool_and(seg = -1) AS on_coordinator
+  FROM gdr_union;
+ two_rows | on_coordinator 
+----------+----------------
+ t        | t
+(1 row)
+
+-- A nested view with a filter must also execute on the segments.
+CREATE TEMP VIEW gdr_union_filtered AS
+  SELECT * FROM gdr_union WHERE branch = 2;
+SELECT bool_and(seg >= 0) AS on_segments,
+       count(*) = (SELECT count(*) FROM gp_segment_configuration
+                    WHERE role = 'p' AND content >= 0)
+         AND count(*) = count(DISTINCT seg) AS once_per_segment
+  FROM gp_dist_random('gdr_union_filtered');
+ on_segments | once_per_segment 
+-------------+------------------
+ t           | t
+(1 row)
+
+-- An impossible filter must remain empty.
+SELECT count(*) FROM gp_dist_random('gdr_union') WHERE branch = 3;
+ count 
+-------
+     0
+(1 row)
+
+ROLLBACK;
+-- RTE equality must distinguish ordinary scans from gp_dist_random scans.
+-- ONLY makes the inheritance flag identical, leaving forceDistRandom different.
+BEGIN;
+SET LOCAL optimizer = off;
+SELECT EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0)
+         AS on_coordinator,
+       EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0)
+         AS on_segments;
+ on_coordinator | on_segments 
+----------------+-------------
+ f              | t
+(1 row)
+
+-- Both OR orders must retain the distributed subquery and return one row.
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0)
+    OR EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0);
+ result 
+--------
+      1
+(1 row)
+
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0)
+    OR EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0);
+ result 
+--------
+      1
+(1 row)
+
+-- Both AND orders must retain the coordinator subquery and return no rows.
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0)
+   AND EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0);
+ result 
+--------
+(0 rows)
+
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0)
+   AND EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0);
+ result 
+--------
+(0 rows)
+
+ROLLBACK;

src/test/regress/sql/gpdist.sql

@@ -535,3 +535,57 @@ CREATE TABLE "gp.dist.random.schema".gp_dist_random_table_with_schema
     AS SELECT * FROM gp_dist_random('"gp_dist_random_table"');
 SELECT * FROM gp_dist_random('"gp.dist.random.schema".gp_dist_random_table_with_schema');
 DROP SCHEMA "gp.dist.random.schema" CASCADE;
+
+-- gp_dist_random must preserve segment execution when a view is a UNION ALL.
+BEGIN;
+SET LOCAL optimizer = off;
+CREATE TEMP VIEW gdr_union AS
+  SELECT gp_execution_segment() AS seg, 1 AS branch FROM gp_id
+  UNION ALL
+  SELECT gp_execution_segment(), 2 FROM gp_id;
+SELECT branch,
+       bool_and(seg >= 0) AS on_segments,
+       count(*) = (SELECT count(*) FROM gp_segment_configuration
+                    WHERE role = 'p' AND content >= 0)
+         AND count(*) = count(DISTINCT seg) AS once_per_segment
+  FROM gp_dist_random('gdr_union')
+ GROUP BY branch
+ ORDER BY branch;
+-- Ordinary access still executes on the coordinator.
+SELECT count(*) = 2 AS two_rows, bool_and(seg = -1) AS on_coordinator
+  FROM gdr_union;
+-- A nested view with a filter must also execute on the segments.
+CREATE TEMP VIEW gdr_union_filtered AS
+  SELECT * FROM gdr_union WHERE branch = 2;
+SELECT bool_and(seg >= 0) AS on_segments,
+       count(*) = (SELECT count(*) FROM gp_segment_configuration
+                    WHERE role = 'p' AND content >= 0)
+         AND count(*) = count(DISTINCT seg) AS once_per_segment
+  FROM gp_dist_random('gdr_union_filtered');
+-- An impossible filter must remain empty.
+SELECT count(*) FROM gp_dist_random('gdr_union') WHERE branch = 3;
+ROLLBACK;
+
+-- RTE equality must distinguish ordinary scans from gp_dist_random scans.
+-- ONLY makes the inheritance flag identical, leaving forceDistRandom different.
+BEGIN;
+SET LOCAL optimizer = off;
+SELECT EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0)
+         AS on_coordinator,
+       EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0)
+         AS on_segments;
+-- Both OR orders must retain the distributed subquery and return one row.
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0)
+    OR EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0);
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0)
+    OR EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0);
+-- Both AND orders must retain the coordinator subquery and return no rows.
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0)
+   AND EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0);
+SELECT 1 AS result
+ WHERE EXISTS (SELECT 1 FROM gp_dist_random('gp_id') WHERE gp_execution_segment() >= 0)
+   AND EXISTS (SELECT 1 FROM ONLY gp_id WHERE gp_execution_segment() >= 0);
+ROLLBACK;