From d9ec75289b57d2dd664a5a9f4775fdd95a37dfb2 Mon Sep 17 00:00:00 2001
From: Zhong ShiHao <zhong950419@gmail.com>
Date: Mon, 7 Sep 2026 21:18:27 -0400
Subject: [PATCH v2 1/2] Add a planner support function for
 generate_subscripts()

Previously the planner had no way to estimate how many rows
generate_subscripts() returns, so it always used the function's
prorows value of 1000.  Its close sibling unnest() has computed a
proper estimate from its array argument since planner support
functions were introduced in v12.

Add generate_subscripts_support(), which handles SupportRequestRows
as follows:

* If any argument is a constant NULL, report zero rows: the function
  is strict.

* If both the array and the dimension number are plan-time constants,
  report the length of the requested dimension exactly.

* If only the dimension number is known and it is 1, use
  estimate_array_length(), which can consult the array column's
  statistics.  For one-dimensional arrays, by far the common case,
  the total element count it reports equals the length of dimension 1.

* Otherwise report nothing, and the planner falls back on prorows as
  before.
---
 src/backend/utils/adt/arrayfuncs.c        | 111 ++++++++++++++
 src/include/catalog/pg_proc.dat           |   9 +-
 src/test/regress/expected/planner_est.out | 167 ++++++++++++++++++++++
 src/test/regress/sql/planner_est.sql      |  97 +++++++++++++
 4 files changed, 382 insertions(+), 2 deletions(-)

diff --git a/src/backend/utils/adt/arrayfuncs.c b/src/backend/utils/adt/arrayfuncs.c
index ef66182b047..04dacacc512 100644
--- a/src/backend/utils/adt/arrayfuncs.c
+++ b/src/backend/utils/adt/arrayfuncs.c
@@ -6000,6 +6000,117 @@ generate_subscripts_nodir(PG_FUNCTION_ARGS)
 	return generate_subscripts(fcinfo);
 }
 
