Allow virtual columns in index expressions or predicate

From c19c72299fd4eb970086f03a2cd2913298e45816 Mon Sep 17 00:00:00 2001
From: Richard Guo <guofenglinux@gmail.com>
Date: Tue, 30 Sep 2025 17:20:24 +0900
Subject: [PATCH v1] Allow virtual columns in index expressions or predicate

---
 src/backend/commands/indexcmds.c              | 36 +++++-------
 .../regress/expected/generated_virtual.out    | 55 +++++++++++++++----
 src/test/regress/sql/generated_virtual.sql    | 24 ++++----
 3 files changed, 69 insertions(+), 46 deletions(-)

diff --git a/src/backend/commands/indexcmds.c b/src/backend/commands/indexcmds.c
index ca2bde62e82..f269cd20227 100644
--- a/src/backend/commands/indexcmds.c
+++ b/src/backend/commands/indexcmds.c
@@ -53,6 +53,7 @@
 #include "parser/parse_utilcmd.h"
 #include "partitioning/partdesc.h"
 #include "pgstat.h"
+#include "rewrite/rewriteHandler.h"
 #include "rewrite/rewriteManip.h"
 #include "storage/lmgr.h"
 #include "storage/proc.h"
@@ -1097,7 +1098,6 @@ DefineIndex(Oid tableId,
 		}
 	}
 
-
 	/*
 	 * We disallow indexes on system columns.  They would not necessarily get
 	 * updated correctly, and they don't seem useful anyway.
@@ -1126,13 +1126,11 @@ DefineIndex(Oid tableId,
 	}
 
 	/*
-	 * Also check for system and generated columns used in expressions or
-	 * predicates.
+	 * Also check for system columns used in expressions or predicates.
 	 */
 	if (indexInfo->ii_Expressions || indexInfo->ii_Predicate)
 	{
 		Bitmapset  *indexattrs = NULL;
-		int			j;
 
 		pull_varattnos((Node *) indexInfo->ii_Expressions, 1, &indexattrs);
 		pull_varattnos((Node *) indexInfo->ii_Predicate, 1, &indexattrs);
@@ -1145,26 +1143,20 @@ DefineIndex(Oid tableId,
 						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						 errmsg("index creation on system columns is not supported")));
 		}
-
-		/*
-		 * XXX Virtual generated columns in index expressions or predicates
-		 * could be supported, but it needs support in
-		 * RelationGetIndexExpressions() and RelationGetIndexPredicate().
-		 */
-		j = -1;
-		while ((j = bms_next_member(indexattrs, j)) >= 0)
-		{
-			AttrNumber	attno = j + FirstLowInvalidHeapAttributeNumber;
-
-			if (TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
-				ereport(ERROR,
-						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-						 stmt->isconstraint ?
-						 errmsg("unique constraints on virtual generated columns are not supported") :
-						 errmsg("indexes on virtual generated columns are not supported")));
-		}
 	}
 
+	/*
+	 * Expand virtual generated columns in expressions or predicates.
+	 */
+	if (indexInfo->ii_Expressions)
+		indexInfo->ii_Expressions = (List *)
+			expand_generated_columns_in_expr((Node *) indexInfo->ii_Expressions,
+											 rel, 1);
+	if (indexInfo->ii_Predicate)
+		indexInfo->ii_Predicate = (List *)
+			expand_generated_columns_in_expr((Node *) indexInfo->ii_Predicate,
+											 rel, 1);
+
 	/* Is index safe for others to ignore?  See set_indexsafe_procflags() */
 	safe_index = indexInfo->ii_Expressions == NIL &&
 		indexInfo->ii_Predicate == NIL;
diff --git a/src/test/regress/expected/generated_virtual.out b/src/test/regress/expected/generated_virtual.out
index c861bd36c5a..5a047ee74e7 100644
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -741,18 +741,49 @@ ERROR:  primary keys on virtual generated columns are not supported
 -- indexes
 CREATE TABLE gtest22c (a int, b int GENERATED ALWAYS AS (a * 2) VIRTUAL);
 --CREATE INDEX gtest22c_b_idx ON gtest22c (b);
