[PATCH] Add support for SAOP in the optimizer for partial index paths

diff --git a/src/backend/optimizer/path/indxpath.c b/src/backend/optimizer/path/indxpath.c
index 51b9d6677d3..f07e3d14086 100644
--- a/src/backend/optimizer/path/indxpath.c
+++ b/src/backend/optimizer/path/indxpath.c
@@ -36,12 +36,23 @@
 #include "optimizer/restrictinfo.h"
 #include "utils/lsyscache.h"
 #include "utils/selfuncs.h"
+#include "utils/array.h"
 
 
 /* XXX see PartCollMatchesExprColl */
 #define IndexCollMatchesExprColl(idxcollation, exprcollation) \
 	((idxcollation) == InvalidOid || (idxcollation) == (exprcollation))
 
+/*
+ * Proof attempts involving large arrays in ScalarArrayOpExpr nodes are
+ * likely to require O(N^2) time, and more often than not fail anyway.
+ * So we set an arbitrary limit on the number of array elements that
+ * we will allow to be treated as an AND or OR clause.
+ * XXX is it worth exposing this as a GUC knob?
+ * This block was copied verbatim from ../utils/predtest.c
+ */
+#define MAX_SAOP_ARRAY_SIZE		100
+
 /* Whether we are looking for plain indexscan, bitmap scan, or either */
 typedef enum
 {
@@ -113,6 +124,8 @@ static List *build_paths_for_OR(PlannerInfo *root, RelOptInfo *rel,
 								List *clauses, List *other_clauses);
 static List *generate_bitmap_or_paths(PlannerInfo *root, RelOptInfo *rel,
 									  List *clauses, List *other_clauses);
+static List *generate_bitmap_saop_paths(PlannerInfo *root, RelOptInfo *rel,
+										List *clauses, List *indexes);
 static Path *choose_bitmap_and(PlannerInfo *root, RelOptInfo *rel,
 							   List *paths);
 static int	path_usage_comparator(const void *a, const void *b);
@@ -325,6 +338,15 @@ create_index_paths(PlannerInfo *root, RelOptInfo *rel)
 										  rel->baserestrictinfo, NIL);
 	bitindexpaths = list_concat(bitindexpaths, indexpaths);
 
