TABLESAMPLE, SQL Standard and extensible

Add a TABLESAMPLE clause to SELECT statements that allows
user to specify random BERNOULLI sampling or block level
SYSTEM sampling. Implementation allows for extensible
sampling functions to be written, using a standard API.
Basic version follows SQLStandard exactly. Usable
concrete use cases for the sampling API follow in later
commits.

Petr Jelinek

Reviewed by Michael Paquier and Simon Riggs

3 files changed, 233 insertions, 2 deletions

diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 2dce87879e7..14397830686 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -454,6 +454,7 @@ static Node *makeRecursiveViewSelect(char *relname, List *aliases, Node *query);
 %type <range>	relation_expr
 %type <range>	relation_expr_opt_alias
 %type <target>	target_el single_set_clause set_target insert_column_item
+%type <node>	relation_expr_tablesample tablesample_clause opt_repeatable_clause
 
 %type <str>		generic_option_name
 %type <node>	generic_option_arg
@@ -625,8 +626,8 @@ static Node *makeRecursiveViewSelect(char *relname, List *aliases, Node *query);
 	STATEMENT STATISTICS STDIN STDOUT STORAGE STRICT_P STRIP_P SUBSTRING
 	SYMMETRIC SYSID SYSTEM_P
 
-	TABLE TABLES TABLESPACE TEMP TEMPLATE TEMPORARY TEXT_P THEN TIME TIMESTAMP
-	TO TRAILING TRANSACTION TRANSFORM TREAT TRIGGER TRIM TRUE_P
+	TABLE TABLES TABLESAMPLE TABLESPACE TEMP TEMPLATE TEMPORARY TEXT_P THEN
+	TIME TIMESTAMP TO TRAILING TRANSACTION TRANSFORM TREAT TRIGGER TRIM TRUE_P
 	TRUNCATE TRUSTED TYPE_P TYPES_P
 
 	UNBOUNDED UNCOMMITTED UNENCRYPTED UNION UNIQUE UNKNOWN UNLISTEN UNLOGGED
@@ -10386,6 +10387,10 @@ table_ref:	relation_expr opt_alias_clause
 					$1->alias = $2;
 					$$ = (Node *) $1;
 				}
+			| relation_expr_tablesample
+				{
+					$$ = (Node *) $1;
+				}
 			| func_table func_alias_clause
 				{
 					RangeFunction *n = (RangeFunction *) $1;
@@ -10711,6 +10716,32 @@ relation_expr_opt_alias: relation_expr					%prec UMINUS
 				}
 		;
 
+
+relation_expr_tablesample: relation_expr opt_alias_clause tablesample_clause
+				{
+					RangeTableSample *n = (RangeTableSample *) $3;
+					n->relation = $1;
+					n->relation->alias = $2;
+					$$ = (Node *) n;
+				}
+		;
+
+tablesample_clause:
+			TABLESAMPLE ColId '(' expr_list ')' opt_repeatable_clause
+				{
+					RangeTableSample *n = makeNode(RangeTableSample);
+					n->method = $2;
+					n->args = $4;
+					n->repeatable = $6;
+					$$ = (Node *) n;
+				}
+		;
+
+opt_repeatable_clause:
+			REPEATABLE '(' a_expr ')'	{ $$ = (Node *) $3; }
+			| /*EMPTY*/					{ $$ = NULL; }
+		;
+
 /*
  * func_table represents a function invocation in a FROM list. It can be
  * a plain function call, like "foo(...)", or a ROWS FROM expression with
@@ -13804,6 +13835,7 @@ type_func_name_keyword:
 			| OVERLAPS
 			| RIGHT
 			| SIMILAR
+			| TABLESAMPLE
 			| VERBOSE
 		;
 
diff --git a/src/backend/parser/parse_clause.c b/src/backend/parser/parse_clause.c
index 73c505ed85b..6b1bbe57d0e 100644
--- a/src/backend/parser/parse_clause.c
+++ b/src/backend/parser/parse_clause.c
@@ -17,6 +17,7 @@
 
 #include "access/heapam.h"
 #include "catalog/catalog.h"
+#include "access/htup_details.h"
 #include "catalog/heap.h"
 #include "catalog/pg_constraint.h"
 #include "catalog/pg_type.h"
@@ -31,6 +32,7 @@
 #include "parser/parse_coerce.h"
 #include "parser/parse_collate.h"
 #include "parser/parse_expr.h"
+#include "parser/parse_func.h"
 #include "parser/parse_oper.h"
 #include "parser/parse_relation.h"
 #include "parser/parse_target.h"
@@ -39,6 +41,7 @@
 #include "utils/guc.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
+#include "utils/syscache.h"
 
 
 /* Convenience macro for the most common makeNamespaceItem() case */
