Search code examples
sqlarrayspostgresqlset-returning-functions

Multiply elements in an array according to their position


What is an efficient function I could use to increase the amount of elements in an array by their position?

For example, if I have an array of:

ARRAY[10,20,30,40,50]::INT[]

I would like that to turn into:

ARRAY[10,20,20,30,30,30,40,40,40,40,50,50,50,50,50]::INT[]

So that I have 1 "10", 2 "20", 3 "30", 4 "40", and 5 "50".

For the array:

ARRAY[10,20,30,40,10]::INT[]

I would like that to turn into:

ARRAY[10,20,20,30,30,30,40,40,40,40,10,10,10,10,10]::INT[]

So that I have 6 "10", 2 "20", 3 "30" and 4 "40".


Solution

  • WITH t AS (SELECT ARRAY[10,20,30,40,50]::INT[] AS arr)   -- variable for demo
    SELECT ARRAY(
       SELECT unnest(array_fill(arr[idx], ARRAY[idx])) AS mult
       FROM  (SELECT arr, generate_subscripts(arr, 1) AS idx FROM t) sub
       );
    

    I would wrap the logic into a simple IMMUTABLE SQL function:

    CREATE OR REPLACE FUNCTION f_expand_arr(_arr anyarray)
      RETURNS anyarray AS
    $func$
    SELECT ARRAY(
       SELECT unnest(array_fill(_arr[idx], ARRAY[idx]))
       FROM  (SELECT generate_subscripts(_arr, 1) AS idx) sub
       )
    $func$ LANGUAGE sql IMMUTABLE;
    

    Works for arrays of any base type due to the polymorphic parameter type anyarray:
    How to write a function that returns text or integer values?

    The manual on generate_subscripts() and array_fill().

    Note: This works with the actual array indexes, which can differ from the ordinal array position in Postgres. You may be interested in @Daniel's method to "normalize" the array index:
    Normalize array subscripts for 1-dimensional array so they start with 1

    The upcoming Postgres 9.4 (currently beta) provides WITH ORDINALITY:
    PostgreSQL unnest() with element number

    Allowing for this even more elegant and reliable solution:

    CREATE OR REPLACE FUNCTION f_expand_arr(_arr anyarray)
      RETURNS anyarray AS
    $func$
    SELECT ARRAY(
       SELECT unnest(array_fill(a, ARRAY[idx]))
       FROM   unnest(_arr) WITH ORDINALITY AS x (a, idx)
       )
    $func$ LANGUAGE sql IMMUTABLE;
    

    One might still argue that proper order is not actually guaranteed. I claim it is ...
    Parallel unnest() and sort order in PostgreSQL

    Call:

    SELECT f_expand_arr(ARRAY[10,20,30,40,10]::INT[]) AS a2;
    

    Or for values from a table:

    SELECT f_expand_arr(a) AS a2 FROM t;
    

    SQL Fiddle.