Draft ready for PostgreSQL aggregation processing speed

1 files changed, 173 insertions, 20 deletions

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)
-        cursor.fetchall()
+        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()
 
-@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,
+    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)