@@ -419,6 +422,39 @@ transformJoinOnClause(ParseState *pstate, JoinExpr *j, List *namespace)
 	return result;
 }
 
+static RangeTblEntry *
+transformTableSampleEntry(ParseState *pstate, RangeTableSample *rv)
+{
+	RangeTblEntry   *rte = NULL;
+	CommonTableExpr *cte = NULL;
+	TableSampleClause *tablesample = NULL;
+
+	/* if relation has an unqualified name, it might be a CTE reference */
+	if (!rv->relation->schemaname)
+	{
+		Index	levelsup;
+		cte = scanNameSpaceForCTE(pstate, rv->relation->relname, &levelsup);
+	}
+
+	/* We first need to build a range table entry */
+	if (!cte)
+		rte = transformTableEntry(pstate, rv->relation);
+
+	if (!rte ||
+		(rte->relkind != RELKIND_RELATION &&
+		rte->relkind != RELKIND_MATVIEW))
+		ereport(ERROR,
+				(errcode(ERRCODE_SYNTAX_ERROR),
+				 errmsg("TABLESAMPLE clause can only be used on tables and materialized views"),
+				 parser_errposition(pstate, rv->relation->location)));
+
+	tablesample = ParseTableSample(pstate, rv->method, rv->repeatable,
+								   rv->args, rv->relation->location);
+	rte->tablesample = tablesample;
+
+	return rte;
+}
+
 /*
  * transformTableEntry --- transform a RangeVar (simple relation reference)
  */
@@ -1127,6 +1163,26 @@ transformFromClauseItem(ParseState *pstate, Node *n,
 
 		return (Node *) j;
 	}
+	else if (IsA(n, RangeTableSample))
+	{
+		/* Tablesample reference */
+		RangeTableSample   *rv = (RangeTableSample *) n;
+		RangeTblRef *rtr;
+		RangeTblEntry *rte = NULL;
+		int			rtindex;
+
+		rte = transformTableSampleEntry(pstate, rv);
+
+		/* assume new rte is at end */
+		rtindex = list_length(pstate->p_rtable);
+		Assert(rte == rt_fetch(rtindex, pstate->p_rtable));
+		*top_rte = rte;
+		*top_rti = rtindex;
+		*namespace = list_make1(makeDefaultNSItem(rte));
+		rtr = makeNode(RangeTblRef);
+		rtr->rtindex = rtindex;
+		return (Node *) rtr;
+	}
 	else
 		elog(ERROR, "unrecognized node type: %d", (int) nodeTag(n));
 	return NULL;				/* can't get here, keep compiler quiet */
diff --git a/src/backend/parser/parse_func.c b/src/backend/parser/parse_func.c
index f7affebf84e..fa50f92d8dd 100644
--- a/src/backend/parser/parse_func.c
+++ b/src/backend/parser/parse_func.c
@@ -18,6 +18,7 @@
 #include "catalog/pg_aggregate.h"
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
+#include "catalog/pg_tablesample_method.h"
 #include "funcapi.h"
 #include "lib/stringinfo.h"
 #include "nodes/makefuncs.h"
@@ -26,6 +27,7 @@
 #include "parser/parse_clause.h"
 #include "parser/parse_coerce.h"
 #include "parser/parse_func.h"
+#include "parser/parse_expr.h"
 #include "parser/parse_relation.h"
 #include "parser/parse_target.h"
 #include "parser/parse_type.h"
@@ -767,6 +769,147 @@ ParseFuncOrColumn(ParseState *pstate, List *funcname, List *fargs,
 }
 
 
