implement and test ANY ELEMENT OF.. WITH INDEX

Author: feodor

src/backend/utils/adt/arrayfuncs.c

@@ -5868,73 +5868,101 @@ array_fill_internal(ArrayType *dims, ArrayType *lbs,
 	return result;
 }
 
+typedef struct ArrayUnnestState
+{
+	AnyArrayType	*array;
+	array_iter		iter;
+	int				nextelem;
+	int				numelems;
+	int16			elmlen;
+	bool			elmbyval;
+	char			elmalign;
+	TupleDesc		tupdesc;
+} ArrayUnnestState;
+
+static void
+init_array_unnest(FuncCallContext *funcctx,  Datum array, FunctionCallInfo fcinfo)
+{
+	MemoryContext		 oldcontext;
+	ArrayUnnestState	*state;
+
+	oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
+
+	state = palloc0(sizeof(*state));
+	/*
+	 * Get the array value and detoast if needed.  We can't do this
+	 * earlier because if we have to detoast, we want the detoasted copy
+	 * to be in multi_call_memory_ctx, so it will go away when we're done
+	 * and not before.  (If no detoast happens, we assume the originally
+	 * passed array will stick around till then.)
+	 */
+	state->array = DatumGetAnyArray(array);
+	array_iter_setup(&state->iter, state->array);
+	state->numelems = ArrayGetNItems(AARR_NDIM(state->array), AARR_DIMS(state->array));
+
+	if (VARATT_IS_EXPANDED_HEADER(state->array))
+	{
+		/* we can just grab the type data from expanded array */
+		state->elmlen = state->array->xpn.typlen;
+		state->elmbyval = state->array->xpn.typbyval;
+		state->elmalign = state->array->xpn.typalign;
+	}
+	else
+		get_typlenbyvalalign(AARR_ELEMTYPE(state->array),
+							 &state->elmlen,
+							 &state->elmbyval,
+							 &state->elmalign);
+
+	if (fcinfo && get_call_result_type(fcinfo, NULL, &state->tupdesc) != TYPEFUNC_COMPOSITE)
+		elog(ERROR, "return type must be a row type");
+
+	funcctx->user_fctx = state;
+	MemoryContextSwitchTo(oldcontext);
+}
 
 /*
  * UNNEST
  */
 Datum
 array_unnest(PG_FUNCTION_ARGS)
 {
-	typedef struct
-	{
-		array_iter	iter;
-		int			nextelem;
-		int			numelems;
-		int16		elmlen;
-		bool		elmbyval;
-		char		elmalign;
-	} array_unnest_fctx;
-
-	FuncCallContext *funcctx;
-	array_unnest_fctx *fctx;
-	MemoryContext oldcontext;
+	FuncCallContext 	*funcctx;
+	ArrayUnnestState 	*fctx;
 
 	/* stuff done only on the first call of the function */
 	if (SRF_IS_FIRSTCALL())
-	{
-		AnyArrayType *arr;
-
-		/* create a function context for cross-call persistence */
-		funcctx = SRF_FIRSTCALL_INIT();
+		init_array_unnest(SRF_FIRSTCALL_INIT(), PG_GETARG_DATUM(0), NULL);
 
-		/*
-		 * switch to memory context appropriate for multiple function calls
-		 */
-		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
+	/* stuff done on every call of the function */
+	funcctx = SRF_PERCALL_SETUP();
+	fctx = funcctx->user_fctx;
 
-		/*
-		 * Get the array value and detoast if needed.  We can't do this
-		 * earlier because if we have to detoast, we want the detoasted copy
-		 * to be in multi_call_memory_ctx, so it will go away when we're done
-		 * and not before.  (If no detoast happens, we assume the originally
-		 * passed array will stick around till then.)
-		 */
-		arr = PG_GETARG_ANY_ARRAY(0);
+	if (fctx->nextelem < fctx->numelems)
+	{
+		int			offset = fctx->nextelem++;
+		Datum		elem;
 
-		/* allocate memory for user context */
-		fctx = (array_unnest_fctx *) palloc(sizeof(array_unnest_fctx));
+		elem = array_iter_next(&fctx->iter, &fcinfo->isnull, offset,
+							   fctx->elmlen, fctx->elmbyval, fctx->elmalign);
 
-		/* initialize state */
-		array_iter_setup(&fctx->iter, arr);
-		fctx->nextelem = 0;
-		fctx->numelems = ArrayGetNItems(AARR_NDIM(arr), AARR_DIMS(arr));
+		SRF_RETURN_NEXT(funcctx, elem);
+	}
+	else
+	{
+		/* do when there is no more left */
+		SRF_RETURN_DONE(funcctx);
+	}
+}
 
-		if (VARATT_IS_EXPANDED_HEADER(arr))
-		{
-			/* we can just grab the type data from expanded array */
-			fctx->elmlen = arr->xpn.typlen;
-			fctx->elmbyval = arr->xpn.typbyval;
-			fctx->elmalign = arr->xpn.typalign;
-		}
-		else
-			get_typlenbyvalalign(AARR_ELEMTYPE(arr),
-								 &fctx->elmlen,
-								 &fctx->elmbyval,
-								 &fctx->elmalign);
+Datum
+array_unnest_element_index(PG_FUNCTION_ARGS)
+{
+	FuncCallContext 	*funcctx;
+	ArrayUnnestState 	*fctx;
 
-		funcctx->user_fctx = fctx;
-		MemoryContextSwitchTo(oldcontext);
-	}
+	/* stuff done only on the first call of the function */
+	if (SRF_IS_FIRSTCALL())
+		init_array_unnest(SRF_FIRSTCALL_INIT(), PG_GETARG_DATUM(0), fcinfo);
 
 	/* stuff done on every call of the function */
 	funcctx = SRF_PERCALL_SETUP();
@@ -5944,11 +5972,20 @@ array_unnest(PG_FUNCTION_ARGS)
 	{
 		int			offset = fctx->nextelem++;
 		Datum		elem;
+		Datum		vals[2];
+		bool		nulls[2] = {false, false};
+		bool		isnull;
 
-		elem = array_iter_next(&fctx->iter, &fcinfo->isnull, offset,
+		elem = array_iter_next(&fctx->iter, &isnull, offset,
 							   fctx->elmlen, fctx->elmbyval, fctx->elmalign);
 
-		SRF_RETURN_NEXT(funcctx, elem);
+		if (isnull)
+			nulls[0] = true;
+		else
+			vals[0] = elem;
+		vals[1] = Int32GetDatum(fctx->nextelem);
+
+		SRF_RETURN_NEXT(funcctx, HeapTupleGetDatum(heap_form_tuple(fctx->tupdesc, vals, nulls)));
 	}
 	else
 	{

src/backend/utils/adt/jsonfuncs.c

@@ -4003,6 +4003,8 @@ finiUnnest(UnnestState *state)
 {
 	if (state)
 	{
+		if (state->tupdesc)
+			FreeTupleDesc(state->tupdesc);
 		pfree(state->jb);
 		pfree(state);
 	}

src/include/catalog/pg_proc.h

@@ -927,6 +927,8 @@ DESCR("expand array to set of rows");
 /* just for compatibility with jsonb_unnest* for any/each element clause */
 DATA(insert OID = 7644 (  unnest_element		   PGNSP PGUID 12 1 100 0 0 f f f f t t i 2 0 2283 "2277 16" _null_ _null_ _null_ _null_ _null_ array_unnest_element _null_ _null_ _null_ ));
 DESCR("expand array to set of rows");
+DATA(insert OID = 7653 (  unnest_element_index		   PGNSP PGUID 12 1 100 0 0 f f f f t t i 2 0 2249 "2277 16" "{2277,16,2283,23}" "{i,i,o,o}" "{array_in,recursive,element,index}" _null_ _null_ array_unnest_element_index _null_ _null_ _null_ ));
+DESCR("expand array to set of rows with index");
 DATA(insert OID = 3167 (  array_remove	   PGNSP PGUID 12 1 0 0 0 f f f f f f i 2 0 2277 "2277 2283" _null_ _null_ _null_ _null_ _null_ array_remove _null_ _null_ _null_ ));
 DESCR("remove any occurrences of an element from an array");
 DATA(insert OID = 3168 (  array_replace    PGNSP PGUID 12 1 0 0 0 f f f f f f i 3 0 2277 "2277 2283 2283" _null_ _null_ _null_ _null_ _null_ array_replace _null_ _null_ _null_ ));

src/include/utils/array.h

@@ -357,6 +357,7 @@ extern Datum array_fill(PG_FUNCTION_ARGS);
 extern Datum array_fill_with_lower_bounds(PG_FUNCTION_ARGS);
 extern Datum array_unnest(PG_FUNCTION_ARGS);
 extern Datum array_unnest_element(PG_FUNCTION_ARGS);
+extern Datum array_unnest_element_index(PG_FUNCTION_ARGS);
 extern Datum array_remove(PG_FUNCTION_ARGS);
 extern Datum array_replace(PG_FUNCTION_ARGS);
 extern Datum width_bucket_array(PG_FUNCTION_ARGS);

src/test/regress/expected/arrays.out

@@ -2084,3 +2084,9 @@ SELECT a, (1 = ALL(a)), EACH ELEMENT OF a AS v SATISFIES (v = 1) FROM anytest;
  {NULL,2} | f        | f
 (8 rows)
 
+SELECT ANY ELEMENT OF '{1,2,3}'::int4[] AS e with index as i SATISFIES (e > 1 and i = 3);
+ bool_or_not_null 
+------------------
+ t
+(1 row)
+

src/test/regress/expected/jsonb.out

@@ -2174,6 +2174,12 @@ SELECT count(*) FROM testjsonb WHERE ANY KEY OF j AS k SATISFIES (
      4
 (1 row)
 
+SELECT ANY ELEMENT OF '[1,2,3]'::jsonb AS e with index as i SATISFIES (e > '1'::jsonb and i = 1);
+ bool_or_not_null 
+------------------
+ t
+(1 row)
+
 CREATE INDEX jidx ON testjsonb USING gin (j);
 SET enable_seqscan = off;
 SELECT count(*) FROM testjsonb WHERE j @> '{"wait":null}';

src/test/regress/expected/jsonb_1.out

@@ -2174,6 +2174,12 @@ SELECT count(*) FROM testjsonb WHERE ANY KEY OF j AS k SATISFIES (
      4
 (1 row)
 
+SELECT ANY ELEMENT OF '[1,2,3]'::jsonb AS e with index as i SATISFIES (e > '1'::jsonb and i = 1);
+ bool_or_not_null 
+------------------
+ t
+(1 row)
+
 CREATE INDEX jidx ON testjsonb USING gin (j);
 SET enable_seqscan = off;
 SELECT count(*) FROM testjsonb WHERE j @> '{"wait":null}';

src/test/regress/sql/arrays.sql

@@ -619,3 +619,5 @@ insert into anytest values
 
 SELECT a, (1 = ANY(a)), ANY  ELEMENT OF a AS v SATISFIES (v = 1) FROM anytest;
 SELECT a, (1 = ALL(a)), EACH ELEMENT OF a AS v SATISFIES (v = 1) FROM anytest;
+
+SELECT ANY ELEMENT OF '{1,2,3}'::int4[] AS e with index as i SATISFIES (e > 1 and i = 3);

src/test/regress/sql/jsonb.sql

@@ -530,6 +530,8 @@ SELECT count(*) FROM testjsonb WHERE ANY ELEMENT OF j->'array' AS e SATISFIES (
 SELECT count(*) FROM testjsonb WHERE ANY KEY OF j AS k SATISFIES (
 	k = 'array' AND ANY ELEMENT OF j->k AS e SATISFIES ( e = '"baz"'::jsonb )
 );
+SELECT ANY ELEMENT OF '[1,2,3]'::jsonb AS e with index as i SATISFIES (e > '1'::jsonb and i = 1);
+
 
 CREATE INDEX jidx ON testjsonb USING gin (j);
 SET enable_seqscan = off;