Adapting asyncpg for numpy

[j-carson]

I'm working on writing some code that will pull large engineering datasets from Postgres and process them using numpy. I've written type-converters to override the defaults and return numpy types from my queries, but I'm concerned that this may be less than ideal performance-wise for large datasets. (I am currently just prototyping using toy data.)

• Is there a "standard way" to accomplish what I'm doing below?
• Assuming I've implemented proper type converters, I'm confused about how to set up type conversion when using a connection pool -- do I just register my converters every time I call pool.aquire()? There doesn't seem to be an API to set the converters for the pool in one step.

Thank you for any hints!

import asyncpg 
import numpy as np

def float32_decoder( buf ):
    return np.frombuffer( buf, dtype='>f4')[0]

def float32_encoder( val ):
    if not isinstance(val, np.float32):
        val = np.float32(val)
    return val.newbyteorder('>').tobytes()

def int32_encoder( val ):
    if not isinstance(val, np.int32):
        val = np.int32(val)
    return val.newbyteorder('>').tobytes()

def int32_decoder( buf ):
    return np.frombuffer( buf, dtype='>i4')[0]

def int16_encoder( val ):
    if not isinstance(val, np.int16):
        val = np.int16(val)
    return val.newbyteorder('>').tobytes()
        
def int16_decoder( buf ):
    return np.frombuffer( buf, dtype='>i2')[0]

async def register_adapters(conn: asyncpg.Connection):
    """
    register type converters to conn
    """
    await conn.set_type_codec(
        'real', 
        encoder = float32_encoder,
        decoder = float32_decoder,
        schema = 'pg_catalog',
        format = 'binary'
    )
    await conn.set_type_codec(
        'integer',
        encoder = int32_encoder,
        decoder = int32_decoder,
        schema = 'pg_catalog',
        format = 'binary'
    )
    await conn.set_type_codec(
        'smallint',
        encoder = int16_encoder,
        decoder = int16_decoder,
        schema = 'pg_catalog',
        format = 'binary'
    )

[elprans]

There doesn't seem to be an API to set the converters for the pool in one step.

There is: init callback passed to create_pool.

[j-carson]

Thanks! I was looking in the type conversion docs and not enough at create_pool! Do you have any suggestions about improvements to performance of type converters?

[elprans]

Your general approach is correct, except if you want absolute best performance, I'd do numpy datatype initialization once, instead of for every value:

def int32_decoder(buf, dt=np.dtype(np.int32).newbyteorder('>')):
    return np.frombuffer(buf, dtype=dt)[0]

[j-carson]

Thanks! I'll try that out!

[j-carson]

Thanks! I'll try that out!

The numpy type initialization definitely helped the performance, but it is still noticeably slower with type converters. Thank you for your help.

[j-carson]

So after trying various options, it seems like the most performant one is to let asyncpg return the native python type with its faster code and cast it back to the desired "right-sized" type later.