Skip to content

Commit f2bc5b1

Browse files
feodorfeodor
committed
implement and test ANY ELEMENT OF.. WITH INDEX
1 parent 95d6f4d commit f2bc5b1

File tree

9 files changed

+117
-53
lines changed

9 files changed

+117
-53
lines changed

‎src/backend/utils/adt/arrayfuncs.c

Lines changed: 90 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -5868,73 +5868,101 @@ array_fill_internal(ArrayType *dims, ArrayType *lbs,
58685868
return result;
58695869
}
58705870

5871+
typedef struct ArrayUnnestState
5872+
{
5873+
AnyArrayType *array;
5874+
array_iter iter;
5875+
int nextelem;
5876+
int numelems;
5877+
int16 elmlen;
5878+
bool elmbyval;
5879+
char elmalign;
5880+
TupleDesc tupdesc;
5881+
} ArrayUnnestState;
5882+
5883+
static void
5884+
init_array_unnest(FuncCallContext *funcctx, Datum array, FunctionCallInfo fcinfo)
5885+
{
5886+
MemoryContext oldcontext;
5887+
ArrayUnnestState *state;
5888+
5889+
oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
5890+
5891+
state = palloc0(sizeof(*state));
5892+
/*
5893+
* Get the array value and detoast if needed. We can't do this
5894+
* earlier because if we have to detoast, we want the detoasted copy
5895+
* to be in multi_call_memory_ctx, so it will go away when we're done
5896+
* and not before. (If no detoast happens, we assume the originally
5897+
* passed array will stick around till then.)
5898+
*/
5899+
state->array = DatumGetAnyArray(array);
5900+
array_iter_setup(&state->iter, state->array);
5901+
state->numelems = ArrayGetNItems(AARR_NDIM(state->array), AARR_DIMS(state->array));
5902+
5903+
if (VARATT_IS_EXPANDED_HEADER(state->array))
5904+
{
5905+
/* we can just grab the type data from expanded array */
5906+
state->elmlen = state->array->xpn.typlen;
5907+
state->elmbyval = state->array->xpn.typbyval;
5908+
state->elmalign = state->array->xpn.typalign;
5909+
}
5910+
else
5911+
get_typlenbyvalalign(AARR_ELEMTYPE(state->array),
5912+
&state->elmlen,
5913+
&state->elmbyval,
5914+
&state->elmalign);
5915+
5916+
if (fcinfo && get_call_result_type(fcinfo, NULL, &state->tupdesc) != TYPEFUNC_COMPOSITE)
5917+
elog(ERROR, "return type must be a row type");
5918+
5919+
funcctx->user_fctx = state;
5920+
MemoryContextSwitchTo(oldcontext);
5921+
}
58715922

