get_partitions_excluded_by.patch

diff --git a/src/backend/catalog/partition.c b/src/backend/catalog/partition.c
index 0d9c774..57ea4a7 100644
--- a/src/backend/catalog/partition.c
+++ b/src/backend/catalog/partition.c
@@ -291,8 +291,8 @@ PG_FUNCTION_INFO_V1(satisfies_hash_partition);
 
 static Bitmapset *get_partitions_from_clauses_recurse(Relation relation,
 								int rt_index, List *clauses);
-static Bitmapset *get_partitions_from_ne_clauses(Relation relation,
-								List *ne_clauses);
+static Bitmapset *get_partitions_excluded_by(Relation relation,
+						   List *ne_clauses);
 static Bitmapset *get_partitions_from_or_clause_args(Relation relation,
 								int rt_index, List *or_clause_args);
 static bool classify_partition_bounding_keys(Relation relation, List *clauses,
@@ -1787,15 +1787,10 @@ get_partitions_from_clauses_recurse(Relation relation, int rt_index,
 	{
 		Bitmapset *ne_clause_parts;
 
-		ne_clause_parts = get_partitions_from_ne_clauses(relation, ne_clauses);
+		ne_clause_parts = get_partitions_excluded_by(relation, ne_clauses);
 
-		/*
-		 * Clauses in ne_clauses are in conjunction with the clauses that
-		 * selected the partitions contained in result, so combine the
-		 * partitions thus selected with those in result using set
-		 * intersection.
-		 */
-		result = bms_int_members(result, ne_clause_parts);
+		/* Remove any matched partitions */
+		result = bms_del_members(result, ne_clause_parts);
 		bms_free(ne_clause_parts);
 	}
 
@@ -1825,151 +1820,105 @@ get_partitions_from_clauses_recurse(Relation relation, int rt_index,
 		(0 == DatumGetInt32(FunctionCall2Coll(&partkey->partsupfunc[0],\
 											  partkey->partcollation[0],\
 											  (d1), (d2))))
-/*
- * Check if d is equal to some member of darray where equality is determined
- * by the partitioning comparison function.
- */
-static bool
-datum_in_array(PartitionKey partkey, Datum d, Datum *darray, int n)
-{
-	int		i;
-
-	if (darray == NULL || n == 0)
-		return false;
-
-	for (i = 0; i < n; i++)
-		if (partkey_datums_equal(d, darray[i]))
-			return true;
-
-	return false;
-}
-
-/*
- * count_partition_datums
- *
- * Returns the number of non-null datums allowed by a non-default list
- * partition with given index.
- */
-static int
-count_partition_datums(Relation rel, int index)
-{
-	PartitionBoundInfo boundinfo = RelationGetPartitionDesc(rel)->boundinfo;
-	int		i,
-			result = 0;
-
-	Assert(index != boundinfo->default_index);
-
-	/*
-	 * The answer is as many as the count of occurrence of the value index
-	 * in boundinfo->indexes[].
-	 */
-	for (i = 0; i < boundinfo->ndatums; i++)
-		if (index == boundinfo->indexes[i])
-			result += 1;
-
-	return result;
-}
 
 /*
- * get_partitions_from_ne_clauses
+ * get_partitions_excluded_by
  *
- * Return partitions of relation that satisfy all <> operator clauses in
- * ne_clauses.  Only ever called if relation is a list partitioned table.
+ * Returns a Bitmapset of partition indexes of any partition that can safely
+ * be removed due to 'ne_clauses' containing not-equal clauses for all
+ * possible values that the partition can contain.
  */
 static Bitmapset *
-get_partitions_from_ne_clauses(Relation relation, List *ne_clauses)
+get_partitions_excluded_by(Relation relation, List *ne_clauses)
 {
-	ListCell   *lc;
-	Bitmapset  *result,
-			   *excluded_parts;
-	PartitionKey partkey = RelationGetPartitionKey(relation);
-	PartitionDesc partdesc = RelationGetPartitionDesc(relation);
+	ListCell	   *lc;
+	Bitmapset	   *excluded_parts = NULL;
+	Bitmapset	   *foundoffsets = NULL;
+	PartitionKey	partkey = RelationGetPartitionKey(relation);
+	PartitionDesc	partdesc = RelationGetPartitionDesc(relation);
 	PartitionBoundInfo boundinfo = partdesc->boundinfo;
-	Datum  *exclude_datums;
-	int	   *count_excluded,
-			n_exclude_datums,
-			i;
+	PartitionBoundCmpArg arg;
+	int			   *datums_in_part;
+	int			   *datums_found;
+	int				i;
 
 	Assert(partkey->strategy == PARTITION_STRATEGY_LIST);
+	Assert(partkey->partnatts == 1);
 
-	/*
-	 * How this works:
-	 *
-	 * For each constant expression, we look up the partition that would
-	 * contain its value and mark the same as excluded partition.  After
-	 * doing the same for all clauses we'll have set of partitions that
-	 * are excluded.  For each excluded partition, check if there exist
-	 * values that it allows but are not specified in the clauses, if so
-	 * the partition won't actually be excluded.
-	 */
+	memset(&arg, 0, sizeof(arg));
 
-	/* De-duplicate constant values. */
-	exclude_datums = (Datum *) palloc0(list_length(ne_clauses) *
-									   sizeof(Datum));
-	n_exclude_datums = 0;
+	/* build a Bitmapset to record the offsets of all datums found */
 	foreach(lc, ne_clauses)
 	{
-		PartClause *pc = lfirst(lc);
+		PartClause *pc = (PartClause *) lfirst(lc);
 		Datum	datum;
 
-		if (partkey_datum_from_expr(partkey, 0, pc->constarg, &datum) &&
-			!datum_in_array(partkey, datum, exclude_datums, n_exclude_datums))
-			exclude_datums[n_exclude_datums++] = datum;
+		if (partkey_datum_from_expr(partkey, 0, pc->constarg, &datum))
+		{
+			int		offset;
+			bool	is_equal;
+
+			arg.datums = &datum;
+			arg.ndatums = 1;
+			offset = partition_bound_bsearch(partkey, boundinfo, &arg,
+											 &is_equal);
+
+			if (offset >= 0 && is_equal && boundinfo->indexes[offset] >= 0)
+				foundoffsets = bms_add_member(foundoffsets, offset);
+		}
 	}
 
+	/* No partitions can be excluded if we found no valid offsets above */
+	if (bms_is_empty(foundoffsets))
+		return NULL;
+
 	/*
-	 * For each value, if it's found in boundinfo, increment the count of its
-	 * partition as excluded due to that value.
+	 * Since each list partition can have multiple values in the IN clause, we
+	 * must ensure that we got all values in that clause before we can
+	 * eliminate the entire partition.
+	 *
+	 * We'll need two arrays for this, one to count the number of unique
+	 * datums we found, and another to record the number of datums permitted
+	 * in each partition.  Once we've counted all this, we can eliminate any
+	 * partition where the number of datums found match the number of datums
+	 * allowed in the partition.
 	 */
-	count_excluded = (int *) palloc0(partdesc->nparts * sizeof(int));
-	for (i = 0; i < n_exclude_datums; i++)
-	{
-		int		offset,
-				excluded_part;
-		bool	is_equal;
-		PartitionBoundCmpArg arg;
-		Datum   argdatums[] = {exclude_datums[i]};
-
-		memset(&arg, 0, sizeof(arg));
-		arg.datums = argdatums;
-		arg.ndatums = 1;
-		offset = partition_bound_bsearch(partkey, boundinfo, &arg, &is_equal);
-		if (offset >= 0 && is_equal && boundinfo->indexes[offset] >= 0)
-		{
-			excluded_part = boundinfo->indexes[offset];
-			count_excluded[excluded_part]++;
-		}
-	}
+	datums_in_part = (int *) palloc0(sizeof(int) * partdesc->nparts);
+	datums_found = (int *) palloc0(sizeof(int) * partdesc->nparts);
 
-	excluded_parts = NULL;
+	i = -1;
+	while ((i = bms_next_member(foundoffsets, i)) >= 0)
+		datums_found[boundinfo->indexes[i]]++;
+
+	/* Now, in a single pass of the partitions, count the datums it permits */
+	for (i = 0; i < boundinfo->ndatums; i++)
+		datums_in_part[boundinfo->indexes[i]]++;
+
+	/*
+	 * Compare the counts, eliminate any partition that we found clause for
+	 * all possible values. We must be careful here not to include any default
+	 * partition.  In this case both arrays will contain zero, so we can just
+	 * simply ensure we only eliminate when we found at least 1 datum.
+	 */
 	for (i = 0; i < partdesc->nparts; i++)
 	{
-		/*
-		 * If all datums of this partition appeared in ne_clauses, exclude
-		 * this partition.
-		 */
-		if (count_excluded[i] > 0 &&
-			count_excluded[i] == count_partition_datums(relation, i))
+		if (datums_found[i] >= datums_in_part[i] && datums_found[i] > 0)
 			excluded_parts = bms_add_member(excluded_parts, i);
 	}
 
 	/*
-	 * Also, exclude the "null-only" partition, because strict clauses in
-	 * ne_clauses will not select any rows from it.
+	 * Because the above clauses are strict, we can also exclude the NULL
+	 * partition providing it does not also allow non-NULL values.
 	 */
 	if (partition_bound_accepts_nulls(boundinfo) &&
-		count_partition_datums(relation, boundinfo->null_index) == 0)
+		datums_in_part[boundinfo->null_index] == 0)
 		excluded_parts = bms_add_member(excluded_parts,
 										boundinfo->null_index);
 
-	pfree(count_excluded);
-	pfree(exclude_datums);
-
-	result = bms_add_range(NULL, 0, partdesc->nparts - 1);
-	result = bms_del_members(result, excluded_parts);
-	bms_free(excluded_parts);
+	pfree(datums_in_part);
+	pfree(datums_found);
 
-	return result;
+	return excluded_parts;
 }
 
 /*
@@ -2251,7 +2200,7 @@ classify_partition_bounding_keys(Relation relation, List *clauses,
 				/*
 				 * We don't turn a <> operator clause into a key right away.
 				 * Instead, the caller will hand over such clauses to
-				 * get_partitions_from_ne_clauses().
+				 * get_partitions_excluded_by().
 				 */
 				if (is_ne_listp)
 					*ne_clauses = lappend(*ne_clauses, pc);