+/*
+ * Planner support function for generate_subscripts(anyarray, int [, bool])
+ *
+ * Unlike unnest(), generate_subscripts() returns one row per subscript of
+ * the *requested dimension*, not one row per array element.  Those are the
+ * same thing for a one-dimensional array, but in general we need the length
+ * of dimension "dim", which we can only determine when the array itself is
+ * available as a constant.
+ */
+Datum
+generate_subscripts_support(PG_FUNCTION_ARGS)
+{
+	Node	   *rawreq = (Node *) PG_GETARG_POINTER(0);
+	Node	   *ret = NULL;
+
+	if (IsA(rawreq, SupportRequestRows))
+	{
+		/* Try to estimate the number of rows returned */
+		SupportRequestRows *req = (SupportRequestRows *) rawreq;
+
+		if (is_funcclause(req->node))	/* be paranoid */
+		{
+			List	   *args = ((FuncExpr *) req->node)->args;
+
+			/* We can use estimated argument values here */
+			Node	   *arg1 = estimate_expression_value(req->root,
+														 linitial(args));
+			Node	   *arg2 = estimate_expression_value(req->root,
+														 lsecond(args));
+			Node	   *arg3 = NULL;
+
+			/* The third argument, "reverse", is optional */
+			if (list_length(args) >= 3)
+				arg3 = estimate_expression_value(req->root, lthird(args));
+
+			/*
+			 * The function is strict, so a constant NULL in any argument
+			 * position means that no rows will be returned.  Otherwise we
+			 * need to know the dimension number to say anything at all.
+			 */
+			if ((IsA(arg1, Const) &&
+				 ((Const *) arg1)->constisnull) ||
+				(IsA(arg2, Const) &&
+				 ((Const *) arg2)->constisnull) ||
+				(arg3 != NULL && IsA(arg3, Const) &&
+				 ((Const *) arg3)->constisnull))
+			{
+				req->rows = 0;
+				ret = (Node *) req;
+			}
+			else if (IsA(arg2, Const))
+			{
+				int32		reqdim = DatumGetInt32(((Const *) arg2)->constvalue);
+
+				if (reqdim <= 0)
+				{
+					/* generate_subscripts() returns no rows for such dims */
+					req->rows = 0;
+					ret = (Node *) req;
+				}
+				else if (IsA(arg1, Const))
+				{
+					ArrayType  *arr;
+
+					/*
+					 * We know the array, so we can report the length of the
+					 * requested dimension exactly.  Note that the dimension's
+					 * lower bound does not enter into it.  A dimension the
+					 * array doesn't have yields no rows, which also covers
+					 * empty arrays, whose ndim is zero.
+					 *
+					 * Unlike estimate_array_length(), we don't bother to look
+					 * through array coercions here; if one is in the way we
+					 * simply fall through to the cases below.
+					 */
+					arr = DatumGetArrayTypeP(((Const *) arg1)->constvalue);
+
+					if (reqdim > ARR_NDIM(arr))
+						req->rows = 0;	/* no such dimension */
+					else
+						req->rows = ARR_DIMS(arr)[reqdim - 1];
+					ret = (Node *) req;
+				}
+				else if (reqdim == 1)
+				{
+					/*
+					 * All we have is estimate_array_length(), which counts
+					 * every element rather than the length of one dimension.
+					 * Those agree for 1-D arrays, and since dimension 1 was
+					 * requested it's fair to suppose that's what we have.  If
+					 * not we'll overestimate, but no per-dimension statistics
+					 * exist that could do better.
+					 */
+					req->rows = estimate_array_length(req->root, arg1);
+					ret = (Node *) req;
+				}
+
+				/*
+				 * Otherwise a dimension above the first was requested for an
+				 * array that isn't a constant.  We know nothing about the
+				 * array's dimensionality, let alone the length of the
+				 * requested dimension, so leave ret as NULL to make the
+				 * caller fall back on prorows.
+				 */
+			}
+		}
+	}
+
+	PG_RETURN_POINTER(ret);
+}
+
 /*
  * array_fill_with_lower_bounds
  *		Create and fill array with defined lower bounds.
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 960763ee50b..c1fa1984bd9 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -1711,13 +1711,18 @@
   proargtypes => 'anycompatiblearray anycompatible',
   prosrc => 'array_positions' },
 { oid => '1191', descr => 'array subscripts generator',
-  proname => 'generate_subscripts', prorows => '1000', proretset => 't',
+  proname => 'generate_subscripts', prorows => '1000',
+  prosupport => 'generate_subscripts_support', proretset => 't',
   prorettype => 'int4', proargtypes => 'anyarray int4 bool',
   prosrc => 'generate_subscripts' },
 { oid => '1192', descr => 'array subscripts generator',
-  proname => 'generate_subscripts', prorows => '1000', proretset => 't',
+  proname => 'generate_subscripts', prorows => '1000',
+  prosupport => 'generate_subscripts_support', proretset => 't',
   prorettype => 'int4', proargtypes => 'anyarray int4',
   prosrc => 'generate_subscripts_nodir' },
+{ oid => '9419', descr => 'planner support for generate_subscripts',
+  proname => 'generate_subscripts_support', prorettype => 'internal',
+  proargtypes => 'internal', prosrc => 'generate_subscripts_support' },
 { oid => '1193', descr => 'array constructor with value',
   proname => 'array_fill', proisstrict => 'f', prorettype => 'anyarray',
   proargtypes => 'anyelement _int4', prosrc => 'array_fill' },
diff --git a/src/test/regress/expected/planner_est.out b/src/test/regress/expected/planner_est.out
index 236cb274a78..44eb5eeca32 100644
--- a/src/test/regress/expected/planner_est.out
+++ b/src/test/regress/expected/planner_est.out
@@ -183,6 +183,173 @@ false, true, false, true);
  Function Scan on generate_series g  (cost=N..N rows=1000 width=N)
 (1 row)
 
+--
+-- Test the SupportRequestRows support function for generate_subscripts()
+--
+-- A constant array gives an exact estimate for the requested dimension
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{10,20,30,40,50}'::int[], 1) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=5 width=N) (actual rows=5.00 loops=1)
+(1 row)
+
+-- As above but for the 3-argument form; "reverse" cannot change the row count
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{10,20,30,40,50}'::int[], 1, true) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=5 width=N) (actual rows=5.00 loops=1)
+(1 row)
+
+-- Ensure a non-zero lower bound does not affect the estimate
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('[5:9]={10,20,30,40,50}'::int[], 1) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=5 width=N) (actual rows=5.00 loops=1)
+(1 row)
+
+-- Ensure each dimension of a multi-dimensional array is estimated separately
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{{1,2,3},{4,5,6}}'::int[], 1) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=2 width=N) (actual rows=2.00 loops=1)
+(1 row)
+
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{{1,2,3},{4,5,6}}'::int[], 2) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=3 width=N) (actual rows=3.00 loops=1)
+(1 row)
+
+-- Ensure a dimension above the second is handled too.  Each dimension of this
+-- array has a different length, so a mix-up could not go unnoticed.
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{{{1,2,3,4},{5,6,7,8},{9,10,11,12}},{{13,14,15,16},{17,18,19,20},{21,22,23,24}}}'::int[], 3) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=4 width=N) (actual rows=4.00 loops=1)
+(1 row)
+
+-- Ensure cases which return no rows estimate 1 row after clamping.  Try an
+-- out-of-range dimension, a dimension below the first, and an empty array.
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{{1,2,3},{4,5,6}}'::int[], 3) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=1 width=N) (actual rows=0.00 loops=1)
+(1 row)
+
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{10,20,30}'::int[], 0) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=1 width=N) (actual rows=0.00 loops=1)
+(1 row)
+
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{}'::int[], 1) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=1 width=N) (actual rows=0.00 loops=1)
+(1 row)
+
+-- Ensure a constant NULL in any argument position estimates no rows, since
+-- generate_subscripts() is strict
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts(NULL::int[], 1) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=1 width=N) (actual rows=0.00 loops=1)
+(1 row)
+
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{10,20,30}'::int[], NULL) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=1 width=N) (actual rows=0.00 loops=1)
+(1 row)
+
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{10,20,30}'::int[], 1, NULL) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=1 width=N) (actual rows=0.00 loops=1)
+(1 row)
+
+-- An ArrayExpr is estimated by way of estimate_array_length()
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts(ARRAY[1, 2, (SELECT 3)], 1) g(s);$$,
+true, true, false, true);
+                                      explain_mask_costs                                       
+-----------------------------------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=3 width=N) (actual rows=3.00 loops=1)
+   InitPlan expr_1
+     ->  Result  (cost=N..N rows=1 width=N) (actual rows=1.00 loops=1)
+(3 rows)
+
+-- Ensure we get the default row estimate when the dimension number isn't a
+-- constant, since then we don't even know which dimension is wanted
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts(ARRAY[1, 2, (SELECT 3)], (SELECT 1)) g(s);$$,
+false, true, false, true);
+                          explain_mask_costs                           
+-----------------------------------------------------------------------
+ Function Scan on generate_subscripts g  (cost=N..N rows=1000 width=N)
+   InitPlan expr_1
+     ->  Result  (cost=N..N rows=1 width=N)
+   InitPlan expr_2
+     ->  Result  (cost=N..N rows=1 width=N)
+(5 rows)
+
+-- Ensure a Var array is estimated from the DECHIST statistics.  Every array
+-- below holds exactly 7 distinct elements, so the average distinct element
+-- count does not depend on which rows ANALYZE happens to sample.
+CREATE TEMP TABLE subscript_table_1 AS
+  SELECT array_agg(g) AS a FROM generate_series(1, 100) i,
+    LATERAL generate_series(1, 7) g GROUP BY i;
+ANALYZE subscript_table_1;
+-- Memoize is disabled here only to keep the plan shape stable; the node's
+-- capacity and hit percentage are not masked by explain_mask_costs().
+SET enable_memoize = off;
+SELECT explain_mask_costs($$
+SELECT * FROM subscript_table_1 t, LATERAL generate_subscripts(t.a, 1) g(s);$$,
+true, true, false, true);
+                                          explain_mask_costs                                           
+-------------------------------------------------------------------------------------------------------
+ Nested Loop  (cost=N..N rows=700 width=N) (actual rows=700.00 loops=1)
+   ->  Seq Scan on subscript_table_1 t  (cost=N..N rows=100 width=N) (actual rows=100.00 loops=1)
+   ->  Function Scan on generate_subscripts g  (cost=N..N rows=7 width=N) (actual rows=7.00 loops=100)
+(3 rows)
+
+-- As above, but a dimension above the first falls back on prorows: the element
+-- count in the statistics says nothing about the length of that dimension
+SELECT explain_mask_costs($$
+SELECT * FROM subscript_table_1 t, LATERAL generate_subscripts(t.a, 2) g(s);$$,
+false, true, false, true);
+                             explain_mask_costs                              
+-----------------------------------------------------------------------------
+ Nested Loop  (cost=N..N rows=100000 width=N)
+   ->  Seq Scan on subscript_table_1 t  (cost=N..N rows=100 width=N)
+   ->  Function Scan on generate_subscripts g  (cost=N..N rows=1000 width=N)
+(3 rows)
+
+RESET enable_memoize;
 --
 -- Test ScalarArrayOpExpr row estimates for <> ALL for arrays with NULLs.  We
 -- expect the planner to estimate 1 row will match in both of the following
diff --git a/src/test/regress/sql/planner_est.sql b/src/test/regress/sql/planner_est.sql
index 2b696a4e4e5..c4e7ba32cf3 100644
--- a/src/test/regress/sql/planner_est.sql
+++ b/src/test/regress/sql/planner_est.sql
@@ -131,6 +131,103 @@ SELECT explain_mask_costs($$
 SELECT * FROM generate_series(25.0, 2.0, 0.0) g(s);$$,
 false, true, false, true);
 
+--
+-- Test the SupportRequestRows support function for generate_subscripts()
+--
+
+-- A constant array gives an exact estimate for the requested dimension
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{10,20,30,40,50}'::int[], 1) g(s);$$,
+true, true, false, true);
+
+-- As above but for the 3-argument form; "reverse" cannot change the row count
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{10,20,30,40,50}'::int[], 1, true) g(s);$$,
+true, true, false, true);
+
+-- Ensure a non-zero lower bound does not affect the estimate
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('[5:9]={10,20,30,40,50}'::int[], 1) g(s);$$,
+true, true, false, true);
+
+-- Ensure each dimension of a multi-dimensional array is estimated separately
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{{1,2,3},{4,5,6}}'::int[], 1) g(s);$$,
+true, true, false, true);
+
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{{1,2,3},{4,5,6}}'::int[], 2) g(s);$$,
+true, true, false, true);
+
+-- Ensure a dimension above the second is handled too.  Each dimension of this
+-- array has a different length, so a mix-up could not go unnoticed.
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{{{1,2,3,4},{5,6,7,8},{9,10,11,12}},{{13,14,15,16},{17,18,19,20},{21,22,23,24}}}'::int[], 3) g(s);$$,
+true, true, false, true);
+
+-- Ensure cases which return no rows estimate 1 row after clamping.  Try an
+-- out-of-range dimension, a dimension below the first, and an empty array.
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{{1,2,3},{4,5,6}}'::int[], 3) g(s);$$,
+true, true, false, true);
+
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{10,20,30}'::int[], 0) g(s);$$,
+true, true, false, true);
+
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{}'::int[], 1) g(s);$$,
+true, true, false, true);
+
+-- Ensure a constant NULL in any argument position estimates no rows, since
+-- generate_subscripts() is strict
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts(NULL::int[], 1) g(s);$$,
+true, true, false, true);
+
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{10,20,30}'::int[], NULL) g(s);$$,
+true, true, false, true);
+
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts('{10,20,30}'::int[], 1, NULL) g(s);$$,
+true, true, false, true);
+
+-- An ArrayExpr is estimated by way of estimate_array_length()
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts(ARRAY[1, 2, (SELECT 3)], 1) g(s);$$,
+true, true, false, true);
+
+-- Ensure we get the default row estimate when the dimension number isn't a
+-- constant, since then we don't even know which dimension is wanted
+SELECT explain_mask_costs($$
+SELECT * FROM generate_subscripts(ARRAY[1, 2, (SELECT 3)], (SELECT 1)) g(s);$$,
+false, true, false, true);
+
+-- Ensure a Var array is estimated from the DECHIST statistics.  Every array
+-- below holds exactly 7 distinct elements, so the average distinct element
+-- count does not depend on which rows ANALYZE happens to sample.
+CREATE TEMP TABLE subscript_table_1 AS
+  SELECT array_agg(g) AS a FROM generate_series(1, 100) i,
+    LATERAL generate_series(1, 7) g GROUP BY i;
+ANALYZE subscript_table_1;
+
+-- Memoize is disabled here only to keep the plan shape stable; the node's
+-- capacity and hit percentage are not masked by explain_mask_costs().
+SET enable_memoize = off;
+
+SELECT explain_mask_costs($$
+SELECT * FROM subscript_table_1 t, LATERAL generate_subscripts(t.a, 1) g(s);$$,
+true, true, false, true);
+
+-- As above, but a dimension above the first falls back on prorows: the element
+-- count in the statistics says nothing about the length of that dimension
+SELECT explain_mask_costs($$
+SELECT * FROM subscript_table_1 t, LATERAL generate_subscripts(t.a, 2) g(s);$$,
+false, true, false, true);
+
+RESET enable_memoize;
+
 --
 -- Test ScalarArrayOpExpr row estimates for <> ALL for arrays with NULLs.  We
 -- expect the planner to estimate 1 row will match in both of the following