+/*
+ * ParseTableSample
+ *
+ * Parse TABLESAMPLE clause and process the arguments
+ */
+TableSampleClause *
+ParseTableSample(ParseState *pstate, char *samplemethod, Node *repeatable,
+				 List *sampleargs, int location)
+{
+	HeapTuple		tuple;
+	Form_pg_tablesample_method tsm;
+	Form_pg_proc procform;
+	TableSampleClause *tablesample;
+	List		   *fargs;
+	ListCell	   *larg;
+	int				nargs, initnargs;
+	Oid				init_arg_types[FUNC_MAX_ARGS];
+
+	/* Load the tablesample method */
+	tuple = SearchSysCache1(TABLESAMPLEMETHODNAME, PointerGetDatum(samplemethod));
+	if (!HeapTupleIsValid(tuple))
+		ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_OBJECT),
+				 errmsg("tablesample method \"%s\" does not exist",
+						samplemethod),
+				 parser_errposition(pstate, location)));
+
+	tablesample = makeNode(TableSampleClause);
+	tablesample->tsmid = HeapTupleGetOid(tuple);
+
+	tsm = (Form_pg_tablesample_method) GETSTRUCT(tuple);
+
+	tablesample->tsmseqscan = tsm->tsmseqscan;
+	tablesample->tsmpagemode = tsm->tsmpagemode;
+	tablesample->tsminit = tsm->tsminit;
+	tablesample->tsmnextblock = tsm->tsmnextblock;
+	tablesample->tsmnexttuple = tsm->tsmnexttuple;
+	tablesample->tsmexaminetuple = tsm->tsmexaminetuple;
+	tablesample->tsmend = tsm->tsmend;
+	tablesample->tsmreset = tsm->tsmreset;
+	tablesample->tsmcost = tsm->tsmcost;
+
+	ReleaseSysCache(tuple);
+
+	/* Validate the parameters against init function definition. */
+	tuple = SearchSysCache1(PROCOID,
+							ObjectIdGetDatum(tablesample->tsminit));
+
+	if (!HeapTupleIsValid(tuple))	/* should not happen */
+		elog(ERROR, "cache lookup failed for function %u",
+			 tablesample->tsminit);
+
+	procform = (Form_pg_proc) GETSTRUCT(tuple);
+	initnargs = procform->pronargs;
+	Assert(initnargs >= 3);
+
+	/*
+	 * First parameter is used to pass the SampleScanState, second is
+	 * seed (REPEATABLE), skip the processing for them here, just assert
+	 * that the types are correct.
+	 */
+	Assert(procform->proargtypes.values[0] == INTERNALOID);
+	Assert(procform->proargtypes.values[1] == INT4OID);
+	initnargs -= 2;
+	memcpy(init_arg_types, procform->proargtypes.values + 2,
+				   initnargs * sizeof(Oid));
+
+	/* Now we are done with the catalog */
+	ReleaseSysCache(tuple);
+
+	/* Process repeatable (seed) */
+	if (repeatable != NULL)
+	{
+		Node   *arg = repeatable;
+
+		if (arg && IsA(arg, A_Const))
+		{
+			A_Const    *con = (A_Const *) arg;
+
+			if (con->val.type == T_Null)
+				ereport(ERROR,
+						(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+						 errmsg("REPEATABLE clause must be NOT NULL numeric value"),
+						 parser_errposition(pstate, con->location)));
+
+		}
+
+		arg = transformExpr(pstate, arg, EXPR_KIND_FROM_FUNCTION);
+		arg = coerce_to_specific_type(pstate, arg, INT4OID, "REPEATABLE");
+		tablesample->repeatable = arg;
+	}
+	else
+		tablesample->repeatable = NULL;
+
+	/* Check user provided expected number of arguments. */
+	if (list_length(sampleargs) != initnargs)
+		ereport(ERROR,
+				(errcode(ERRCODE_TOO_MANY_ARGUMENTS),
+		 errmsg_plural("tablesample method \"%s\" expects %d argument got %d",
+					   "tablesample method \"%s\" expects %d arguments got %d",
+					   initnargs,
+					   samplemethod,
+					   initnargs, list_length(sampleargs)),
+		 parser_errposition(pstate, location)));
+
+	/* Transform the arguments, typecasting them as needed. */
+	fargs = NIL;
+	nargs = 0;
+	foreach(larg, sampleargs)
+	{
+		Node   *inarg = (Node *) lfirst(larg);
+		Node   *arg = transformExpr(pstate, inarg, EXPR_KIND_FROM_FUNCTION);
+		Oid		argtype = exprType(arg);
+
+		if (argtype != init_arg_types[nargs])
+		{
+			if (!can_coerce_type(1, &argtype, &init_arg_types[nargs],
+								 COERCION_IMPLICIT))
+				ereport(ERROR,
+						(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("wrong parameter %d for tablesample method \"%s\"",
+							nargs + 1, samplemethod),
+					 errdetail("Expected type %s got %s.",
+							   format_type_be(init_arg_types[nargs]),
+							   format_type_be(argtype)),
+					 parser_errposition(pstate, exprLocation(inarg))));
+
+			arg = coerce_type(pstate, arg, argtype, init_arg_types[nargs], -1,
+							  COERCION_IMPLICIT, COERCE_IMPLICIT_CAST, -1);
+		}
+
+		fargs = lappend(fargs, arg);
+		nargs++;
+	}
+
+	/* Pass the arguments down */
+	tablesample->args = fargs;
+
+	return tablesample;
+}
+
 /* func_match_argtypes()
  *
  * Given a list of candidate functions (having the right name and number