From 808ca8dcc72ce3aacd18653359e39b1dbcc45be5 Mon Sep 17 00:00:00 2001
From: Marko Kreen <markokr@gmail.com>
Date: Fri, 27 Jun 2008 08:55:19 +0000
Subject: [PATCH] accept int2 and int8 from hash function

---
 expected/plproxy_test.out | 46 +++++++++++++++++++++++++++++++++++++++
 sql/plproxy_test.sql      | 36 ++++++++++++++++++++++++++++++
 src/execute.c             | 20 ++++++++++++-----
 3 files changed, 96 insertions(+), 6 deletions(-)

diff --git a/expected/plproxy_test.out b/expected/plproxy_test.out
index 9cc596d..b79f5ca 100644
--- a/expected/plproxy_test.out
+++ b/expected/plproxy_test.out
@@ -264,3 +264,49 @@ select * from "testQuoting"('user', '1', 'dat');
      1 | BazOoka
 (1 row)
 
+-- test hash types function
+create or replace function t_hash16(int4) returns int2 as $$
+declare
+    res int2;
+begin
+    res = $1::int2;
+    return res;
+end;
+$$ language plpgsql;
+create or replace function t_hash64(int4) returns int8 as $$
+declare
+    res int8;
+begin
+    res = $1;
+    return res;
+end;
+$$ language plpgsql;
+create function test_hash16(id integer, data text)
+returns text as $$ cluster 'testcluster'; run on t_hash16(id); select data; $$ language plproxy;
+select * from test_hash16('0', 'hash16');
+ test_hash16 
+-------------
+ hash16
+(1 row)
+
+create function test_hash64(id integer, data text)
+returns text as $$ cluster 'testcluster'; run on t_hash64(id); select data; $$ language plproxy;
+select * from test_hash64('0', 'hash64');
+ test_hash64 
+-------------
+ hash64
+(1 row)
+
+-- test argument difference
+\c test_part
+create function test_difftypes(username text, out val1 int2, out val2 float8)
+as $$ begin val1 = 1; val2 = 3;return; end; $$ language plpgsql;
+\c regression
+create function test_difftypes(username text, out val1 int4, out val2 float4)
+as $$ cluster 'testcluster'; run on 0; $$ language plproxy;
+select * from test_difftypes('types');
+ val1 | val2 
+------+------
+    1 |    3
+(1 row)
+
diff --git a/sql/plproxy_test.sql b/sql/plproxy_test.sql
index 43cd7d4..5fb025b 100644
--- a/sql/plproxy_test.sql
+++ b/sql/plproxy_test.sql
@@ -162,3 +162,39 @@ returns "RetWeird" as $$ select 1::int4, 'BazOoka'::text $$ language sql;
 \c regression
 select * from "testQuoting"('user', '1', 'dat');
 
+-- test hash types function
+create or replace function t_hash16(int4) returns int2 as $$
+declare
+    res int2;
+begin
+    res = $1::int2;
+    return res;
+end;
+$$ language plpgsql;
+
+create or replace function t_hash64(int4) returns int8 as $$
+declare
+    res int8;
+begin
+    res = $1;
+    return res;
+end;
+$$ language plpgsql;
+
+create function test_hash16(id integer, data text)
+returns text as $$ cluster 'testcluster'; run on t_hash16(id); select data; $$ language plproxy;
+select * from test_hash16('0', 'hash16');
+
+create function test_hash64(id integer, data text)
+returns text as $$ cluster 'testcluster'; run on t_hash64(id); select data; $$ language plproxy;
+select * from test_hash64('0', 'hash64');
+
+-- test argument difference
+\c test_part
+create function test_difftypes(username text, out val1 int2, out val2 float8)
+as $$ begin val1 = 1; val2 = 3;return; end; $$ language plpgsql;
+\c regression
+create function test_difftypes(username text, out val1 int4, out val2 float4)
+as $$ cluster 'testcluster'; run on 0; $$ language plproxy;
+select * from test_difftypes('types');
+
diff --git a/src/execute.c b/src/execute.c
index 0e75c6a..5bba3a8 100644
--- a/src/execute.c
+++ b/src/execute.c
@@ -619,6 +619,7 @@ tag_hash_partitions(ProxyFunction *func, FunctionCallInfo fcinfo)
 {
 	int			i;
 	TupleDesc	desc;
+	Oid			htype;
 	ProxyCluster *cluster = func->cur_cluster;
 
 	/* execute cached plan */
@@ -626,22 +627,29 @@ tag_hash_partitions(ProxyFunction *func, FunctionCallInfo fcinfo)
 
 	/* get header */
 	desc = SPI_tuptable->tupdesc;
-
-	/* check if type is ok */
-	if (SPI_gettypeid(desc, 1) != INT4OID)
-		plproxy_error(func, "Hash result must be int4");
+	htype = SPI_gettypeid(desc, 1);
 
 	/* tag connections */
 	for (i = 0; i < SPI_processed; i++)
 	{
 		bool		isnull;
-		int			hashval;
+		uint32		hashval = 0;
 		HeapTuple	row = SPI_tuptable->vals[i];
 		Datum		val = SPI_getbinval(row, desc, 1, &isnull);
 
 		if (isnull)
 			plproxy_error(func, "Hash function returned NULL");
-		hashval = DatumGetInt32(val) & cluster->part_mask;
+
+		if (htype == INT4OID)
+			hashval = DatumGetInt32(val);
+		else if (htype == INT8OID)
+			hashval = DatumGetInt64(val);
+		else if (htype == INT2OID)
+			hashval = DatumGetInt16(val);
+		else
+			plproxy_error(func, "Hash result must be int2, int4 or int8");
+
+		hashval &= cluster->part_mask;
 		cluster->part_map[hashval]->run_on = 1;
 	}
 
-- 
2.39.5