Add FOREACH IN ARRAY looping to plpgsql.

(I'm not entirely sure that we've finished bikeshedding the syntax details,
but the functionality seems OK.)

Pavel Stehule, reviewed by Stephen Frost and Tom Lane

Author: tglsfdc

doc/src/sgml/plpgsql.sgml

@@ -1263,7 +1263,7 @@ EXECUTE 'UPDATE tbl SET '
 <programlisting>
 EXECUTE format('UPDATE tbl SET %I = %L WHERE key = %L', colname, newvalue, keyvalue);
 </programlisting>
-     The <function>format</function> function can be used in conjunction with 
+     The <function>format</function> function can be used in conjunction with
      the <literal>USING</literal> clause:
 <programlisting>
 EXECUTE format('UPDATE tbl SET %I = $1 WHERE key = $2', colname)
@@ -1356,19 +1356,15 @@ GET DIAGNOSTICS integer_var = ROW_COUNT;
             true if it successfully repositions the cursor, false otherwise.
            </para>
           </listitem>
-
           <listitem>
            <para>
-            A <command>FOR</> statement sets <literal>FOUND</literal> true
-            if it iterates one or more times, else false.  This applies to
-            all four variants of the <command>FOR</> statement (integer
-            <command>FOR</> loops, record-set <command>FOR</> loops,
-            dynamic record-set <command>FOR</> loops, and cursor
-            <command>FOR</> loops).
+            A <command>FOR</> or <command>FOREACH</> statement sets
+            <literal>FOUND</literal> true
+            if it iterates one or more times, else false.
             <literal>FOUND</literal> is set this way when the
-            <command>FOR</> loop exits; inside the execution of the loop,
+            loop exits; inside the execution of the loop,
             <literal>FOUND</literal> is not modified by the
-            <command>FOR</> statement, although it might be changed by the
+            loop statement, although it might be changed by the
             execution of other statements within the loop body.
            </para>
           </listitem>
@@ -1910,9 +1906,9 @@ END CASE;
 
     <para>
      With the <literal>LOOP</>, <literal>EXIT</>,
-     <literal>CONTINUE</>, <literal>WHILE</>, and <literal>FOR</>
-     statements, you can arrange for your <application>PL/pgSQL</>
-     function to repeat a series of commands.
+     <literal>CONTINUE</>, <literal>WHILE</>, <literal>FOR</>,
+     and <literal>FOREACH</> statements, you can arrange for your
+     <application>PL/pgSQL</> function to repeat a series of commands.
     </para>
 
     <sect3>
@@ -2238,6 +2234,90 @@ END LOOP <optional> <replaceable>label</replaceable> </optional>;
     </para>
    </sect2>
 
+   <sect2 id="plpgsql-foreach-array">
+    <title>Looping Through Arrays</title>
+
+    <para>
+     The <literal>FOREACH</> loop is much like a <literal>FOR</> loop,
+     but instead of iterating through the rows returned by a SQL query,
+     it iterates through the elements of an array value.
+     (In general, <literal>FOREACH</> is meant for looping through
+     components of a composite-valued expression; variants for looping
+     through composites besides arrays may be added in future.)
+     The <literal>FOREACH</> statement to loop over an array is:
+
+<synopsis>
+<optional> &lt;&lt;<replaceable>label</replaceable>&gt;&gt; </optional>
+FOREACH <replaceable>target</replaceable> <optional> SLICE <replaceable>number</replaceable> </optional> IN ARRAY <replaceable>expression</replaceable> LOOP
+    <replaceable>statements</replaceable>
+END LOOP <optional> <replaceable>label</replaceable> </optional>;
+</synopsis>
+    </para>
+
+    <para>
+     Without <literal>SLICE</>, or if <literal>SLICE 0</> is specified,
+     the loop iterates through individual elements of the array produced
+     by evaluating the <replaceable>expression</replaceable>.
+     The <replaceable>target</replaceable> variable is assigned each
+     element value in sequence, and the loop body is executed for each element.
+     Here is an example of looping through the elements of an integer
+     array:
+
+<programlisting>
+CREATE FUNCTION sum(int[]) RETURNS int8 AS $$
+DECLARE
+  s int8 := 0;
+  x int;
+BEGIN
+  FOREACH x IN ARRAY $1
+  LOOP
+    s := s + x;
+  END LOOP;
+  RETURN s;
+END;
+$$ LANGUAGE plpgsql;
+</programlisting>
+
+     The elements are visited in storage order, regardless of the number of
+     array dimensions.  Although the <replaceable>target</replaceable> is
+     usually just a single variable, it can be a list of variables when
+     looping through an array of composite values (records).  In that case,
+     for each array element, the variables are assigned from successive
+     columns of the composite value.
+    </para>
+
+    <para>
+     With a positive <literal>SLICE</> value, <literal>FOREACH</>
+     iterates through slices of the array rather than single elements.
+     The <literal>SLICE</> value must be an integer constant not larger
+     than the number of dimensions of the array.  The
+     <replaceable>target</replaceable> variable must be an array,
+     and it receives successive slices of the array value, where each slice
+     is of the number of dimensions specified by <literal>SLICE</>.
+     Here is an example of iterating through one-dimensional slices:
+
+<programlisting>
+CREATE FUNCTION scan_rows(int[]) RETURNS void AS $$
+DECLARE
+  x int[];
+BEGIN
+  FOREACH x SLICE 1 IN ARRAY $1
+  LOOP
+    RAISE NOTICE 'row = %', x;
+  END LOOP;
+END;
+$$ LANGUAGE plpgsql;
+
+SELECT scan_rows(ARRAY[[1,2,3],[4,5,6],[7,8,9],[10,11,12]]);
+
+NOTICE:  row = {1,2,3}
+NOTICE:  row = {4,5,6}
+NOTICE:  row = {7,8,9}
+NOTICE:  row = {10,11,12}
+</programlisting>
+    </para>
+   </sect2>
+
    <sect2 id="plpgsql-error-trapping">
     <title>Trapping Errors</title>
 