58725923
/*
58735924
* UNNEST
58745925
*/
58755926
Datum
58765927
array_unnest(PG_FUNCTION_ARGS)
58775928
{
5878-
typedef struct
5879-
{
5880-
array_iter iter;
5881-
int nextelem;
5882-
int numelems;
5883-
int16 elmlen;
5884-
bool elmbyval;
5885-
char elmalign;
5886-
} array_unnest_fctx;
5887-
5888-
FuncCallContext *funcctx;
5889-
array_unnest_fctx *fctx;
5890-
MemoryContext oldcontext;
5929+
FuncCallContext *funcctx;
5930+
ArrayUnnestState *fctx;
58915931

58925932
/* stuff done only on the first call of the function */
58935933
if (SRF_IS_FIRSTCALL())
5894-
{
5895-
AnyArrayType *arr;
5896-
5897-
/* create a function context for cross-call persistence */
5898-
funcctx = SRF_FIRSTCALL_INIT();
5934+
init_array_unnest(SRF_FIRSTCALL_INIT(), PG_GETARG_DATUM(0), NULL);
58995935

5900-
/*
5901-
* switch to memory context appropriate for multiple function calls
5902-
*/
5903-
oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
5936+
/* stuff done on every call of the function */
5937+
funcctx = SRF_PERCALL_SETUP();
5938+
fctx = funcctx->user_fctx;
59045939

5905-
/*
5906-
* Get the array value and detoast if needed. We can't do this
5907-
* earlier because if we have to detoast, we want the detoasted copy
5908-
* to be in multi_call_memory_ctx, so it will go away when we're done
5909-
* and not before. (If no detoast happens, we assume the originally
5910-
* passed array will stick around till then.)
5911-
*/
5912-
arr = PG_GETARG_ANY_ARRAY(0);
5940+
if (fctx->nextelem < fctx->numelems)
5941+
{
5942+
int offset = fctx->nextelem++;
5943+
Datum elem;
59135944

5914-
/* allocate memory for user context */
5915-
fctx = (array_unnest_fctx *) palloc(sizeof(array_unnest_fctx));
5945+
elem = array_iter_next(&fctx->iter, &fcinfo->isnull, offset,
5946+
fctx->elmlen, fctx->elmbyval, fctx->elmalign);
59165947

5917-
/* initialize state */
5918-
array_iter_setup(&fctx->iter, arr);
5919-
fctx->nextelem = 0;
5920-
fctx->numelems = ArrayGetNItems(AARR_NDIM(arr), AARR_DIMS(arr));
5948+
SRF_RETURN_NEXT(funcctx, elem);
5949+
}
5950+
else
5951+
{
5952+
/* do when there is no more left */
5953+
SRF_RETURN_DONE(funcctx);
5954+
}
5955+
}
59215956

5922-
if (VARATT_IS_EXPANDED_HEADER(arr))
5923-
{
5924-
/* we can just grab the type data from expanded array */
5925-
fctx->elmlen = arr->xpn.typlen;
5926-
fctx->elmbyval = arr->xpn.typbyval;
5927-
fctx->elmalign = arr->xpn.typalign;
5928-
}
5929-
else
5930-
get_typlenbyvalalign(AARR_ELEMTYPE(arr),
5931-
&fctx->elmlen,
5932-
&fctx->elmbyval,
5933-
&fctx->elmalign);
5957+
Datum
5958+
array_unnest_element_index(PG_FUNCTION_ARGS)
5959+
{
5960+
FuncCallContext *funcctx;
5961+
ArrayUnnestState *fctx;
59345962

5935-
funcctx->user_fctx = fctx;
5936-
MemoryContextSwitchTo(oldcontext);
5937-
}
5963+
/* stuff done only on the first call of the function */
5964+
if (SRF_IS_FIRSTCALL())
5965+
init_array_unnest(SRF_FIRSTCALL_INIT(), PG_GETARG_DATUM(0), fcinfo);
59385966

59395967
/* stuff done on every call of the function */
59405968
funcctx = SRF_PERCALL_SETUP();
@@ -5944,11 +5972,20 @@ array_unnest(PG_FUNCTION_ARGS)
59445972
{
59455973
int offset = fctx->nextelem++;
59465974
Datum elem;
5975+
Datum vals[2];
5976+
bool nulls[2] = {false, false};
5977+
bool isnull;
59475978

5948-
elem = array_iter_next(&fctx->iter, &fcinfo->isnull, offset,
5979+
elem = array_iter_next(&fctx->iter, &isnull, offset,
59495980
fctx->elmlen, fctx->elmbyval, fctx->elmalign);
59505981

5951-
SRF_RETURN_NEXT(funcctx, elem);
5982+
if (isnull)
5983+
nulls[0] = true;
5984+
else
5985+
vals[0] = elem;
5986+
vals[1] = Int32GetDatum(fctx->nextelem);
5987+
5988+
SRF_RETURN_NEXT(funcctx, HeapTupleGetDatum(heap_form_tuple(fctx->tupdesc, vals, nulls)));
59525989
}
59535990
else
59545991
{

‎src/backend/utils/adt/jsonfuncs.c

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4003,6 +4003,8 @@ finiUnnest(UnnestState *state)
40034003
{
40044004
if (state)
40054005
{
4006+
if (state->tupdesc)
4007+
FreeTupleDesc(state->tupdesc);
40064008
pfree(state->jb);
40074009
pfree(state);
40084010
}

‎src/include/catalog/pg_proc.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -927,6 +927,8 @@ DESCR("expand array to set of rows");
927927
/* just for compatibility with jsonb_unnest* for any/each element clause */
928928
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_ ));
929929
DESCR("expand array to set of rows");
930+
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_ ));
931+
DESCR("expand array to set of rows with index");
930932
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_ ));
931933
DESCR("remove any occurrences of an element from an array");
932934
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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,7 @@ extern Datum array_fill(PG_FUNCTION_ARGS);
357357
extern Datum array_fill_with_lower_bounds(PG_FUNCTION_ARGS);
358358
extern Datum array_unnest(PG_FUNCTION_ARGS);
359359
extern Datum array_unnest_element(PG_FUNCTION_ARGS);
360+
extern Datum array_unnest_element_index(PG_FUNCTION_ARGS);
360361
extern Datum array_remove(PG_FUNCTION_ARGS);
361362
extern Datum array_replace(PG_FUNCTION_ARGS);
362363
extern Datum width_bucket_array(PG_FUNCTION_ARGS);

