From 4720b07e9f18915cf1f6466245a92dedd11fdc55 Mon Sep 17 00:00:00 2001 From: Martin Kagamino Lehoux Date: Mon, 29 Sep 2025 18:48:48 +0200 Subject: Moved post to better named folder --- content/posts/postgres-aggregation-speed/tests.py | 190 ++++++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 content/posts/postgres-aggregation-speed/tests.py (limited to 'content/posts/postgres-aggregation-speed/tests.py') diff --git a/content/posts/postgres-aggregation-speed/tests.py b/content/posts/postgres-aggregation-speed/tests.py new file mode 100644 index 0000000..3c32db5 --- /dev/null +++ b/content/posts/postgres-aggregation-speed/tests.py @@ -0,0 +1,190 @@ +""" +uv run pytest --benchmark-columns=mean,stddev,iqr +""" + +from decimal import Decimal +from datetime import date +from typing import Generator +import psycopg +from psycopg import Connection +from psycopg.abc import ConnectionType +from pytest_benchmark.fixture import BenchmarkFixture +import pytest + + +@pytest.fixture(scope="module") +def conn() -> Generator[Connection]: + with psycopg.connect( + dbname="postgres", user="postgres", password="", host="127.0.0.1", port="5432" + ) as connection: + yield connection + + +BASE_QUERY = """SELECT accommodation_id, {} AS daily_configs + FROM overrides + WHERE '[2025-07-01,2025-08-01)'::daterange @> date + GROUP BY overrides.accommodation_id""" +COMPOSITE_QUERY = BASE_QUERY.format("array_agg((date, price))") +JSON_OBJECT_EXPR = "json_agg(json_build_object('date', date, 'price', price::text))" +JSONB_OBJECT_EXPR = "jsonb_agg(jsonb_build_object('date', date, 'price', price::text))" +JSONB_ARRAY_EXPR = "jsonb_agg(jsonb_build_array(date, price::text))" +JSON_ARRAY_EXPR = "json_agg(json_build_array(date, price::text))" + +type PythonOut = list[tuple[str, list[tuple[date, Decimal]]]] +type JSONOut = list[tuple[str, list[dict[str, str]]]] + + +def load_any_json_array_python( + conn: Connection, query: bytes, *, binary: bool +) -> PythonOut: + with conn.cursor() as cursor: + cursor.execute(query, binary=binary) + rows = cursor.fetchall() + + return [ + ( + row[0], + [(date.fromisoformat(item[0]), Decimal(item[1])) for item in row[1]], + ) + for row in rows + ] + + +def load_any_json_array_json( + conn: Connection, query: bytes, *, binary: bool +) -> JSONOut: + with conn.cursor() as cursor: + cursor.execute(query, binary=binary) + rows = cursor.fetchall() + + return [ + ( + row[0], + [{"date": item[0], "price": item[1]} for item in row[1]], + ) + for row in rows + ] + + +def load_any_json_object_python( + conn: Connection, query: bytes, *, binary: bool +) -> PythonOut: + with conn.cursor() as cursor: + cursor.execute(query, binary=binary) + rows = cursor.fetchall() + + return [ + ( + row[0], + [ + (date.fromisoformat(item["date"]), Decimal(item["price"])) + for item in row[1] + ], + ) + for row in rows + ] + + +def load_any_json_object_json( + conn: Connection, query: bytes, *, binary: bool +) -> PythonOut: + with conn.cursor() as cursor: + cursor.execute(query, binary=binary) + return cursor.fetchall() + + +def load_composite_text_python(conn: Connection) -> PythonOut: + with conn.cursor() as cursor: + cursor.execute(COMPOSITE_QUERY, binary=False) + rows = cursor.fetchall() + + return [ + ( + row[0], + [(date.fromisoformat(item[0]), Decimal(item[1])) for item in row[1]], + ) + for row in rows + ] + + +def load_composite_text_json(conn: Connection) -> JSONOut: + with conn.cursor() as cursor: + cursor.execute(COMPOSITE_QUERY, binary=False) + rows = cursor.fetchall() + + return [ + (row[0], [{"date": item[0], "price": item[1]} for item in row[1]]) + for row in rows + ] + + +def load_composite_binary_python(conn: Connection) -> PythonOut: + with conn.cursor() as cursor: + cursor.execute(COMPOSITE_QUERY, binary=True) + return cursor.fetchall() + + +def load_composite_binary_json(conn: Connection) -> JSONOut: + with conn.cursor() as cursor: + cursor.execute(COMPOSITE_QUERY, binary=True) + rows = cursor.fetchall() + + return [ + ( + row[0], + [{"date": item[0].isoformat(), "price": str(item[1])} for item in row[1]], + ) + for row in rows + ] + + +@pytest.mark.benchmark(group="json out") +class TestJSONOut: + @pytest.mark.parametrize("binary", [False, True]) + @pytest.mark.parametrize( + "expr", + [JSON_OBJECT_EXPR, JSONB_OBJECT_EXPR], + ids=["json object", "jsonb object"], + ) + def test_json_object( + self, + conn: Connection, + benchmark: BenchmarkFixture, + binary: bool, + expr: str, + ): + benchmark( + load_any_json_object_json, conn, BASE_QUERY.format(expr), binary=binary + ) + + def test_composite_binary(self, conn: Connection, benchmark: BenchmarkFixture): + benchmark(load_composite_binary_json, conn) + + def test_composite_text(self, conn: Connection, benchmark: BenchmarkFixture): + benchmark(load_composite_text_json, conn) + + +@pytest.mark.benchmark(group="python out") +class TestPythonOut: + @pytest.mark.parametrize("binary", [False, True]) + @pytest.mark.parametrize( + "expr", + [JSON_ARRAY_EXPR, JSONB_ARRAY_EXPR], + ids=["json array", "jsonb array"], + ) + def test_json_array( + self, + conn: Connection, + benchmark: BenchmarkFixture, + binary: bool, + expr: str, + ): + benchmark( + load_any_json_array_python, conn, BASE_QUERY.format(expr), binary=binary + ) + + def test_composite_binary(self, conn: Connection, benchmark: BenchmarkFixture): + benchmark(load_composite_binary_python, conn) + + def test_composite_text(self, conn: Connection, benchmark: BenchmarkFixture): + benchmark(load_composite_text_python, conn) -- cgit v1.3