Support binary protocol for hstore

[Zaczero]

When using the binary protocol, hstore columns are not correctly decoded. Related code. This results in type errors because the data is received as raw bytes rather than a Python dictionary. I went ahead and implemented missing binary loader/dumper classes:

_U32_STRUCT = Struct('!I')
_I2B = {i: i.to_bytes(4) for i in range(256)}

class HstoreBinaryLoader(RecursiveLoader):
    format = Format.BINARY
    _encoding: str

    def load(self, data: Buffer) -> dict[str, str | None]:
        if len(data) < 12:  # Fast-path if too small to contain any data.
            return {}

        unpack_from = _U32_STRUCT.unpack_from
        encoding = self._encoding
        result = {}

        view = bytes(data)
        size, = unpack_from(view)
        pos = 4

        for _ in range(size):
            key_size, = unpack_from(view, pos)
            pos += 4

            key = view[pos : pos + key_size].decode(encoding)
            pos += key_size

            value_size, = unpack_from(view, pos)
            pos += 4

            if value_size == 0xFFFFFFFF:
                value = None
            else:
                value = view[pos : pos + value_size].decode(encoding)
                pos += value_size

            result[key] = value

        return result


class HstoreBinaryDumper(RecursiveDumper):
    format = Format.BINARY
    _encoding: str

    def dump(self, obj: dict[str, str | None]) -> Buffer:
        if not obj:
            return b'\x00\x00\x00\x00'

        pack = _U32_STRUCT.pack
        i2b = _I2B
        encoding = self._encoding
        buffer: list[bytes] = [i2b.get(l := len(obj)) or pack(l)]

        for key, value in obj.items():
            key_bytes = key.encode(encoding)
            buffer.append(i2b.get(l := len(key_bytes)) or pack(l))
            buffer.append(key_bytes)

            if value is None:
                buffer.append(b'\xFF\xFF\xFF\xFF')
            else:
                value_bytes = value.encode(encoding)
                buffer.append(i2b.get(l := len(value_bytes)) or pack(l))
                buffer.append(value_bytes)

        return b''.join(buffer)

_encoding is the current connection encoding (from psycopg._encodings import conn_encoding).

Note

If you're planning to use any part or all of the suggested code, please add me to the commit's Co-authored-by list. I enjoy collecting contributions to open source projects 🙂.

PS.

I saw #1028 (comment) but I can't guarantee I'll be able to submit a PR soon. I may or may not be able to - it depends on my other work. So feel free to integrate it early (or not) if you have spare capacity.

[dvarrazzo]

I was a fan of the pack_into=struct.pack_into hack, but from some measurements I took not many years ago it doesn't make anymore a substantial difference, in modern versions of Python.

For the rest, thank you very much for this contribution, but I don't have time either to verify it works correctly. After sorting out other priorities I can look into it, but if a MR reaches this project sooner that would be great.

Cheers!

[Zaczero]

I was curious and checked it. It seems to speed up execution by about 2% in this test (10k dict size):

import struct
from timeit import timeit


def dump(obj: dict[str, str | None], *, pack_into=struct.pack_into):
    if not obj:
        return b'\x00\x00\x00\x00'

    pos = 4
    encoding = 'utf-8'
    encoded: list[tuple[bytes, bytes | None]] = []

    for key, value in obj.items():
        key_bytes = bytes(key, encoding)
        pos += 4 + len(key_bytes)

        if value is None:
            value_bytes = None
            pos += 4
        else:
            value_bytes = bytes(value, encoding)
            pos += 4 + len(value_bytes)

        encoded.append((key_bytes, value_bytes))

    buffer = bytearray(pos)
    pack_into('>I', buffer, 0, len(obj))
    pos = 4

    for key_bytes, value_bytes in encoded:
        key_len = len(key_bytes)
        pack_into('>I', buffer, pos, key_len)
        pos += 4

        buffer[pos : pos + key_len] = key_bytes
        pos += key_len

        if value_bytes is None:
            pack_into('>I', buffer, pos, 0xFFFFFFFF)
            pos += 4
        else:
            value_len = len(value_bytes)
            pack_into('>I', buffer, pos, value_len)
            pos += 4

            buffer[pos : pos + value_len] = value_bytes
            pos += value_len

    return buffer

def dump2(obj: dict[str, str | None]):
    if not obj:
        return b'\x00\x00\x00\x00'

    pos = 4
    encoding = 'utf-8'
    encoded: list[tuple[bytes, bytes | None]] = []

    for key, value in obj.items():
        key_bytes = bytes(key, encoding)
        pos += 4 + len(key_bytes)

        if value is None:
            value_bytes = None
            pos += 4
        else:
            value_bytes = bytes(value, encoding)
            pos += 4 + len(value_bytes)

        encoded.append((key_bytes, value_bytes))

    buffer = bytearray(pos)
    struct.pack_into('>I', buffer, 0, len(obj))
    pos = 4

    for key_bytes, value_bytes in encoded:
        key_len = len(key_bytes)
        struct.pack_into('>I', buffer, pos, key_len)
        pos += 4

        buffer[pos : pos + key_len] = key_bytes
        pos += key_len

        if value_bytes is None:
            struct.pack_into('>I', buffer, pos, 0xFFFFFFFF)
            pos += 4
        else:
            value_len = len(value_bytes)
            struct.pack_into('>I', buffer, pos, value_len)
            pos += 4

            buffer[pos : pos + value_len] = value_bytes
            pos += value_len

    return buffer


