Add hooks to let plugins override the planner's lookups in pg_statistic.

Simon Riggs, with some editorialization by me.

4 files changed, 105 insertions, 19 deletions

diff --git a/src/backend/utils/adt/selfuncs.c b/src/backend/utils/adt/selfuncs.c
index 4ad58e56ab..171bb0f091 100644
--- a/src/backend/utils/adt/selfuncs.c
+++ b/src/backend/utils/adt/selfuncs.c
@@ -119,6 +119,10 @@
 #include "utils/syscache.h"
 
 
+/* Hooks for plugins to get control when we ask for stats */
+get_relation_stats_hook_type get_relation_stats_hook = NULL;
+get_index_stats_hook_type get_index_stats_hook = NULL;
+
 static double var_eq_const(VariableStatData *vardata, Oid operator,
 			 Datum constval, bool constisnull,
 			 bool varonleft);
@@ -2935,7 +2939,7 @@ estimate_num_groups(PlannerInfo *root, List *groupExprs, double input_rows)
 		 * complicated.
 		 */
 		examine_variable(root, groupexpr, 0, &vardata);
-		if (vardata.statsTuple != NULL || vardata.isunique)
+		if (HeapTupleIsValid(vardata.statsTuple) || vardata.isunique)
 		{
 			varinfos = add_unique_group_var(root, varinfos,
 											groupexpr, &vardata);
@@ -3942,6 +3946,7 @@ get_join_variables(PlannerInfo *root, List *args, SpecialJoinInfo *sjinfo,
  *		subquery, not one in the current query).
  *	statsTuple: the pg_statistic entry for the variable, if one exists;
  *		otherwise NULL.
+ *	freefunc: pointer to a function to release statsTuple with.
  *	vartype: exposed type of the expression; this should always match
  *		the declared input type of the operator we are estimating for.
  *	atttype, atttypmod: type data to pass to get_attstatsslot().  This is
@@ -3986,7 +3991,18 @@ examine_variable(PlannerInfo *root, Node *node, int varRelid,
 
 		rte = root->simple_rte_array[var->varno];
 
-		if (rte->inh)
+		if (get_relation_stats_hook &&
+			(*get_relation_stats_hook) (root, rte, var->varattno, vardata))
+		{
+			/*
+			 * The hook took control of acquiring a stats tuple.  If it
+			 * did supply a tuple, it'd better have supplied a freefunc.
+			 */
+			if (HeapTupleIsValid(vardata->statsTuple) &&
+				!vardata->freefunc)
+				elog(ERROR, "no function provided to release variable stats with");
+		}
+		else if (rte->inh)
 		{
 			/*
 			 * XXX This means the Var represents a column of an append
@@ -4000,6 +4016,7 @@ examine_variable(PlannerInfo *root, Node *node, int varRelid,
 												 ObjectIdGetDatum(rte->relid),
 												 Int16GetDatum(var->varattno),
 												 0, 0);
+			vardata->freefunc = ReleaseSysCache;
 		}
 		else
 		{
@@ -4116,10 +4133,28 @@ examine_variable(PlannerInfo *root, Node *node, int varRelid,
 							index->indpred == NIL)
 							vardata->isunique = true;
 						/* Has it got stats? */
-						vardata->statsTuple = SearchSysCache(STATRELATT,
-										   ObjectIdGetDatum(index->indexoid),
-													  Int16GetDatum(pos + 1),
-															 0, 0);
+						if (get_index_stats_hook &&
+							(*get_index_stats_hook) (root, index->indexoid,
+													 pos + 1, vardata))
+						{
+							/*
+							 * The hook took control of acquiring a stats
+							 * tuple.  If it did supply a tuple, it'd better
+							 * have supplied a freefunc.
+							 */
+							if (HeapTupleIsValid(vardata->statsTuple) &&
+								!vardata->freefunc)
+								elog(ERROR, "no function provided to release variable stats with");
+						}
+						else
+						{
+							vardata->statsTuple =
+								SearchSysCache(STATRELATT,
+											   ObjectIdGetDatum(index->indexoid),
+											   Int16GetDatum(pos + 1),
+											   0, 0);
+							vardata->freefunc = ReleaseSysCache;
+						}
 						if (vardata->statsTuple)
 							break;
 					}
@@ -5551,7 +5586,7 @@ btcostestimate(PG_FUNCTION_ARGS)
 	double	   *indexCorrelation = (double *) PG_GETARG_POINTER(7);
 	Oid			relid;
 	AttrNumber	colnum;
-	HeapTuple	tuple;
+	VariableStatData vardata;
 	double		numIndexTuples;
 	List	   *indexBoundQuals;
 	int			indexcol;
@@ -5756,17 +5791,34 @@ btcostestimate(PG_FUNCTION_ARGS)
 		colnum = 1;
 	}
 
-	tuple = SearchSysCache(STATRELATT,
-						   ObjectIdGetDatum(relid),
-						   Int16GetDatum(colnum),
-						   0, 0);
+	MemSet(&vardata, 0, sizeof(vardata));
 
-	if (HeapTupleIsValid(tuple))
+	if (get_index_stats_hook &&
+		(*get_index_stats_hook) (root, relid, colnum, &vardata))
+	{
+		/*
+		 * The hook took control of acquiring a stats tuple.  If it did supply
+		 * a tuple, it'd better have supplied a freefunc.
+		 */
+		if (HeapTupleIsValid(vardata.statsTuple) &&
+			!vardata.freefunc)
+			elog(ERROR, "no function provided to release variable stats with");
+	}
+	else
+	{
+		vardata.statsTuple = SearchSysCache(STATRELATT,
+											ObjectIdGetDatum(relid),
+											Int16GetDatum(colnum),
+											0, 0);
+		vardata.freefunc = ReleaseSysCache;
+	}
+
+	if (HeapTupleIsValid(vardata.statsTuple))
 	{
 		float4	   *numbers;
 		int			nnumbers;
 
-		if (get_attstatsslot(tuple, InvalidOid, 0,
+		if (get_attstatsslot(vardata.statsTuple, InvalidOid, 0,
 							 STATISTIC_KIND_CORRELATION,
 							 index->fwdsortop[0],
 							 NULL, NULL, &numbers, &nnumbers))
@@ -5783,7 +5835,7 @@ btcostestimate(PG_FUNCTION_ARGS)
 
 			free_attstatsslot(InvalidOid, NULL, 0, numbers, nnumbers);
 		}
-		else if (get_attstatsslot(tuple, InvalidOid, 0,
+		else if (get_attstatsslot(vardata.statsTuple, InvalidOid, 0,
 								  STATISTIC_KIND_CORRELATION,
 								  index->revsortop[0],
 								  NULL, NULL, &numbers, &nnumbers))
@@ -5800,9 +5852,10 @@ btcostestimate(PG_FUNCTION_ARGS)
 
 			free_attstatsslot(InvalidOid, NULL, 0, numbers, nnumbers);
 		}
-		ReleaseSysCache(tuple);
 	}
 
+	ReleaseVariableStats(vardata);
+
 	PG_RETURN_VOID();
 }
 
diff --git a/src/backend/utils/cache/lsyscache.c b/src/backend/utils/cache/lsyscache.c
index 1328ae3053..e87c3edee3 100644
--- a/src/backend/utils/cache/lsyscache.c
+++ b/src/backend/utils/cache/lsyscache.c
@@ -35,6 +35,9 @@
 #include "utils/lsyscache.h"
 #include "utils/syscache.h"
 
+/* Hook for plugins to get control in get_attavgwidth() */
+get_attavgwidth_hook_type get_attavgwidth_hook = NULL;
+
 
 /*				---------- AMOP CACHES ----------						 */
 
@@ -2492,20 +2495,30 @@ get_typmodout(Oid typid)
  *
  *	  Given the table and attribute number of a column, get the average
  *	  width of entries in the column.  Return zero if no data available.
+ *
+ * Calling a hook at this point looks somewhat strange, but is required
+ * because the optimizer calls this function without any other way for
+ * plug-ins to control the result.
  */
 int32
 get_attavgwidth(Oid relid, AttrNumber attnum)
 {
 	HeapTuple	tp;
+	int32		stawidth;
 
+	if (get_attavgwidth_hook)
+	{
+		stawidth = (*get_attavgwidth_hook) (relid, attnum);
+		if (stawidth > 0)
+			return stawidth;
+	}
 	tp = SearchSysCache(STATRELATT,
 						ObjectIdGetDatum(relid),
 						Int16GetDatum(attnum),
 						0, 0);
 	if (HeapTupleIsValid(tp))
 	{
-		int32		stawidth = ((Form_pg_statistic) GETSTRUCT(tp))->stawidth;
-
+		stawidth = ((Form_pg_statistic) GETSTRUCT(tp))->stawidth;
 		ReleaseSysCache(tp);
 		if (stawidth > 0)
 			return stawidth;
@@ -2523,6 +2536,9 @@ get_attavgwidth(Oid relid, AttrNumber attnum)
  * already-looked-up tuple in the pg_statistic cache.  We do this since
  * most callers will want to extract more than one value from the cache
  * entry, and we don't want to repeat the cache lookup unnecessarily.
+ * Also, this API allows this routine to be used with statistics tuples
+ * that have been provided by a stats hook and didn't really come from
+ * pg_statistic.
  *
  * statstuple: pg_statistics tuple to be examined.
  * atttype: type OID of attribute (can be InvalidOid if values == NULL).
diff --git a/src/include/utils/lsyscache.h b/src/include/utils/lsyscache.h
index ec589dfb67..9f638967d4 100644
--- a/src/include/utils/lsyscache.h
+++ b/src/include/utils/lsyscache.h
@@ -26,6 +26,10 @@ typedef enum IOFuncSelector
 	IOFunc_send
 } IOFuncSelector;
 
+/* Hook for plugins to get control in get_attavgwidth() */
+typedef int32 (*get_attavgwidth_hook_type) (Oid relid, AttrNumber attnum);
+extern PGDLLIMPORT get_attavgwidth_hook_type get_attavgwidth_hook;
+
 extern bool op_in_opfamily(Oid opno, Oid opfamily);
 extern int	get_op_opfamily_strategy(Oid opno, Oid opfamily);
 extern void get_op_opfamily_properties(Oid opno, Oid opfamily,
diff --git a/src/include/utils/selfuncs.h b/src/include/utils/selfuncs.h
index f2bf161c58..d4dd11b45a 100644
--- a/src/include/utils/selfuncs.h
+++ b/src/include/utils/selfuncs.h
@@ -64,12 +64,13 @@
 
 
 /* Return data from examine_variable and friends */
-typedef struct
+typedef struct VariableStatData
 {
 	Node	   *var;			/* the Var or expression tree */
 	RelOptInfo *rel;			/* Relation, or NULL if not identifiable */
 	HeapTuple	statsTuple;		/* pg_statistic tuple, or NULL if none */
 	/* NB: if statsTuple!=NULL, it must be freed when caller is done */
+	void		(*freefunc) (HeapTuple tuple);	/* how to free statsTuple */
 	Oid			vartype;		/* exposed type of expression */
 	Oid			atttype;		/* type to pass to get_attstatsslot */
 	int32		atttypmod;		/* typmod to pass to get_attstatsslot */
@@ -79,7 +80,7 @@ typedef struct
 #define ReleaseVariableStats(vardata)  \
 	do { \
 		if (HeapTupleIsValid((vardata).statsTuple)) \
-			ReleaseSysCache((vardata).statsTuple); \
+			(* (vardata).freefunc) ((vardata).statsTuple); \
 	} while(0)
 
 
@@ -97,6 +98,18 @@ typedef enum
 
 /* selfuncs.c */
 
+/* Hooks for plugins to get control when we ask for stats */
+typedef bool (*get_relation_stats_hook_type) (PlannerInfo *root,
+											  RangeTblEntry *rte,
+											  AttrNumber attnum,
+											  VariableStatData *vardata);
+extern PGDLLIMPORT get_relation_stats_hook_type get_relation_stats_hook;
+typedef bool (*get_index_stats_hook_type) (PlannerInfo *root,
+										   Oid indexOid,
+										   AttrNumber indexattnum,
+										   VariableStatData *vardata);
+extern PGDLLIMPORT get_index_stats_hook_type get_index_stats_hook;
+
 extern void examine_variable(PlannerInfo *root, Node *node, int varRelid,
 				 VariableStatData *vardata);
 extern bool get_restriction_variable(PlannerInfo *root, List *args,