PL/Python array support

Support arrays as parameters and return values of PL/Python functions.

4 files changed, 343 insertions, 4 deletions

diff --git a/doc/src/sgml/plpython.sgml b/doc/src/sgml/plpython.sgml
index 4944b8ad6c..5820872039 100644
--- a/doc/src/sgml/plpython.sgml
+++ b/doc/src/sgml/plpython.sgml
@@ -1,4 +1,4 @@
-<!-- $PostgreSQL: pgsql/doc/src/sgml/plpython.sgml,v 1.40 2009/03/30 16:15:43 alvherre Exp $ -->
+<!-- $PostgreSQL: pgsql/doc/src/sgml/plpython.sgml,v 1.41 2009/12/10 20:43:40 petere Exp $ -->
 
 <chapter id="plpython">
  <title>PL/Python - Python Procedural Language</title>
@@ -135,6 +135,43 @@ $$ LANGUAGE plpythonu;
   </para>
 
   <para>
+   SQL array values are passed into PL/Python as a Python list.  To
+   return an SQL array value out of a PL/Python function, return a
+   Python sequence, for example a list or tuple:
+
+<programlisting>
+CREATE FUNCTION return_arr()
+  RETURNS int[]
+AS $$
+return (1, 2, 3, 4, 5)
+$$ LANGUAGE plpythonu;
+
+SELECT return_arr();
+ return_arr  
+-------------
+ {1,2,3,4,5}
+(1 row)
+</programlisting>
+
+   Note that in Python, strings are sequences, which can have
+   undesirable effects that might be familiar to Python programmers:
+
+<programlisting>
+CREATE FUNCTION return_str_arr()
+  RETURNS varchar[]
+AS $$
+return "hello"
+$$ LANGUAGE plpythonu;
+
+SELECT return_str_arr();
+ return_str_arr
+----------------
+ {h,e,l,l,o}
+(1 row)
+</programlisting>
+  </para>
+
+  <para>
    Composite-type arguments are passed to the function as Python mappings. The
    element names of the mapping are the attribute names of the composite type.
    If an attribute in the passed row has the null value, it has the value
diff --git a/src/pl/plpython/expected/plpython_types.out b/src/pl/plpython/expected/plpython_types.out
index 2dd498c058..9cda31b13d 100644
--- a/src/pl/plpython/expected/plpython_types.out
+++ b/src/pl/plpython/expected/plpython_types.out
@@ -477,3 +477,113 @@ CONTEXT:  PL/Python function "test_type_conversion_bytea10"
 ERROR:  value for domain bytea10 violates check constraint "bytea10_check"
 CONTEXT:  while creating return value
 PL/Python function "test_type_conversion_bytea10"
