From 637302a61ee9b4ef84e5c152a17af442599a8e2a Mon Sep 17 00:00:00 2001 From: Martin Kagamino Lehoux Date: Mon, 29 Sep 2025 18:42:44 +0200 Subject: Draft ready for PostgreSQL aggregation processing speed --- content/posts/nested-postgres-results/tests.py | 197 ++++++++++++++++++++++--- 1 file changed, 175 insertions(+), 22 deletions(-) (limited to 'content/posts/nested-postgres-results/tests.py') diff --git a/content/posts/nested-postgres-results/tests.py b/content/posts/nested-postgres-results/tests.py index c1f1506..3c32db5 100644 --- a/content/posts/nested-postgres-results/tests.py +++ b/content/posts/nested-postgres-results/tests.py @@ -1,37 +1,190 @@ -from pathlib import Path +""" +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 connection() -> Generator[psycopg.Connection]: +def conn() -> Generator[Connection]: with psycopg.connect( dbname="postgres", user="postgres", password="", host="127.0.0.1", port="5432" ) as connection: yield connection -def load(connection: psycopg.Connection, query: bytes, *, binary: bool) -> None: - with connection.cursor() as cursor: +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) - cursor.fetchall() - - -@pytest.mark.parametrize("binary", [False, True]) -@pytest.mark.parametrize( - "query", ["json-array-agg.sql", "jsonb-array-agg.sql", "composite-array-agg.sql"] -) -def test_benchmark_load( - connection: psycopg.Connection, - benchmark: BenchmarkFixture, - binary: bool, - query: str, -): - benchmark( - load, - connection, - (Path("content/posts/nested-postgres-results") / query).read_bytes(), - 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