+	/*
+	 * Now generate BitmapOrPaths for any suitable SAOP-clauses present in the
+	 * restriction list.  Add these to bitindexpaths.
+	 */
+	indexpaths = generate_bitmap_saop_paths(root, rel,
+											rel->baserestrictinfo,
+											rel->indexlist);
+	bitindexpaths = list_concat(bitindexpaths, indexpaths);
+
 	/*
 	 * Likewise, generate BitmapOrPaths for any suitable OR-clauses present in
 	 * the joinclause list.  Add these to bitjoinpaths.
@@ -1772,6 +1794,227 @@ generate_bitmap_or_paths(PlannerInfo *root, RelOptInfo *rel,
 }
 
 
+/*
+ * generate_bitmap_saop_paths
+ *		Look through the list of clauses to find IN/ANY clauses, and generate
+ *		a BitmapOrPath for each one we can handle via an index.  Return a list
+ *		of the generated BitmapOrPaths.
+ */
+static List *
+generate_bitmap_saop_paths(PlannerInfo *root, RelOptInfo *rel,
+			   List *clauses, List *indexes)
+{
+	ListCell   *lc;
+	List	   *result = NIL;
+
+	/*
+	 * Iterate through each restriction clause to find ScalarArrayOpExprs.
+	 */
+	foreach(lc, clauses)
+	{
+		bool			elem_typbyval;
+		char			elem_typalign;
+		int16			elem_typlen;
+		int			nelems = 0;
+		int			index_pos = 0;
+		Oid			elem_type;
+		ListCell		*idx_lc;
+		Node			*leftop;
+		Node			*rightop;
+		const Const		*arrayconst;
+		ArrayType		*arrayval;
+		Datum			*elem_values;
+		bool			*elem_nulls;
+		ScalarArrayOpExpr	*saop;
+		List			*per_saop_paths = NIL;
+		Bitmapset		*suitable_indexes = NULL;
+		RestrictInfo		*rinfo = lfirst_node(RestrictInfo, lc);
+
+		/* We only care about ScalarArrayOpExpr */
+		if (!IsA(rinfo->clause, ScalarArrayOpExpr))
+			continue;
+
+		saop = (ScalarArrayOpExpr *) rinfo->clause;
+
+		/* We only handle IN clauses (operator ANY), not ALL clauses. */
+		if (!saop->useOr)
+			continue;
+
+		/*
+		 * The logic to build bitmap paths requires concrete values.
+		 * Therefore, we can only handle the case where the array is a
+		 * non-null Const.
+		 */
+		leftop = (Node *) linitial(saop->args);
+		rightop = (Node *) lsecond(saop->args);
+		if (!rightop || !IsA(rightop, Const) || ((Const *) rightop)->constisnull)
+			continue;
+
+		arrayconst = (const Const *) rightop;
+
+		/*
+		 * Deconstruct the array constant into its constituent elements.
+		 */
+		arrayval = DatumGetArrayTypeP(arrayconst->constvalue);
+		if (ARR_NDIM(arrayval) != 1)
+			continue;
+
+		/* Safety valve for huge IN lists */
+		nelems = ArrayGetNItems(1, ARR_DIMS(arrayval));
+		if (nelems < 1 || nelems > MAX_SAOP_ARRAY_SIZE)
+			continue;
+
+		/*
+		 * Search through all available indexes to find a partial index that
+		 * is satisfied by our newly constructed clause.
+		 */
+		foreach(idx_lc, indexes)
+		{
+			IndexOptInfo *index = (IndexOptInfo *) lfirst(idx_lc);
+
+			/*
+			 * Only consider indexes that support bitmaps and are either
+			 * partial or match the current SAOP.
+			 */
+			if (index->amhasgetbitmap && (index->indpred != NIL || index->predOK))
+			{
+				suitable_indexes = bms_add_member(suitable_indexes, index_pos);
+			}
+			index_pos++;
+		}
+
+		/*
+		 * If no suitable indexes exist, skip expensive array parsing
+		 */
+		if (bms_is_empty(suitable_indexes))
+			continue;
+
+		elem_type = ARR_ELEMTYPE(arrayval);
+		get_typlenbyvalalign(elem_type,
+					&elem_typlen, &elem_typbyval, &elem_typalign);
+
+		deconstruct_array(arrayval, elem_type,
+					elem_typlen, elem_typbyval, elem_typalign,
+					&elem_values, &elem_nulls, &nelems);
+
+		/*
+		 * For each element of the IN list, try to find a matching partial
+		 * index and generate a bitmap path for it.
+		 */
+		for (int i = 0; i < nelems; i++)
+		{
+			Expr		*elem_const;
+			OpExpr		*opclause;
+			RestrictInfo	*new_rinfo;
+			bool		found_path_for_elem = false;
+			const Datum	elem_value = elem_values[i];
+
+			/* We can't build an index qual from a NULL element. */
+			if (elem_nulls[i])
+				break;
+
+			/*
+			 * Construct a new clause: "indexkey = const_element".
+			 */
+			elem_const = (Expr *) makeConst(elem_type,
+							-1,
+							arrayconst->constcollid,
+							elem_typlen,
+							elem_value, false, elem_typbyval);
+
+			opclause = (OpExpr *) make_opclause(saop->opno, BOOLOID, false,
+								(Expr *) copyObject(leftop),
+								elem_const,
+								saop->inputcollid, saop->inputcollid);
+
+			new_rinfo = make_simple_restrictinfo(root, (Expr *) opclause);
+
+			/*
+			 * Search through all candidate indexes we pre-filtered, to find a
+			 * partial index that is satisfied by our newly constructed
+			 * clause.
+			 */
+			index_pos = -1;
+			while ((index_pos = bms_next_member(suitable_indexes, index_pos)) >= 0)
+			{
+				IndexOptInfo *index = (IndexOptInfo *) list_nth(indexes, index_pos);
+
+				/*
+				 * If the index predicate is implied by our new clause, we can
+				 * use it
+				 */
+				if (predicate_implied_by(index->indpred, list_make1(new_rinfo), false))
+				{
+					/*
+					 * We found a usable index. Now build a bitmap path for it
+					 * using the new clause as the index qual.
+					 */
+					List		*indexpaths;
+					IndexClauseSet	clauseset;
+
+					MemSet(&clauseset, 0, sizeof(clauseset));
+					match_clause_to_index(root, new_rinfo, index, &clauseset);
+
+					/* We must have found a match to proceed */
+					if (!clauseset.nonempty)
+						continue;
+
+					indexpaths = build_index_paths(root, rel,
+									index, &clauseset,
+									true,	/* predicate is useful */
+									ST_BITMAPSCAN,
+									NULL);
+
+					if (indexpaths)
+					{
+						/*
+						 * Success! Add the path and stop searching for other
+						 * indexes for this element. We only need one.
+						 */
+						per_saop_paths = list_concat(per_saop_paths, indexpaths);
+						found_path_for_elem = true;
+						break;	/* out of inner index loop */
+					}
+				}
+			}
+
+			/*
+			 * If we could not find any partial index for this element, then
+			 * we cannot satisfy the whole IN clause this way. Abort.
+			 */
+			if (!found_path_for_elem)
+			{
+				per_saop_paths = NIL;	/* Discard any paths found so far */
+				break;			/* out of outer element loop */
+			}
+		}
+
+		bms_free(suitable_indexes);
+
+		/*
+		 * If we successfully found a path for every element of the IN list,
+		 * we can combine them into a BitmapOrPath.
+		 */
+		if (per_saop_paths != NIL)
+		{
+			Path	*bitmapqual = list_length(per_saop_paths) > 1 ?
+				(Path *) create_bitmap_or_path(root, rel, per_saop_paths) :
+				(Path *) linitial(per_saop_paths);
+
+			result = lappend(result, bitmapqual);
+		}
+
+		if (elem_values)
+			pfree(elem_values);
+
+		if (elem_nulls)
+			pfree(elem_nulls);
+	}
+
+	return result;
+}
+
+
 /*
  * choose_bitmap_and
  *		Given a nonempty list of bitmap paths, AND them into one path.
diff --git a/src/test/regress/expected/saop_bitmap-1.out b/src/test/regress/expected/saop_bitmap-1.out
new file mode 100644
index 00000000000..96fd5c058d7
--- /dev/null
+++ b/src/test/regress/expected/saop_bitmap-1.out
@@ -0,0 +1,94 @@
+-- Test ScalarArrayOpExpr for partial indexes
+-- Generate enough data that we can test the use of an index
+CREATE TABLE tpairs(
+   key TEXT,
+   value int
+)
+WITH (autovacuum_enabled = false);
+-- Install the partial index on these values
+CREATE INDEX ON tpairs(key) WHERE key = 'foo';
+CREATE INDEX ON tpairs(key) WHERE key = 'bar';
+CREATE INDEX ON tpairs(key) WHERE key = 'baz';
+-- Dummy data to surround ours
+INSERT INTO tpairs(key, value)
+SELECT n::TEXT, n
+FROM generate_series(1, 10000) AS tmp(n);
+-- Our specific test data to force the partial index
+INSERT INTO tpairs(key, value)
+VALUES
+('foo', 25),
+('bar', 50),
+('baz', 75);
+-- Update statistics
+VACUUM (ANALYZE, FREEZE) tpairs;
+-- We want ensure that an index use is always attempted - just for the test
+SET enable_seqscan=false;
+-- Variant 1 (works without this patch)
+EXPLAIN (COSTS OFF)
+SELECT value
+FROM tpairs
+WHERE
+value > 10 AND
+(
+   key = 'foo' OR
+   key = 'bar' OR
+   key = 'baz'
+);
+                                     QUERY PLAN                                      
+-------------------------------------------------------------------------------------
+ Bitmap Heap Scan on tpairs
+   Recheck Cond: ((key = 'foo'::text) OR (key = 'bar'::text) OR (key = 'baz'::text))
+   Filter: (value > 10)
+   ->  BitmapOr
+         ->  Bitmap Index Scan on tpairs_key_idx
+               Index Cond: (key = 'foo'::text)
+         ->  Bitmap Index Scan on tpairs_key_idx1
+               Index Cond: (key = 'bar'::text)
+         ->  Bitmap Index Scan on tpairs_key_idx2
+               Index Cond: (key = 'baz'::text)
+(10 rows)
+
+-- Variant 2 (works only with this patch)
+EXPLAIN (COSTS OFF)
+SELECT value
+FROM tpairs
+WHERE
+value > 10 AND
+key IN ('foo', 'bar', 'baz');
+                                     QUERY PLAN                                      
+-------------------------------------------------------------------------------------
+ Bitmap Heap Scan on tpairs
+   Recheck Cond: ((key = 'foo'::text) OR (key = 'bar'::text) OR (key = 'baz'::text))
+   Filter: (value > 10)
+   ->  BitmapOr
+         ->  Bitmap Index Scan on tpairs_key_idx
+               Index Cond: (key = 'foo'::text)
+         ->  Bitmap Index Scan on tpairs_key_idx1
+               Index Cond: (key = 'bar'::text)
+         ->  Bitmap Index Scan on tpairs_key_idx2
+               Index Cond: (key = 'baz'::text)
+(10 rows)
+
+-- Variant 3 (works only with this patch)
+EXPLAIN (COSTS OFF)
+SELECT value
+FROM tpairs
+WHERE
+value > 10 AND
+key = ANY('{foo,bar,baz}');
+                                     QUERY PLAN                                      
+-------------------------------------------------------------------------------------
+ Bitmap Heap Scan on tpairs
+   Recheck Cond: ((key = 'foo'::text) OR (key = 'bar'::text) OR (key = 'baz'::text))
+   Filter: (value > 10)
+   ->  BitmapOr
+         ->  Bitmap Index Scan on tpairs_key_idx
+               Index Cond: (key = 'foo'::text)
+         ->  Bitmap Index Scan on tpairs_key_idx1
+               Index Cond: (key = 'bar'::text)
+         ->  Bitmap Index Scan on tpairs_key_idx2
+               Index Cond: (key = 'baz'::text)
+(10 rows)
+
+-- Clean up
+DROP TABLE tpairs;
diff --git a/src/test/regress/expected/saop_bitmap-2.out b/src/test/regress/expected/saop_bitmap-2.out
new file mode 100644
index 00000000000..0fa4054f407
--- /dev/null
+++ b/src/test/regress/expected/saop_bitmap-2.out
@@ -0,0 +1,101 @@
+-- Test ScalarArrayOpExpr for partial indexes
+-- Generate enough data that we can test the use of an index
+CREATE TABLE ipairs(
+   key INT,
+   value INT
+)
+WITH (autovacuum_enabled = false);
+-- Install the partial index on these values
+CREATE INDEX ON ipairs(key) WHERE key = 7;
+CREATE INDEX ON ipairs(key) WHERE key = 8;
+CREATE INDEX ON ipairs(key) WHERE key = 9;
+-- Dummy data to surround ours
+INSERT INTO ipairs(key, value)
+SELECT n, -n
+FROM generate_series(1, 10000) AS tmp(n);
+-- Update statistics
+VACUUM (ANALYZE) ipairs;
+-- Our specific test data to force the partial index
+EXPLAIN (ANALYZE, COSTS OFF, BUFFERS OFF, SUMMARY OFF, TIMING OFF)
+UPDATE ipairs
+SET value = 0
+WHERE key = ANY('{7,8,9}'::INT[]);
+                                    QUERY PLAN                                     
+-----------------------------------------------------------------------------------
+ Update on ipairs (actual rows=0.00 loops=1)
+   ->  Bitmap Heap Scan on ipairs (actual rows=3.00 loops=1)
+         Recheck Cond: ((key = 7) OR (key = 8) OR (key = 9))
+         Heap Blocks: exact=1
+         ->  BitmapOr (actual rows=0.00 loops=1)
+               ->  Bitmap Index Scan on ipairs_key_idx (actual rows=1.00 loops=1)
+                     Index Cond: (key = 7)
+                     Index Searches: 1
+               ->  Bitmap Index Scan on ipairs_key_idx1 (actual rows=1.00 loops=1)
+                     Index Cond: (key = 8)
+                     Index Searches: 1
+               ->  Bitmap Index Scan on ipairs_key_idx2 (actual rows=1.00 loops=1)
+                     Index Cond: (key = 9)
+                     Index Searches: 1
+(14 rows)
+
+-- We want ensure that an index use is always attempted - just for the test
+SET enable_seqscan=false;
+-- Variant 1 (works without this patch)
+EXPLAIN (COSTS OFF)
+SELECT value=0
+FROM ipairs
+WHERE
+   key = 7 OR
+   key = 8 OR
+   key = 9;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Bitmap Heap Scan on ipairs
+   Recheck Cond: ((key = 7) OR (key = 8) OR (key = 9))
+   ->  BitmapOr
+         ->  Bitmap Index Scan on ipairs_key_idx
+               Index Cond: (key = 7)
+         ->  Bitmap Index Scan on ipairs_key_idx1
+               Index Cond: (key = 8)
+         ->  Bitmap Index Scan on ipairs_key_idx2
+               Index Cond: (key = 9)
+(9 rows)
+
+-- Variant 2 (works only with this patch)
+EXPLAIN (COSTS OFF)
+SELECT value=0
+FROM ipairs
+WHERE key IN (7, 8, 9);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Bitmap Heap Scan on ipairs
+   Recheck Cond: ((key = 7) OR (key = 8) OR (key = 9))
+   ->  BitmapOr
+         ->  Bitmap Index Scan on ipairs_key_idx
+               Index Cond: (key = 7)
+         ->  Bitmap Index Scan on ipairs_key_idx1
+               Index Cond: (key = 8)
+         ->  Bitmap Index Scan on ipairs_key_idx2
+               Index Cond: (key = 9)
+(9 rows)
+
+-- Variant 3 (works only with this patch)
+EXPLAIN (COSTS OFF)
+SELECT value=0
+FROM ipairs
+WHERE key = ANY('{7,8,9}'::INT[]);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Bitmap Heap Scan on ipairs
+   Recheck Cond: ((key = 7) OR (key = 8) OR (key = 9))
+   ->  BitmapOr
+         ->  Bitmap Index Scan on ipairs_key_idx
+               Index Cond: (key = 7)
+         ->  Bitmap Index Scan on ipairs_key_idx1
+               Index Cond: (key = 8)
+         ->  Bitmap Index Scan on ipairs_key_idx2
+               Index Cond: (key = 9)
+(9 rows)
+
+-- Clean up
+DROP TABLE ipairs;
diff --git a/src/test/regress/parallel_schedule b/src/test/regress/parallel_schedule
index 021d57f66bb..e17f03d8d04 100644
--- a/src/test/regress/parallel_schedule
+++ b/src/test/regress/parallel_schedule
@@ -102,7 +102,7 @@ test: publication subscription
 # Another group of parallel tests
 # select_views depends on create_view
 # ----------
-test: select_views portals_p2 foreign_key cluster dependency guc bitmapops combocid tsearch tsdicts foreign_data window xmlmap functional_deps advisory_lock indirect_toast equivclass stats_rewrite
+test: select_views portals_p2 foreign_key cluster dependency guc bitmapops combocid tsearch tsdicts foreign_data window xmlmap functional_deps advisory_lock indirect_toast equivclass stats_rewrite saop_bitmap-1 saop_bitmap-2
 
 # ----------
 # Another group of parallel tests (JSON related)
diff --git a/src/test/regress/sql/saop_bitmap-1.sql b/src/test/regress/sql/saop_bitmap-1.sql
new file mode 100644
index 00000000000..167c391e15d
--- /dev/null
+++ b/src/test/regress/sql/saop_bitmap-1.sql
@@ -0,0 +1,61 @@
+-- Test ScalarArrayOpExpr for partial indexes
+
+-- Generate enough data that we can test the use of an index
+CREATE TABLE tpairs(
+   key TEXT,
+   value int
+)
+WITH (autovacuum_enabled = false);
+-- Install the partial index on these values
+CREATE INDEX ON tpairs(key) WHERE key = 'foo';
+CREATE INDEX ON tpairs(key) WHERE key = 'bar';
+CREATE INDEX ON tpairs(key) WHERE key = 'baz';
+
+-- Dummy data to surround ours
+INSERT INTO tpairs(key, value)
+SELECT n::TEXT, n
+FROM generate_series(1, 10000) AS tmp(n);
+
+-- Our specific test data to force the partial index
+INSERT INTO tpairs(key, value)
+VALUES
+('foo', 25),
+('bar', 50),
+('baz', 75);
+
+-- Update statistics
+VACUUM (ANALYZE, FREEZE) tpairs;
+
+-- We want ensure that an index use is always attempted - just for the test
+SET enable_seqscan=false;
+
+-- Variant 1 (works without this patch)
+EXPLAIN (COSTS OFF)
+SELECT value
+FROM tpairs
+WHERE
+value > 10 AND
+(
+   key = 'foo' OR
+   key = 'bar' OR
+   key = 'baz'
+);
+
+-- Variant 2 (works only with this patch)
+EXPLAIN (COSTS OFF)
+SELECT value
+FROM tpairs
+WHERE
+value > 10 AND
+key IN ('foo', 'bar', 'baz');
+
+-- Variant 3 (works only with this patch)
+EXPLAIN (COSTS OFF)
+SELECT value
+FROM tpairs
+WHERE
+value > 10 AND
+key = ANY('{foo,bar,baz}');
+
+-- Clean up
+DROP TABLE tpairs;
diff --git a/src/test/regress/sql/saop_bitmap-2.sql b/src/test/regress/sql/saop_bitmap-2.sql
new file mode 100644
index 00000000000..12c9f7e3db9
--- /dev/null
+++ b/src/test/regress/sql/saop_bitmap-2.sql
@@ -0,0 +1,53 @@
+-- Test ScalarArrayOpExpr for partial indexes
+
+-- Generate enough data that we can test the use of an index
+CREATE TABLE ipairs(
+   key INT,
+   value INT
+)
+WITH (autovacuum_enabled = false);
+-- Install the partial index on these values
+CREATE INDEX ON ipairs(key) WHERE key = 7;
+CREATE INDEX ON ipairs(key) WHERE key = 8;
+CREATE INDEX ON ipairs(key) WHERE key = 9;
+
+-- Dummy data to surround ours
+INSERT INTO ipairs(key, value)
+SELECT n, -n
+FROM generate_series(1, 10000) AS tmp(n);
+
+-- Update statistics
+VACUUM (ANALYZE) ipairs;
+
+-- Our specific test data to force the partial index
+EXPLAIN (ANALYZE, COSTS OFF, BUFFERS OFF, SUMMARY OFF, TIMING OFF)
+UPDATE ipairs
+SET value = 0
+WHERE key = ANY('{7,8,9}'::INT[]);
+
+-- We want ensure that an index use is always attempted - just for the test
+SET enable_seqscan=false;
+
+-- Variant 1 (works without this patch)
+EXPLAIN (COSTS OFF)
+SELECT value=0
+FROM ipairs
+WHERE
+   key = 7 OR
+   key = 8 OR
+   key = 9;
+
+-- Variant 2 (works only with this patch)
+EXPLAIN (COSTS OFF)
+SELECT value=0
+FROM ipairs
+WHERE key IN (7, 8, 9);
+
+-- Variant 3 (works only with this patch)
+EXPLAIN (COSTS OFF)
+SELECT value=0
+FROM ipairs
+WHERE key = ANY('{7,8,9}'::INT[]);
+
+-- Clean up
+DROP TABLE ipairs;