binary representation of numpy.complex and postgresql composite

[GFuhr]

• asyncpg version: 0.28.0
• PostgreSQL version: 15 (latest docker container)
• Do you use a PostgreSQL SaaS? If so, which? Can you reproduce the issue with a local PostgreSQL install?:
• Python version: 3.10
• Platform: windows or debian in WSL
• Do you use pgbouncer?: no
• Did you install asyncpg with pip?: yes
• If you built asyncpg locally, which version of Cython did you use?:
• Can the issue be reproduced under both asyncio and uvloop?: yes

Hi,

I need for a project to store complex numbers in a database (numpy.complex64 exactly) and for that I created a postgresql composite datatype.

I was able to make it work with asyncpg in text format for simple INSERT, SELECT... after some works and patches found on previous issues report.

However for the binary format I have weird errors : asyncpg.exceptions.DatatypeMismatchError: wrong number of columns: 1082549862, expected 2

I made a "simple" python script to reproduce the issue :

from __future__ import annotations

import asyncpg
import asyncio
import numpy as np

from ast import literal_eval
import struct

SQL_CREATE = """
DO $$ BEGIN
    IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'complex') THEN
        CREATE TYPE complex AS (
            r float4,
            i float4
        );
    END IF;
    CREATE TABLE dummy_table (dummy_column complex);
    DROP TABLE dummy_table;
END $$;

DROP TABLE IF EXISTS "poc_asyncpg";

CREATE TABLE "poc_asyncpg" (
    "id" SERIAL PRIMARY KEY,
    "float_value" float4 NULL,
    "complex_value" complex NULL,
    "complex_array" complex[] NULL
)
"""

def _cplx_decode(val) -> np.complex64:
    cplx = complex(*literal_eval(val))
    return np.complex64(cplx)

def _cplx_encode(val: np.complex64 | complex) -> str:
    return str((np.float32(val.real), np.float32(val.imag)))

async def set_type_codec(conn):
    """
    had to use this patch since the conn.set_type_codec does not work for scalar variables
    """
    schema = 'public'
    format = 'text'
    conn._check_open()
    typenames = ('complex',)
    for typename in typenames:
        typeinfo = await conn.fetchrow(
            asyncpg.introspection.TYPE_BY_NAME, typename, schema)
        if not typeinfo:
            raise ValueError('unknown type: {}.{}'.format(schema, typename))

        oid = typeinfo['oid']
        conn._protocol.get_settings().add_python_codec(
            oid, typename, schema, 'scalar',
            lambda a: _cplx_encode(a), lambda a: _cplx_decode(a), format)

# if this part is commented, error message is : 
# asyncpg.exceptions._base.InternalClientError: no binary format encoder for type complex    
        conn._protocol.get_settings().add_python_codec(
            oid, typename, schema, 'scalar',
            encoder=lambda x: struct.pack('!2f', x.real, x.imag),
            decoder=lambda x: np.frombuffer(x, dtype=np.complex64)[0],
            format="binary")

    # Statement cache is no longer valid due to codec changes.
    conn._drop_local_statement_cache()

async def init_connection(conn):
    await set_type_codec(conn)
    await conn.set_type_codec(
        'numeric', encoder=str, decoder=np.float32,
        schema='pg_catalog', format='text'
    )
    await conn.set_type_codec(
        'float4', encoder=str, decoder=np.float32,
        schema='pg_catalog', format='text'
    )
    await conn.set_type_codec(
        'float4', encoder=struct.Struct("!f").pack, decoder=struct.Struct("!f").unpack,
        schema='pg_catalog', format='binary'
    )

async def trunc(pool):
    async with pool.acquire() as conn:
        async with conn.transaction():
            query = "TRUNCATE poc_asyncpg"
            await conn.execute(query)

async def worker_copy(pool, column_name, data):
    await trunc(pool)
    async with pool.acquire() as conn:
        async with conn.transaction():
            await conn.copy_records_to_table("poc_asyncpg",
                                             records=[(data,)],
                                             columns=(column_name,)
                                             )

try:
    from common import dbinfo
except ImportError as e:
    class DB:
        user = "user"
        password = "password"
        database = "db"

    dbinfo = DB()

def create_pool(db_info):
    pool = asyncpg.create_pool(
        user=db_info.user,
        password=db_info.password,
        database=dbinfo.database,
        host="127.0.0.1",
        init=init_connection
    )
    return pool

async def main(info):
    pool = await create_pool(info)
    async with pool.acquire() as conn:
        async with conn.transaction():
            await conn.execute(SQL_CREATE)

    await worker_copy(pool, "float_value", np.float32(4.2))
    await worker_copy(pool, "complex_value", np.complex64(4.2 + 1j * 4.2))

if __name__ == '__main__':
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(main(dbinfo))

[elprans]

This is because a composite type isn't a "scalar" and requires specific representation on the wire. Ideally, asyncpg should learn to support format="tuple" codecs for user-defined composites, where you would only need to convert np.complex64 to a tuple of real and imaginary parts and vice-versa, e.g:

    await conn.set_type_codec(
        'complex',
        schema='public',
        encoder=lambda x: (x.real, x.imag),
        decoder=lambda t: np.complex64(t[0] + 1j * t[1]),
        format='tuple',
    )

[GFuhr]

thanks a lot, I can't wait for this update :)