Fix null-bitmap combining in array_agg_array_combine().

This code missed the need to update the combined state's
nullbitmap if state1 already had a bitmap but state2 didn't.
We need to extend the existing bitmap with 1's but didn't.
This could result in wrong output from a parallelized
array_agg(anyarray) calculation, if the input has a mix of
null and non-null elements.  The errors depended on timing
of the parallel workers, and therefore would vary from one
run to another.

Also install guards against integer overflow when calculating
the combined object's sizes, and make some trivial cosmetic
improvements.

Author: Dmytro Astapov <dastapov@gmail.com>
Reviewed-by: Tom Lane <tgl@sss.pgh.pa.us>
Discussion: https://postgr.es/m/CAFQUnFj2pQ1HbGp69+w2fKqARSfGhAi9UOb+JjyExp7kx3gsqA@mail.gmail.com
Backpatch-through: 16

diff --git a/src/backend/utils/adt/array_userfuncs.c b/src/backend/utils/adt/array_userfuncs.c
index 5c4fdcfba4693dcf6d83047806bd62cbe3adfd1d..69e9cf13e76d471fe0f61604d9849496296ced45 100644 (file)
--- a/src/backend/utils/adt/array_userfuncs.c
+++ b/src/backend/utils/adt/array_userfuncs.c
@@ -21,6 +21,7 @@
 #include "utils/datum.h"
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
+#include "utils/memutils.h"
 #include "utils/typcache.h"
 
 /*
@@ -969,10 +970,11 @@ array_agg_array_combine(PG_FUNCTION_ARGS)
        }
 
        /* We only need to combine the two states if state2 has any items */
-       else if (state2->nitems > 0)
+       if (state2->nitems > 0)
        {
                MemoryContext oldContext;
-               int                     reqsize = state1->nbytes + state2->nbytes;
+               int                     reqsize;
+               int                     newnitems;
                int                     i;
 
                /*
@@ -995,6 +997,17 @@ array_agg_array_combine(PG_FUNCTION_ARGS)
                                                 errmsg("cannot accumulate arrays of different dimensionality")));
                }
 
+               /* Types should match already. */
+               Assert(state1->array_type == state2->array_type);
+               Assert(state1->element_type == state2->element_type);
+
+               /* Calculate new sizes, guarding against overflow. */
+               if (pg_add_s32_overflow(state1->nbytes, state2->nbytes, &reqsize) ||
+                       pg_add_s32_overflow(state1->nitems, state2->nitems, &newnitems))
+                       ereport(ERROR,
+                                       (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+                                        errmsg("array size exceeds the maximum allowed (%zu)",
+                                                       MaxArraySize)));
 
                oldContext = MemoryContextSwitchTo(state1->mcontext);
 
@@ -1009,17 +1022,16 @@ array_agg_array_combine(PG_FUNCTION_ARGS)
                        state1->data = (char *) repalloc(state1->data, state1->abytes);
                }
 
-               if (state2->nullbitmap)
+               /* Combine the null bitmaps, if present. */
+               if (state1->nullbitmap || state2->nullbitmap)
                {
-                       int                     newnitems = state1->nitems + state2->nitems;
-
                        if (state1->nullbitmap == NULL)
                        {
                                /*
                                 * First input with nulls; we must retrospectively handle any
                                 * previous inputs by marking all their items non-null.
                                 */
-                               state1->aitems = pg_nextpower2_32(Max(256, newnitems + 1));
+                               state1->aitems = pg_nextpower2_32(Max(256, newnitems));
                                state1->nullbitmap = (bits8 *) palloc((state1->aitems + 7) / 8);
                                array_bitmap_copy(state1->nullbitmap, 0,
                                                                  NULL, 0,
@@ -1027,17 +1039,17 @@ array_agg_array_combine(PG_FUNCTION_ARGS)
                        }
                        else if (newnitems > state1->aitems)
                        {
-                               int                     newaitems = state1->aitems + state2->aitems;
-
-                               state1->aitems = pg_nextpower2_32(newaitems);
+                               state1->aitems = pg_nextpower2_32(newnitems);
                                state1->nullbitmap = (bits8 *)
                                        repalloc(state1->nullbitmap, (state1->aitems + 7) / 8);
                        }
+                       /* This will do the right thing if state2->nullbitmap is NULL: */
                        array_bitmap_copy(state1->nullbitmap, state1->nitems,
                                                          state2->nullbitmap, 0,
                                                          state2->nitems);
                }
 
+               /* Finally, combine the data and adjust sizes. */
                memcpy(state1->data + state1->nbytes, state2->data, state2->nbytes);
                state1->nbytes += state2->nbytes;
                state1->nitems += state2->nitems;
@@ -1045,9 +1057,6 @@ array_agg_array_combine(PG_FUNCTION_ARGS)
                state1->dims[0] += state2->dims[0];
                /* remaining dims already match, per test above */
 
-               Assert(state1->array_type == state2->array_type);
-               Assert(state1->element_type == state2->element_type);
-
                MemoryContextSwitchTo(oldContext);
        }