+--
+-- Arrays
+--
+CREATE FUNCTION test_type_conversion_array_int4(x int4[]) RETURNS int4[] AS $$
+plpy.info(x, type(x))
+return x
+$$ LANGUAGE plpythonu;
+SELECT * FROM test_type_conversion_array_int4(ARRAY[0, 100]);
+INFO:  ([0, 100], <type 'list'>)
+CONTEXT:  PL/Python function "test_type_conversion_array_int4"
+ test_type_conversion_array_int4 
+---------------------------------
+ {0,100}
+(1 row)
+
+SELECT * FROM test_type_conversion_array_int4(ARRAY[0,-100,55]);
+INFO:  ([0, -100, 55], <type 'list'>)
+CONTEXT:  PL/Python function "test_type_conversion_array_int4"
+ test_type_conversion_array_int4 
+---------------------------------
+ {0,-100,55}
+(1 row)
+
+SELECT * FROM test_type_conversion_array_int4(ARRAY[NULL,1]);
+INFO:  ([None, 1], <type 'list'>)
+CONTEXT:  PL/Python function "test_type_conversion_array_int4"
+ test_type_conversion_array_int4 
+---------------------------------
+ {NULL,1}
+(1 row)
+
+SELECT * FROM test_type_conversion_array_int4(ARRAY[]::integer[]);
+INFO:  ([], <type 'list'>)
+CONTEXT:  PL/Python function "test_type_conversion_array_int4"
+ test_type_conversion_array_int4 
+---------------------------------
+ {}
+(1 row)
+
+SELECT * FROM test_type_conversion_array_int4(NULL);
+INFO:  (None, <type 'NoneType'>)
+CONTEXT:  PL/Python function "test_type_conversion_array_int4"
+ test_type_conversion_array_int4 
+---------------------------------
+ 
+(1 row)
+
+SELECT * FROM test_type_conversion_array_int4(ARRAY[[1,2,3],[4,5,6]]);
+ERROR:  cannot convert multidimensional array to Python list
+DETAIL:  PL/Python only supports one-dimensional arrays.
+CONTEXT:  PL/Python function "test_type_conversion_array_int4"
+CREATE FUNCTION test_type_conversion_array_bytea(x bytea[]) RETURNS bytea[] AS $$
+plpy.info(x, type(x))
+return x
+$$ LANGUAGE plpythonu;
+SELECT * FROM test_type_conversion_array_bytea(ARRAY[E'\\xdeadbeef'::bytea, NULL]);
+INFO:  (['\xde\xad\xbe\xef', None], <type 'list'>)
+CONTEXT:  PL/Python function "test_type_conversion_array_bytea"
+ test_type_conversion_array_bytea 
+----------------------------------
+ {"\\xdeadbeef",NULL}
+(1 row)
+
+CREATE FUNCTION test_type_conversion_array_mixed1() RETURNS text[] AS $$
+return [123, 'abc']
+$$ LANGUAGE plpythonu;
+SELECT * FROM test_type_conversion_array_mixed1();
+ test_type_conversion_array_mixed1 
+-----------------------------------
+ {123,abc}
+(1 row)
+
+CREATE FUNCTION test_type_conversion_array_mixed2() RETURNS int[] AS $$
+return [123, 'abc']
+$$ LANGUAGE plpythonu;
+SELECT * FROM test_type_conversion_array_mixed2();
+ERROR:  invalid input syntax for integer: "abc"
+CONTEXT:  while creating return value
+PL/Python function "test_type_conversion_array_mixed2"
+CREATE FUNCTION test_type_conversion_array_record() RETURNS type_record[] AS $$
+return [None]
+$$ LANGUAGE plpythonu;
+SELECT * FROM test_type_conversion_array_record();
+ERROR:  PL/Python functions cannot return type type_record[]
+DETAIL:  PL/Python does not support conversion to arrays of row types.
+CREATE FUNCTION test_type_conversion_array_string() RETURNS text[] AS $$
+return 'abc'
+$$ LANGUAGE plpythonu;
+SELECT * FROM test_type_conversion_array_string();
+ test_type_conversion_array_string 
+-----------------------------------
+ {a,b,c}
+(1 row)
+
+CREATE FUNCTION test_type_conversion_array_tuple() RETURNS text[] AS $$
+return ('abc', 'def')
+$$ LANGUAGE plpythonu;
+SELECT * FROM test_type_conversion_array_tuple();
+ test_type_conversion_array_tuple 
+----------------------------------
+ {abc,def}
+(1 row)
+
+CREATE FUNCTION test_type_conversion_array_error() RETURNS int[] AS $$
+return 5
+$$ LANGUAGE plpythonu;
+SELECT * FROM test_type_conversion_array_error();
+ERROR:  PL/Python: return value of function with array return type is not a Python sequence
+CONTEXT:  while creating return value
+PL/Python function "test_type_conversion_array_error"
diff --git a/src/pl/plpython/plpython.c b/src/pl/plpython/plpython.c
index 820641708c..8e1c8689c9 100644
--- a/src/pl/plpython/plpython.c
+++ b/src/pl/plpython/plpython.c
@@ -1,7 +1,7 @@
 /**********************************************************************
  * plpython.c - python as a procedural language for PostgreSQL
  *
- *	$PostgreSQL: pgsql/src/pl/plpython/plpython.c,v 1.132 2009/11/03 11:05:02 petere Exp $
+ *	$PostgreSQL: pgsql/src/pl/plpython/plpython.c,v 1.133 2009/12/10 20:43:40 petere Exp $
  *
  *********************************************************************
  */