‎src/test/regress/expected/arrays.out

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2084,3 +2084,9 @@ SELECT a, (1 = ALL(a)), EACH ELEMENT OF a AS v SATISFIES (v = 1) FROM anytest;
20842084
{NULL,2} | f | f
20852085
(8 rows)
20862086

2087+
SELECT ANY ELEMENT OF '{1,2,3}'::int4[] AS e with index as i SATISFIES (e > 1 and i = 3);
2088+
bool_or_not_null
2089+
------------------
2090+
t
2091+
(1 row)
2092+

‎src/test/regress/expected/jsonb.out

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2174,6 +2174,12 @@ SELECT count(*) FROM testjsonb WHERE ANY KEY OF j AS k SATISFIES (
21742174
4
21752175
(1 row)
21762176

2177+
SELECT ANY ELEMENT OF '[1,2,3]'::jsonb AS e with index as i SATISFIES (e > '1'::jsonb and i = 1);
2178+
bool_or_not_null
2179+
------------------
2180+
t
2181+
(1 row)
2182+
21772183
CREATE INDEX jidx ON testjsonb USING gin (j);
21782184
SET enable_seqscan = off;
21792185
SELECT count(*) FROM testjsonb WHERE j @> '{"wait":null}';

‎src/test/regress/expected/jsonb_1.out

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2174,6 +2174,12 @@ SELECT count(*) FROM testjsonb WHERE ANY KEY OF j AS k SATISFIES (
21742174
4
21752175
(1 row)
21762176

2177+
SELECT ANY ELEMENT OF '[1,2,3]'::jsonb AS e with index as i SATISFIES (e > '1'::jsonb and i = 1);
2178+
bool_or_not_null
2179+
------------------
2180+
t
2181+
(1 row)
2182+
21772183
CREATE INDEX jidx ON testjsonb USING gin (j);
21782184
SET enable_seqscan = off;
21792185
SELECT count(*) FROM testjsonb WHERE j @> '{"wait":null}';

‎src/test/regress/sql/arrays.sql

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -619,3 +619,5 @@ insert into anytest values
619619

620620
SELECT a, (1 = ANY(a)), ANY ELEMENT OF a AS v SATISFIES (v = 1) FROM anytest;
621621
SELECT a, (1 = ALL(a)), EACH ELEMENT OF a AS v SATISFIES (v = 1) FROM anytest;
622+
623+
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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -530,6 +530,8 @@ SELECT count(*) FROM testjsonb WHERE ANY ELEMENT OF j->'array' AS e SATISFIES (
530530
SELECT count(*) FROM testjsonb WHERE ANY KEY OF j AS k SATISFIES (
531531
k = 'array' AND ANY ELEMENT OF j->k AS e SATISFIES ( e = '"baz"'::jsonb )
532532
);
533+
SELECT ANY ELEMENT OF '[1,2,3]'::jsonb AS e with index as i SATISFIES (e > '1'::jsonb and i = 1);
534+
533535

534536
CREATE INDEX jidx ON testjsonb USING gin (j);
535537
SET enable_seqscan = off;

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy