PL/Python array support

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

doc/src/sgml/plpython.sgml

-<!-- $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>
@@ -134,6 +134,43 @@ $$ LANGUAGE plpythonu;
    function is strict or not.
   </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.

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"

src/pl/plpython/plpython.c

 /**********************************************************************
  * 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
@@ -1832,6 +1877,45 @@ PLyString_FromDatum(PLyDatumToOb *arg, Datum d)
 	return r;
 }
 
+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)
 {
@@ -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)
 {

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();