_DICT = {str(i): str(i) for i in range(10000)}

print(timeit(lambda: dump(_DICT), number=1000))
# 8.355752156000108
print(timeit(lambda: dump2(_DICT), number=1000))
# 8.540168364999772

And the related bytecode difference (because I like looking at it):

dump:
 25           LOAD_FAST                1 (pack_into)
              PUSH_NULL
              LOAD_CONST               4 ('>I')
              LOAD_FAST                9 (buffer)
              LOAD_CONST               5 (0)
              LOAD_GLOBAL              5 (len + NULL)
              LOAD_FAST                0 (obj)
              CALL                     1
              CALL                     4
              POP_TOP

dump2:
 25           LOAD_GLOBAL             10 (struct)
              LOAD_ATTR               12 (pack_into)
              PUSH_NULL
              LOAD_CONST               4 ('>I')
              LOAD_FAST                8 (buffer)
              LOAD_CONST               5 (0)
              LOAD_GLOBAL              5 (len + NULL)
              LOAD_FAST                0 (obj)
              CALL                     1
              CALL                     4
              POP_TOP

Personally, I would leave it there because it reduces code verbosity and (minimally) improves the performance. LOAD_FAST will always be slightly faster than LOAD_GLOBAL+LOAD_ATTR. But I'll also understand if you would rather not use such HACKs in the codebase.

[dvarrazzo]

You can speed up the execution 20x by writing this stuff in Cython, just sayin'.

Using pack_into() from a single Struct(">I") instance is more than 40% faster than using struct.pack_into(">I").

Finding little %'s is simple.

[Zaczero]

TIL there exists Struct class 😮. I'll try it out and update the sample code if it works as you say!

Yes Cython is great (and its pure python mode is lovely). But I also find beauty in well performing pure Python code 🌞. I think it's more future proofed that way: no need for binaries, which is even more of an issue with the new free-threaded builds.

[dvarrazzo]

Also, to clarify, what I meant for

it doesn't make anymore a substantial difference, in modern versions of Python.

is that in older Python versions the difference was much more substantial than 2% 🙂 (the difference between LOAD_GLOBAL and LOAD_FAST was much more marked)

[Zaczero]

I went ahead and did some testing and optimized the code with focus on performance. Here are my findings:

Dumper

• Initial time was ~8.5 sec for my test case
• Switching to Struct('!I') reduced the time to ~7.5 sec
• Replacing bytearray with b''.join(list[bytes]) reduced the time to ~3.5 sec
• Replacing bytes(str) with str.encode reduced the time to ~2.8 sec
• EDIT: Caching int->bytes for all sizes below 256, reduced the time to ~2.15 sec. I think this is a sensible change given how common such sizes are, and how little memory it takes.

Loader

• Initial time was ~7.3 sec for my test case
• Switching to Struct('!I') reduced the time to ~5.6 sec
• Interestingly, replacing memoryview with bytes(data) reduced the time to ~4.7 sec. I suspect Python's internal bytes-type optimizations outweigh the bytes slice copy cost.
• Replacing str(bytes) with bytes.decode reduced the time minimally to just under 4.7 sec.

I have already updated the issue's suggested code to the new version. I'll do one more test to see whether caching would be any useful here. EDIT: I have did some tests with caching and only Dumper class saw an improvement. Implementing caching for Loader size decoding required additional bytes slice copy + bytes hash, which seems to have caused too much of a slowdown.

[dvarrazzo]

I am really baffled about why you are willing to spend time on this issue without opening a MR which would allow to track the history of the changes.

Your code above doesn't include the adapters registration and the tests aren't adapted to run in binary mode. As it is, it cannot be tested neither by the CI nor by a dev.

If you just want to share a recipe, that's fine, I really appreciate your contribution, and your involvement can finish here. As you respectfully request we would include your name in the contribution, and we will act on these ticket by converting them in complete code when we have time to do it.

If you want to refine the solution until it's good enough for inclusion, we need a MR and we can review the code and its behaviour. If you are not available to do that, I appreciate it and it's absolutely no problem. However, I am not available to look further or to code-review an ever-changing ticket description.

[Zaczero]

Continued in #1031.

why you are willing to spend time on this issue without opening a MR

Because "Finding little %'s is simple." sounded like a challenge and I also got into the Struct rabbit hole. I didn't intend to work on this today but somehow I did.

[dvarrazzo]

Because "Finding little %'s is simple." sounded like a challenge

I totally resonate with this 😄 I have at least two other full time job going on, but few things beat the pleasure of micro-optimization...

Thank you very much for opening the MR. See you there!