Connection.set_type_codec example for PostGIS geometry/geography types

[adamrothman]

I imagine there are at least a few people (like me, lol) using PostgreSQL with PostGIS wondering if there's an elegant, precise way to encode geometry/geography types for queries.

A naive approach might be to use the string-based functions:

await conn.fetchrow(
    'INSERT INTO locations (name, coords) VALUES ($1, ST_GeographyFromText($2))',
    'Empire State Building',
    'Point(-73.985661 40.748447)',
)

These string-based functions have the potential for precision loss (due to string conversions), and there are faster alternatives available, like ST_MakePoint:

ST_MakePoint while not being OGC compliant is generally faster and more precise than ST_GeomFromText and ST_PointFromText. It is also easier to use if you have raw coordinates rather than WKT.

You could do something like this:

await conn.fetchrow(
    'INSERT INTO locations (name, coords) VALUES ($1, ST_MakePoint($2, $3))',
    'Empire State Building',
    -73.985661,
    40.748447,
)

Better, but specific to the point type.

A fully general solution that suffers no loss of precision can be implemented using Connection.set_type_codec. Geometry/geography objects are encoded in the well-known binary format. It works for any Python object that conforms to the geo interface specification, and relies on the excellent Shapely library (but anything that speaks WKB will do).

from typing import Any

import shapely.geometry
import shapely.wkb
from shapely.geometry.base import BaseGeometry


def encode_geometry(geometry: Any) -> bytes:
    if not hasattr(geometry, '__geo_interface__'):
        raise TypeError(f'{geometry} does not conform to geo interface')
    shape = shapely.geometry.asShape(geometry)
    return shapely.wkb.dumps(shape)


def decode_geometry(wkb: bytes) -> BaseGeometry:
    return shapely.wkb.loads(wkb)


await conn.set_type_codec(
    'geometry',  # also works for 'geography'
    encoder=encode_geometry,
    decoder=decode_geometry,
    format='binary',
)

With a connection configured as above, the query is as simple as:

await conn.fetchrow(
    'INSERT INTO locations (name, coords) VALUES ($1, $2)',
    'Empire State Building',
    shapely.geometry.Point(-73.985661, 40.748447),  # or any other geo-interface compliant type
)

And it works for polygons and line strings too!

This may be rather niche but it took me a good deal of trial-and-error to figure out. It would be really cool to see an example like the above in the documentation. Maybe alongside the JSON example under Custom Type Conversions?

[elprans]

Thanks for providing this. I added an example in #247.

[adamrothman]

Thank you @elprans for maintaining such a great library!

[adamrothman]

Oops, didn’t mean to close this. Will let it auto-close when #247 is merged.

[sunilgopinath]

Hi, I love your library. I was trying to figure out how to add the SRID to shapely. Do you have any suggestions. I can't find where to insert it. The above code (and the one from the doc gives me):

asyncpg.exceptions.InvalidParameterValueError: Geometry SRID (0) does not match column SRID (4326)

[elprans]

I don't think that shapely has APIs to support SRIDs properly, but you can work on the underlying geos object like this:

from shapely import geos

geos.lgeos.GEOSSetSRID(shape._geom, 4326)

You'll also need to call dumps like this:

shapely.wkb.dumps(shape, include_srid=True)

[danni-tarrant-mwv]

If anyone ever finds themselves in the position where they are trying to absolutely maximise the performance of a bulk insert using an asyncpg copy_records_to_table or other such functions they might find that encoding the geometries using shapely adds a significant performance hit when you're trying to inserts millions of geometries.
I adapted some code from this repo to just encode points with as little overhead as possible. My testing found an improvement of loading a 400mb hdf5 file with ~5 million rows of which there were 4 geometry columns from taking 3m37s with shapely to 44s using the following method.

# The below code for encoding the postgis geometry for use with asyncpg was taken and edited from:
# https://github.com/thehappycheese/parse_wkb
# Some static variables used for encoding/decoding the postgis geometry
# The below are required settings used by shapely to give the byte object the correct endian form.
ByteOrderChar = str  # Literal[">", "<"]  # 0 = big = >, 1 = little = <
byte_order_char: ByteOrderChar = "<"
byte_order_byte = 1
UNINT32 = byte_order_char + "I"
BYTE_UNINT32 = byte_order_char + "bI"
BYTE_UNINT32_UNINT32 = byte_order_char + "bII"
map_geom_type_thousands_and_point_encoder = {
    2: (0, byte_order_char + "2d"),
    3: (1, byte_order_char + "3d"),
}



def get_geom_type_and_point_encoder(number_of_dimensions):
    geom_type_thousands, point_encoder = map_geom_type_thousands_and_point_encoder[number_of_dimensions]
    geom_type_units = 1
    geom_type = 1000 * geom_type_thousands + geom_type_units
    return geom_type, point_encoder


def encode_postgis_point(coordinates: Sequence[float]):
    number_of_dimensions = 2
    geom_type, point_encoder = get_geom_type_and_point_encoder(number_of_dimensions)
    result = pack(BYTE_UNINT32, byte_order_byte, geom_type)
    return result + pack(point_encoder, *coordinates)

I still used the above decode function from shapely as you need to pass both a decode/encode function to asyncpg as I was only encoding data in my code.
When adding a geometry I then pass it into asyncpg as a tuple like so:
(point.longitude, point.latitude)
Hopefully this might help anybody else that is looking to maximise insert performance for geometries using asyncpg. Hopefully one day I can come back and finish it off for other geometry types as it shouldn't be too much further work.