@@ -89,6 +89,9 @@ typedef struct PLyDatumToOb
 	Oid			typoid;			/* The OID of the type */
 	Oid			typioparam;
 	bool		typbyval;
+	int16		typlen;
+	char		typalign;
+	struct PLyDatumToOb *elm;
 } PLyDatumToOb;
 
 typedef struct PLyTupleToOb
@@ -120,6 +123,9 @@ typedef struct PLyObToDatum
 	Oid			typoid;			/* The OID of the type */
 	Oid			typioparam;
 	bool		typbyval;
+	int16		typlen;
+	char		typalign;
+	struct PLyObToDatum *elm;
 } PLyObToDatum;
 
 typedef struct PLyObToTuple
@@ -284,6 +290,7 @@ static PyObject *PLyInt_FromInt32(PLyDatumToOb *arg, Datum d);
 static PyObject *PLyLong_FromInt64(PLyDatumToOb *arg, Datum d);
 static PyObject *PLyString_FromBytea(PLyDatumToOb *arg, Datum d);
 static PyObject *PLyString_FromDatum(PLyDatumToOb *arg, Datum d);
+static PyObject *PLyList_FromArray(PLyDatumToOb *arg, Datum d);
 
 static PyObject *PLyDict_FromTuple(PLyTypeInfo *, HeapTuple, TupleDesc);
 
@@ -293,6 +300,8 @@ static Datum PLyObject_ToBytea(PLyTypeInfo *, PLyObToDatum *,
 							   PyObject *);
 static Datum PLyObject_ToDatum(PLyTypeInfo *, PLyObToDatum *,
 							   PyObject *);
+static Datum PLySequence_ToArray(PLyTypeInfo *, PLyObToDatum *,
+								 PyObject *);
 
 static HeapTuple PLyMapping_ToTuple(PLyTypeInfo *, PyObject *);
 static HeapTuple PLySequence_ToTuple(PLyTypeInfo *, PyObject *);
@@ -1653,18 +1662,21 @@ static void
 PLy_output_datum_func2(PLyObToDatum *arg, HeapTuple typeTup)
 {
 	Form_pg_type typeStruct = (Form_pg_type) GETSTRUCT(typeTup);
+	Oid element_type;
 
 	perm_fmgr_info(typeStruct->typinput, &arg->typfunc);
 	arg->typoid = HeapTupleGetOid(typeTup);
 	arg->typioparam = getTypeIOParam(typeTup);
 	arg->typbyval = typeStruct->typbyval;
 
+	element_type = get_element_type(arg->typoid);
+
 	/*
 	 * Select a conversion function to convert Python objects to
 	 * PostgreSQL datums.  Most data types can go through the generic
 	 * function.
 	 */
-	switch (getBaseType(arg->typoid))
+	switch (getBaseType(element_type ? element_type : arg->typoid))
 	{
 		case BOOLOID:
 			arg->func = PLyObject_ToBool;
@@ -1676,6 +1688,29 @@ PLy_output_datum_func2(PLyObToDatum *arg, HeapTuple typeTup)
 			arg->func = PLyObject_ToDatum;
 			break;
 	}
+
+	if (element_type)
+	{
+		char dummy_delim;
+		Oid funcid;
+
+		if (type_is_rowtype(element_type))
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("PL/Python functions cannot return type %s",
+							format_type_be(arg->typoid)),
+					 errdetail("PL/Python does not support conversion to arrays of row types.")));
+
+		arg->elm = PLy_malloc0(sizeof(*arg->elm));
+		arg->elm->func = arg->func;
+		arg->func = PLySequence_ToArray;
+
+		arg->elm->typoid = element_type;
+		get_type_io_data(element_type, IOFunc_input,
+						 &arg->elm->typlen, &arg->elm->typbyval, &arg->elm->typalign, &dummy_delim,
+						 &arg->elm->typioparam, &funcid);
+		perm_fmgr_info(funcid, &arg->elm->typfunc);
+	}
 }
 
 static void
@@ -1691,15 +1726,17 @@ static void
 PLy_input_datum_func2(PLyDatumToOb *arg, Oid typeOid, HeapTuple typeTup)
 {
 	Form_pg_type typeStruct = (Form_pg_type) GETSTRUCT(typeTup);
+	Oid element_type = get_element_type(typeOid);
 
 	/* Get the type's conversion information */
 	perm_fmgr_info(typeStruct->typoutput, &arg->typfunc);
 	arg->typoid = HeapTupleGetOid(typeTup);
 	arg->typioparam = getTypeIOParam(typeTup);
 	arg->typbyval = typeStruct->typbyval;
+	arg->typlen = typeStruct->typlen;
 
 	/* Determine which kind of Python object we will convert to */
-	switch (getBaseType(typeOid))
+	switch (getBaseType(element_type ? element_type : typeOid))
 	{
 		case BOOLOID:
 			arg->func = PLyBool_FromBool;
@@ -1729,6 +1766,14 @@ PLy_input_datum_func2(PLyDatumToOb *arg, Oid typeOid, HeapTuple typeTup)
 			arg->func = PLyString_FromDatum;
 			break;
 	}
+
+	if (element_type)
+	{
+		arg->elm = PLy_malloc0(sizeof(*arg->elm));
+		arg->elm->func = arg->func;
+		arg->func = PLyList_FromArray;
+		get_typlenbyvalalign(element_type, &arg->elm->typlen, &arg->elm->typbyval, &arg->elm->typalign);
+	}
 }
 
 static void
@@ -1833,6 +1878,45 @@ PLyString_FromDatum(PLyDatumToOb *arg, Datum d)
 }
 
 static PyObject *
+PLyList_FromArray(PLyDatumToOb *arg, Datum d)
+{
+	ArrayType  *array = DatumGetArrayTypeP(d);
+	PyObject   *list;
+	int			length;
+	int			lbound;
+	int			i;
+
+	if (ARR_NDIM(array) == 0)
+		return PyList_New(0);
+
+	if (ARR_NDIM(array) != 1)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("cannot convert multidimensional array to Python list"),
+				 errdetail("PL/Python only supports one-dimensional arrays.")));
+
+	length = ARR_DIMS(array)[0];
+	lbound = ARR_LBOUND(array)[0];
+	list = PyList_New(length);
+
+	for (i = 0; i < length; i++)
+	{
+		Datum elem;
+		bool isnull;
+		int offset;
+
+		offset = lbound + i;
+		elem = array_ref(array, 1, &offset, arg->typlen, arg->elm->typlen, arg->elm->typbyval, arg->elm->typalign, &isnull);
+		if (isnull)
+			PyList_SET_ITEM(list, i, Py_None);
+		else
+			PyList_SET_ITEM(list, i, arg->elm->func(arg, elem));
+	}
+
+	return list;
+}
+
+static PyObject *
 PLyDict_FromTuple(PLyTypeInfo *info, HeapTuple tuple, TupleDesc desc)
 {
 	PyObject   *volatile dict;
@@ -1994,6 +2078,49 @@ PLyObject_ToDatum(PLyTypeInfo *info,
 	return rv;
 }
 
+static Datum
+PLySequence_ToArray(PLyTypeInfo *info,
+					PLyObToDatum *arg,
+					PyObject *plrv)
+{
+	ArrayType *array;
+	int			i;
+	Datum		*elems;
+	bool		*nulls;
+	int			len;
+	int			lbs;
+
+	Assert(plrv != Py_None);
+
+	if (!PySequence_Check(plrv))
+		PLy_elog(ERROR, "return value of function with array return type is not a Python sequence");
+
+	len = PySequence_Length(plrv);
+	elems = palloc(sizeof(*elems) * len);
+	nulls = palloc(sizeof(*nulls) * len);
+
+	for (i = 0; i < len; i++)
+	{
+		PyObject *obj = PySequence_GetItem(plrv, i);
+
+		if (obj == Py_None)
+			nulls[i] = true;
+		else
+		{
+			nulls[i] = false;
+			/* We don't support arrays of row types yet, so the first
+			 * argument can be NULL. */
+			elems[i] = arg->elm->func(NULL, arg->elm, obj);
+		}
+		Py_XDECREF(obj);
+	}
+
+	lbs = 1;
+	array = construct_md_array(elems, nulls, 1, &len, &lbs,
+							   get_element_type(arg->typoid), arg->elm->typlen, arg->elm->typbyval, arg->elm->typalign);
+	return PointerGetDatum(array);
+}
+
 static HeapTuple
 PLyMapping_ToTuple(PLyTypeInfo *info, PyObject *mapping)
 {
diff --git a/src/pl/plpython/sql/plpython_types.sql b/src/pl/plpython/sql/plpython_types.sql
index becf5cf088..2afbc87058 100644
--- a/src/pl/plpython/sql/plpython_types.sql
+++ b/src/pl/plpython/sql/plpython_types.sql
@@ -204,3 +204,68 @@ SELECT * FROM test_type_conversion_bytea10('hello world', 'hello wold');
 SELECT * FROM test_type_conversion_bytea10('hello word', 'hello world');
 SELECT * FROM test_type_conversion_bytea10(null, 'hello word');
 SELECT * FROM test_type_conversion_bytea10('hello word', null);
+
+
+--
+-- Arrays
+--
+
+CREATE FUNCTION test_type_conversion_array_int4(x int4[]) RETURNS int4[] AS $$
+plpy.info(x, type(x))
+return x
+$$ LANGUAGE plpythonu;
+
+SELECT * FROM test_type_conversion_array_int4(ARRAY[0, 100]);
+SELECT * FROM test_type_conversion_array_int4(ARRAY[0,-100,55]);
+SELECT * FROM test_type_conversion_array_int4(ARRAY[NULL,1]);
+SELECT * FROM test_type_conversion_array_int4(ARRAY[]::integer[]);
+SELECT * FROM test_type_conversion_array_int4(NULL);
+SELECT * FROM test_type_conversion_array_int4(ARRAY[[1,2,3],[4,5,6]]);
+
+
+CREATE FUNCTION test_type_conversion_array_bytea(x bytea[]) RETURNS bytea[] AS $$
+plpy.info(x, type(x))
+return x
+$$ LANGUAGE plpythonu;
+
+SELECT * FROM test_type_conversion_array_bytea(ARRAY[E'\\xdeadbeef'::bytea, NULL]);
+
+
+CREATE FUNCTION test_type_conversion_array_mixed1() RETURNS text[] AS $$
+return [123, 'abc']
+$$ LANGUAGE plpythonu;
+
+SELECT * FROM test_type_conversion_array_mixed1();
+
+
+CREATE FUNCTION test_type_conversion_array_mixed2() RETURNS int[] AS $$
+return [123, 'abc']
+$$ LANGUAGE plpythonu;
+
+SELECT * FROM test_type_conversion_array_mixed2();
+
+
+CREATE FUNCTION test_type_conversion_array_record() RETURNS type_record[] AS $$
+return [None]
+$$ LANGUAGE plpythonu;
+
+SELECT * FROM test_type_conversion_array_record();
+
+
+CREATE FUNCTION test_type_conversion_array_string() RETURNS text[] AS $$
+return 'abc'
+$$ LANGUAGE plpythonu;
+
+SELECT * FROM test_type_conversion_array_string();
+
+CREATE FUNCTION test_type_conversion_array_tuple() RETURNS text[] AS $$
+return ('abc', 'def')
+$$ LANGUAGE plpythonu;
+
+SELECT * FROM test_type_conversion_array_tuple();
+
+CREATE FUNCTION test_type_conversion_array_error() RETURNS int[] AS $$
+return 5
+$$ LANGUAGE plpythonu;
+
+SELECT * FROM test_type_conversion_array_error();