Re: [HACKERS] extend pgbench expressions with functions

[Fabien COELHO]

This patch extends pgbench expression with functions. Currently only one "abs" function is added. The point is rather to bootstrap the infrastructure for other functions (such as hash, random variants...) to be added later.

Here is a small v2 update:
 - with better error messages on non existing functions
 - a minimal documentation

--
Fabien.
diff --git a/contrib/pgbench/exprparse.y b/contrib/pgbench/exprparse.y
index 243c6b9..254f081 100644
--- a/contrib/pgbench/exprparse.y
+++ b/contrib/pgbench/exprparse.y
@@ -20,6 +20,7 @@ static PgBenchExpr *make_integer_constant(int64 ival);
 static PgBenchExpr *make_variable(char *varname);
 static PgBenchExpr *make_op(char operator, PgBenchExpr *lexpr,
 		PgBenchExpr *rexpr);
+static PgBenchExpr *make_func(const char *fname, PgBenchExpr *arg1);
 
 %}
 
@@ -35,8 +36,8 @@ static PgBenchExpr *make_op(char operator, PgBenchExpr *lexpr,
 
 %type <expr> expr
 %type <ival> INTEGER
-%type <str> VARIABLE
-%token INTEGER VARIABLE
+%type <str> VARIABLE FUNCTION
+%token INTEGER VARIABLE FUNCTION
 %token CHAR_ERROR /* never used, will raise a syntax error */
 
 %left	'+' '-'
@@ -57,6 +58,7 @@ expr: '(' expr ')'			{ $$ = $2; }
 	| expr '%' expr			{ $$ = make_op('%', $1, $3); }
 	| INTEGER				{ $$ = make_integer_constant($1); }
 	| VARIABLE 				{ $$ = make_variable($1); }
+	| FUNCTION '(' expr ')' { $$ = make_func($1, $3); pg_free($1); }
 	;
 
 %%
@@ -93,4 +95,33 @@ make_op(char operator, PgBenchExpr *lexpr, PgBenchExpr *rexpr)
 	return expr;
 }
 
+static struct {
+	char * fname;
+	int nargs;
+	PgBenchFunction tag;
+} PGBENCH_FUNCTIONS[] = {
+	{ "abs", 1, PGBENCH_ABS }
+};
+
+static PgBenchExpr *
+make_func(const char * fname, PgBenchExpr *arg1)
+{
+	PgBenchExpr *expr = pg_malloc(sizeof(PgBenchExpr));
+	int nfunctions = sizeof(PGBENCH_FUNCTIONS) / sizeof(PGBENCH_FUNCTIONS[0]);
+	int i;
+
+	expr->etype = ENODE_FUNCTION;
+	expr->u.function.function = PGBENCH_NONE;
+
+	for (i = 0; i < nfunctions; i++)
+		if (pg_strcasecmp(fname, PGBENCH_FUNCTIONS[i].fname) == 0)
+			expr->u.function.function = PGBENCH_FUNCTIONS[i].tag;
+
+	if (expr->u.function.function == PGBENCH_NONE)
+		yyerror("unexpected function name");
+
+	expr->u.function.arg1 = arg1;
+	return expr;
+}
+
 #include "exprscan.c"
diff --git a/contrib/pgbench/exprscan.l b/contrib/pgbench/exprscan.l
index 4c9229c..a276444 100644
--- a/contrib/pgbench/exprscan.l
+++ b/contrib/pgbench/exprscan.l
@@ -46,6 +46,7 @@ space			[ \t\r\f]
 ")"				{ yycol += yyleng; return ')'; }
 :[a-zA-Z0-9_]+	{ yycol += yyleng; yylval.str = pg_strdup(yytext + 1); return VARIABLE; }
 [0-9]+			{ yycol += yyleng; yylval.ival = strtoint64(yytext); return INTEGER; }
+[a-zA-Z]+       { yycol += yyleng; yylval.str = pg_strdup(yytext); return FUNCTION; }
 
 [\n]			{ yycol = 0; yyline++; }
 {space}			{ yycol += yyleng; /* ignore */ }
diff --git a/contrib/pgbench/pgbench.c b/contrib/pgbench/pgbench.c
index 706fdf5..e0cffc2 100644
--- a/contrib/pgbench/pgbench.c
+++ b/contrib/pgbench/pgbench.c
@@ -959,6 +959,26 @@ evaluateExpr(CState *st, PgBenchExpr *expr, int64 *retval)
 				return false;
 			}
 
+		case ENODE_FUNCTION:
+			{
+				switch (expr->u.function.function)
+				{
+					case PGBENCH_ABS:
+					{
+						int64 arg1;
+						if (!evaluateExpr(st, expr->u.function.arg1, &arg1))
+							return false;
+
+						*retval = arg1 > 0? arg1: -arg1;
+						return true;
+					}
+					default:
+						fprintf(stderr, "bad function (%d)\n",
+								expr->u.function.function);
+						return false;
+				}
+			}
+
 		default:
 			break;
 	}
diff --git a/contrib/pgbench/pgbench.h b/contrib/pgbench/pgbench.h
index 128bf11..b7d30ba 100644
--- a/contrib/pgbench/pgbench.h
+++ b/contrib/pgbench/pgbench.h
@@ -15,12 +15,19 @@ typedef enum PgBenchExprType
 {
 	ENODE_INTEGER_CONSTANT,
 	ENODE_VARIABLE,
-	ENODE_OPERATOR
+	ENODE_OPERATOR,
+	ENODE_FUNCTION
 } PgBenchExprType;
 
 struct PgBenchExpr;
 typedef struct PgBenchExpr PgBenchExpr;
 
+typedef enum PgBenchFunction
+{
+	PGBENCH_NONE,
+	PGBENCH_ABS
+} PgBenchFunction;
+
 struct PgBenchExpr
 {
 	PgBenchExprType	etype;
@@ -40,6 +47,11 @@ struct PgBenchExpr
 			PgBenchExpr	*lexpr;
 			PgBenchExpr *rexpr;
 		} operator;
+		struct
+		{
+			PgBenchFunction function;
+			PgBenchExpr *arg1;
+		} function;
 	} u;
 };
 
diff --git a/doc/src/sgml/pgbench.sgml b/doc/src/sgml/pgbench.sgml
index ed12e27..0c58412 100644
--- a/doc/src/sgml/pgbench.sgml
+++ b/doc/src/sgml/pgbench.sgml
@@ -762,7 +762,7 @@ pgbench <optional> <replaceable>options</> </optional> <replaceable>dbname</>
       references to variables <literal>:</><replaceable>variablename</>,
       and expressions composed of unary (<literal>-</>) or binary operators
       (<literal>+</>, <literal>-</>, <literal>*</>, <literal>/</>, <literal>%</>)
-      with their usual associativity, and parentheses.
+      with their usual associativity, function <literal>abs</> and parentheses.
      </para>
 
      <para>

-- 
Sent via pgsql-hackers mailing list (pgsql-hackers@postgresql.org)
To make changes to your subscription:
http://www.postgresql.org/mailpref/pgsql-hackers