src/backend/utils/adt/arrayfuncs.c

@@ -50,6 +50,30 @@ typedef enum
 	ARRAY_LEVEL_DELIMITED
 } ArrayParseState;
 
+/* Working state for array_iterate() */
+typedef struct ArrayIteratorData
+{
+	/* basic info about the array, set up during array_create_iterator() */
+	ArrayType  *arr;			/* array we're iterating through */
+	bits8	   *nullbitmap;		/* its null bitmap, if any */
+	int			nitems;			/* total number of elements in array */
+	int16		typlen;			/* element type's length */
+	bool		typbyval;		/* element type's byval property */
+	char		typalign;		/* element type's align property */
+
+	/* information about the requested slice size */
+	int			slice_ndim;		/* slice dimension, or 0 if not slicing */
+	int			slice_len;		/* number of elements per slice */
+	int		   *slice_dims;		/* slice dims array */
+	int		   *slice_lbound;	/* slice lbound array */
+	Datum	   *slice_values;	/* workspace of length slice_len */
+	bool	   *slice_nulls;	/* workspace of length slice_len */
+
+	/* current position information, updated on each iteration */
+	char	   *data_ptr;		/* our current position in the array */
+	int			current_item;	/* the item # we're at in the array */
+} ArrayIteratorData;
+
 static bool array_isspace(char ch);
 static int	ArrayCount(const char *str, int *dim, char typdelim);
 static void ReadArrayStr(char *arrayStr, const char *origStr,
@@ -3833,6 +3857,188 @@ arraycontained(PG_FUNCTION_ARGS)
 }
 
 
