Sunday, December 20, 2015

FORTRAN 90 like vector operations in PostgreSQL

Let's say you have two vectors A and B and want to do element wise operations, creating a new vector C, FORTRAN 90 (and many other languages now) has a concise notation for that. E.g.

C = A + B

does the element wise addition of A and B into C.

Can we have that in PostgreSQL too?

CREATE OR REPLACE FUNCTION _vectoradd(
    v1 anyarray,
    v2 anyarray)
  RETURNS anyarray AS
$BODY$
DECLARE
BEGIN
if array_ndims(v1) != 1 OR array_ndims(v2) != 1 Then
RAISE EXCEPTION 'Only vectors supported';
end if;

if array_length(v1,1) != array_length(v2,1)  Then
RAISE EXCEPTION 'Vector size mismatch';
end if;

-- return array(select v.a + v.b from (select unnest(v1) as a, unnest(v2) as b) v); --9.4
return array(select a+b from unnest(v1,v2) v (a,b)); -- 9.5
END;
$BODY$
  LANGUAGE plpgsql immutable strict;

CREATE OPERATOR + (
    LEFTARG = anyarray,
    RIGHTARG = anyarray,
    PROCEDURE = _vectoradd,
    COMMUTATOR = +);

select ARRAY[1,2,3,4] + ARRAY[1,2,3,4];

{2,4,6,8}

Yes, we can! :-)

2 comments:

  1. FWIW - here is an extension that achieves similar goals - https://github.com/pramsey/pgsql-arraymath

    ReplyDelete
    Replies
    1. Damn. Every time I come up with something, somebody else has already done it. :-)

      And in C.

      All I can say in defense is, that my example is 100% SQL and the use of unnest() instead of a slow element-by-index loop.

      BTW: That C extension is probably a good candidate for true parallelization.

      Delete