---CREATE INDEX gtest22c_expr_idx ON gtest22c ((b * 3));
---CREATE INDEX gtest22c_pred_idx ON gtest22c (a) WHERE b > 0;
---\d gtest22c
---INSERT INTO gtest22c VALUES (1), (2), (3);
---SET enable_seqscan TO off;
---SET enable_bitmapscan TO off;
+CREATE INDEX gtest22c_expr_idx ON gtest22c ((b * 3));
+CREATE INDEX gtest22c_pred_idx ON gtest22c (a) WHERE b > 0;
+\d gtest22c
+               Table "generated_virtual_tests.gtest22c"
+ Column |  Type   | Collation | Nullable |           Default           
+--------+---------+-----------+----------+-----------------------------
+ a      | integer |           |          | 
+ b      | integer |           |          | generated always as (a * 2)
+Indexes:
+    "gtest22c_expr_idx" btree ((a * 2 * 3))
+    "gtest22c_pred_idx" btree (a) WHERE (a * 2) > 0
+
+INSERT INTO gtest22c VALUES (1), (2), (3);
+SET enable_seqscan TO off;
+SET enable_bitmapscan TO off;
 --EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE b = 4;
 --SELECT * FROM gtest22c WHERE b = 4;
---EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE b * 3 = 6;
---SELECT * FROM gtest22c WHERE b * 3 = 6;
---EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
---SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
+EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE b * 3 = 6;
+                   QUERY PLAN                   
+------------------------------------------------
+ Index Scan using gtest22c_expr_idx on gtest22c
+   Index Cond: (((a * 2) * 3) = 6)
+(2 rows)
+
+SELECT * FROM gtest22c WHERE b * 3 = 6;
+ a | b 
+---+---
+ 1 | 2
+(1 row)
+
+EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Index Only Scan using gtest22c_pred_idx on gtest22c
+   Index Cond: (a = 1)
+(2 rows)
+
+SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
+ a | b 
+---+---
+ 1 | 2
+(1 row)
+
 --ALTER TABLE gtest22c ALTER COLUMN b SET EXPRESSION AS (a * 4);
 --ANALYZE gtest22c;
 --EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE b = 8;
@@ -761,8 +792,8 @@ CREATE TABLE gtest22c (a int, b int GENERATED ALWAYS AS (a * 2) VIRTUAL);
 --SELECT * FROM gtest22c WHERE b * 3 = 12;
 --EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
 --SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
---RESET enable_seqscan;
---RESET enable_bitmapscan;
+RESET enable_seqscan;
+RESET enable_bitmapscan;
 -- foreign keys
 CREATE TABLE gtest23a (x int PRIMARY KEY, y int);
 --INSERT INTO gtest23a VALUES (1, 11), (2, 22), (3, 33);
diff --git a/src/test/regress/sql/generated_virtual.sql b/src/test/regress/sql/generated_virtual.sql
index adfe88d74ae..1843c3006f6 100644
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -394,19 +394,19 @@ CREATE TABLE gtest22b (a int, b int GENERATED ALWAYS AS (a / 2) VIRTUAL, PRIMARY
 -- indexes
 CREATE TABLE gtest22c (a int, b int GENERATED ALWAYS AS (a * 2) VIRTUAL);
 --CREATE INDEX gtest22c_b_idx ON gtest22c (b);
---CREATE INDEX gtest22c_expr_idx ON gtest22c ((b * 3));
---CREATE INDEX gtest22c_pred_idx ON gtest22c (a) WHERE b > 0;
---\d gtest22c
+CREATE INDEX gtest22c_expr_idx ON gtest22c ((b * 3));
+CREATE INDEX gtest22c_pred_idx ON gtest22c (a) WHERE b > 0;
+\d gtest22c
 
---INSERT INTO gtest22c VALUES (1), (2), (3);
---SET enable_seqscan TO off;
---SET enable_bitmapscan TO off;
+INSERT INTO gtest22c VALUES (1), (2), (3);
+SET enable_seqscan TO off;
+SET enable_bitmapscan TO off;
 --EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE b = 4;
 --SELECT * FROM gtest22c WHERE b = 4;
---EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE b * 3 = 6;
---SELECT * FROM gtest22c WHERE b * 3 = 6;
---EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
---SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
+EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE b * 3 = 6;
+SELECT * FROM gtest22c WHERE b * 3 = 6;
+EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
+SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
 
 --ALTER TABLE gtest22c ALTER COLUMN b SET EXPRESSION AS (a * 4);
 --ANALYZE gtest22c;
@@ -416,8 +416,8 @@ CREATE TABLE gtest22c (a int, b int GENERATED ALWAYS AS (a * 2) VIRTUAL);
 --SELECT * FROM gtest22c WHERE b * 3 = 12;
 --EXPLAIN (COSTS OFF) SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
 --SELECT * FROM gtest22c WHERE a = 1 AND b > 0;
---RESET enable_seqscan;
---RESET enable_bitmapscan;
+RESET enable_seqscan;
+RESET enable_bitmapscan;
 
 -- foreign keys
 CREATE TABLE gtest23a (x int PRIMARY KEY, y int);
-- 
2.39.5 (Apple Git-154)