+/*-----------------------------------------------------------------------------
+ * Array iteration functions
+ *		These functions are used to iterate efficiently through arrays
+ *-----------------------------------------------------------------------------
+ */
+
+/*
+ * array_create_iterator --- set up to iterate through an array
+ *
+ * If slice_ndim is zero, we will iterate element-by-element; the returned
+ * datums are of the array's element type.
+ *
+ * If slice_ndim is 1..ARR_NDIM(arr), we will iterate by slices: the
+ * returned datums are of the same array type as 'arr', but of size
+ * equal to the rightmost N dimensions of 'arr'.
+ *
+ * The passed-in array must remain valid for the lifetime of the iterator.
+ */
+ArrayIterator
+array_create_iterator(ArrayType *arr, int slice_ndim)
+{
+	ArrayIterator iterator = palloc0(sizeof(ArrayIteratorData));
+
+	/*
+	 * Sanity-check inputs --- caller should have got this right already
+	 */
+	Assert(PointerIsValid(arr));
+	if (slice_ndim < 0 || slice_ndim > ARR_NDIM(arr))
+		elog(ERROR, "invalid arguments to array_create_iterator");
+
+	/*
+	 * Remember basic info about the array and its element type
+	 */
+	iterator->arr = arr;
+	iterator->nullbitmap = ARR_NULLBITMAP(arr);
+	iterator->nitems = ArrayGetNItems(ARR_NDIM(arr), ARR_DIMS(arr));
+	get_typlenbyvalalign(ARR_ELEMTYPE(arr),
+						 &iterator->typlen,
+						 &iterator->typbyval,
+						 &iterator->typalign);
+
+	/*
+	 * Remember the slicing parameters.
+	 */
+	iterator->slice_ndim = slice_ndim;
+
+	if (slice_ndim > 0)
+	{
+		/*
+		 * Get pointers into the array's dims and lbound arrays to represent
+		 * the dims/lbound arrays of a slice.  These are the same as the
+		 * rightmost N dimensions of the array.
+		 */
+		iterator->slice_dims = ARR_DIMS(arr) + ARR_NDIM(arr) - slice_ndim;
+		iterator->slice_lbound = ARR_LBOUND(arr) + ARR_NDIM(arr) - slice_ndim;
+
+		/*
+		 * Compute number of elements in a slice.
+		 */
+		iterator->slice_len = ArrayGetNItems(slice_ndim,
+											 iterator->slice_dims);
+
+		/*
+		 * Create workspace for building sub-arrays.
+		 */
+		iterator->slice_values = (Datum *)
+			palloc(iterator->slice_len * sizeof(Datum));
+		iterator->slice_nulls = (bool *)
+			palloc(iterator->slice_len * sizeof(bool));
+	}
+
+	/*
+	 * Initialize our data pointer and linear element number.  These will
+	 * advance through the array during array_iterate().
+	 */
+	iterator->data_ptr = ARR_DATA_PTR(arr);
+	iterator->current_item = 0;
+
+	return iterator;
+}
+
+/*
+ * Iterate through the array referenced by 'iterator'.
+ *
+ * As long as there is another element (or slice), return it into
+ * *value / *isnull, and return true.  Return false when no more data.
+ */
+bool
+array_iterate(ArrayIterator iterator, Datum *value, bool *isnull)
+{
+	/* Done if we have reached the end of the array */
+	if (iterator->current_item >= iterator->nitems)
+		return false;
+
+	if (iterator->slice_ndim == 0)
+	{
+		/*
+		 * Scalar case: return one element.
+		 */
+		if (array_get_isnull(iterator->nullbitmap, iterator->current_item++))
+		{
+			*isnull = true;
+			*value = (Datum) 0;
+		}
+		else
+		{
+			/* non-NULL, so fetch the individual Datum to return */
+			char	   *p = iterator->data_ptr;
+
+			*isnull = false;
+			*value = fetch_att(p, iterator->typbyval, iterator->typlen);
+
+			/* Move our data pointer forward to the next element */
+			p = att_addlength_pointer(p, iterator->typlen, p);
+			p = (char *) att_align_nominal(p, iterator->typalign);
+			iterator->data_ptr = p;
+		}
+	}
+	else
+	{
+		/*
+		 * Slice case: build and return an array of the requested size.
+		 */
+		ArrayType  *result;
+		Datum	   *values = iterator->slice_values;
+		bool	   *nulls = iterator->slice_nulls;
+		char	   *p = iterator->data_ptr;
+		int			i;
+
+		for (i = 0; i < iterator->slice_len; i++)
+		{
+			if (array_get_isnull(iterator->nullbitmap,
+								 iterator->current_item++))
+			{
+				nulls[i] = true;
+				values[i] = (Datum) 0;
+			}
+			else
+			{
+				nulls[i] = false;
+				values[i] = fetch_att(p, iterator->typbyval, iterator->typlen);
+
+				/* Move our data pointer forward to the next element */
+				p = att_addlength_pointer(p, iterator->typlen, p);
+				p = (char *) att_align_nominal(p, iterator->typalign);
+			}
+		}
+
+		iterator->data_ptr = p;
+
+		result = construct_md_array(values,
+									nulls,
+									iterator->slice_ndim,
+									iterator->slice_dims,
+									iterator->slice_lbound,
+									ARR_ELEMTYPE(iterator->arr),
+									iterator->typlen,
+									iterator->typbyval,
+									iterator->typalign);
+
+		*isnull = false;
+		*value = PointerGetDatum(result);
+	}
+
+	return true;
+}
+
+/*
+ * Release an ArrayIterator data structure
+ */
+void
+array_free_iterator(ArrayIterator iterator)
+{
+	if (iterator->slice_ndim > 0)
+	{
+		pfree(iterator->slice_values);
+		pfree(iterator->slice_nulls);
+	}
+	pfree(iterator);
+}
+
+
 /***************************************************************************/
 /******************|		  Support  Routines			  |*****************/
 /***************************************************************************/

src/include/utils/array.h

@@ -114,6 +114,9 @@ typedef struct ArrayMapState
 	ArrayMetaState ret_extra;
 } ArrayMapState;
 
+/* ArrayIteratorData is private in arrayfuncs.c */
+typedef struct ArrayIteratorData *ArrayIterator;
+
 /*
  * fmgr macros for array objects
  */
@@ -254,6 +257,10 @@ extern Datum makeArrayResult(ArrayBuildState *astate,
 extern Datum makeMdArrayResult(ArrayBuildState *astate, int ndims,
 				  int *dims, int *lbs, MemoryContext rcontext, bool release);
 
+extern ArrayIterator array_create_iterator(ArrayType *arr, int slice_ndim);
+extern bool array_iterate(ArrayIterator iterator, Datum *value, bool *isnull);
+extern void array_free_iterator(ArrayIterator iterator);
+
 /*
  * prototypes for functions defined in arrayutils.c
  */