aggregation.qmd

# Aggregation

## 3rd exercice

Let's continue with another book endpoint
This time we want to retrieve paginated list of books, containing some additional information

- List of tag names
- Count of reviews
- Author

```json
[
  {
    "id": 1213190,
    "title": "Country performance attack approach painting soon us product sometimes should yeah.",
    "author_id": 9199,
    "release_date": "1923-05-31",
    "library_id": 6,
    "review_count": 20,
    "tag_names": ["braille", "german", "movie", "audio"],
    "author": {
      "id": 9199,
      "name": "Katherine Miller"
    }
  }
]
```

## 1st try

```{python}
# | echo: false

import django_init

from rich import print
from rich.console import Console
from rich.table import Table
from rich.markdown import Markdown

from django.db import connection
from django.db.models import QuerySet, Q
from books.models import *
from utils.perf_display import perf_counter, format_duration
from rest_framework.test import APIClient
from collections import defaultdict


from rest_framework.test import APIClient

from utils.sql import use_indexes, disable_indexes
console = Console()
client = APIClient()

```

```{python}
# | echo: false
from django.contrib.postgres.aggregates import ArrayAgg
from django.contrib.postgres.expressions import ArraySubquery
from django.db.models import Count, OuterRef, Q, QuerySet
from django.db.models.functions import Coalesce

from books.models import Book, BookTag, Review
from books.views.book.list_books_aggregate import serialize_books
```

```{python}
# | echo: false
from django.db.models import Count

# id not from the 10 libraries with a lot of books
personal = Library.objects.get(id=100)
personal_id = personal.id

public = Library.objects.order_by("-id").first()
public_id = public.id

# Get the library with the most books
alexandria_id, book_count = (
    Book.objects.values("library_id")
    .annotate(book_count=Count("id"))
    .order_by("-book_count")
    .values_list("library_id", "book_count")
    .first()
)

alexandria = Library.objects.get(id=alexandria_id)
```

```{python}
with perf_counter(time_sql=True):
    client.get(f"/books/{public_id}/aggregate?ordering=release_date")
```

## Looking at the first implementation {.scrollable }

```{python}
def get_queryset(library_id, page_size=20):
    return (
        Book.objects.filter(library_id=library_id)
        .order_by("release_date")
        .select_related("author")[:page_size]
    )

def annotate_books_aggregate(book_qs: QuerySet[Book]):
    """
    return a list of book objects,
    annotated with a list of tag names, and count of reviews
    """
    return book_qs.annotate(
        review_count=Count("reviews"), tag_names=ArrayAgg("tags__name")
    )
```

```{python}
# | echo: false
with perf_counter(time_sql=True, print_sql=True):
    qs = annotate_books_aggregate(get_queryset(public_id))
    result = serialize_books(qs)
print(result[0])
```

## Aggregation

Do not forget about duplicates when doing multiple aggregations!

Luckily, both `Count` and `ArrayAgg` accept a `distinct` argument

```{python}
#| code-line-numbers: "3,4"
def annotate_books_aggregate(book_qs: QuerySet[Book]):
    return book_qs.annotate(
        review_count=Count("reviews", distinct=True),
        tag_names=ArrayAgg("tags__name", distinct=True),
    )
```

```{python}
# | echo: false
with perf_counter(time_sql=True, print_sql=True):
    qs = get_queryset(public_id)
    qs = annotate_books_aggregate(qs)
    result = serialize_books(qs)
print(result[0])
```

Now we have a baseline for our endpoint, let's try to improve it.

## Query Plan {.scrollable}

```{python}
# |echo: false
qs = get_queryset(public_id)
qs = annotate_books_aggregate(qs)
result = qs.explain(analyze=True)

print(result)
```

## Query Plan

:::: {.columns}
::: {.column width=40%}

Right, this plan starts to be really complex. Several tools are available online to decrypt PostgreSQL plans, such as [this one](https://explain.dalibo.com/plan/g5h9cad9a38869gc)

<br>

We can now see, as expected, that the slow operations are the scan on `books_review` and `books_booktag`.

<br>

Even though the planner uses an index, it retrieves almost every row on both those tables, before merging them.

<br>

The aggregation is applied only at the end, just before sorting by `release_date`.

:::

::: {.column width=60%}
![](../images/planner.png)
:::
::::

## Ordering by id

:::: {.columns}
::: {.column width=40%}
By comparison, we can look at [the plan](https://explain.dalibo.com/plan/31ae7aa178g5f1da) when ordering by `id`. It's completely different.
1 key difference is the sort, which is now incremental
<br>

> Compared to regular sorts, sorting incrementally allows returning tuples before the entire result set has been sorted, which particularly enables optimizations with LIMIT queries. It may also reduce memory usage and the likelihood of spilling sorts to disk, but it comes at the cost of the increased overhead of splitting the result set into multiple sorting batches.

:::

::: {.column width=60%}
![](../images/agg_plan_by_id.png)

:::

::::

# Splitting filtering and aggregations

While we could try to find clever indexes to improve our queries, we'll try to find other solutions to reach the same result.

Since the aggregates are expensive, and only used for the output, but not to filter book rows, we can try to split the logic in two parts

```{python}
#| code-line-numbers: "|1-6|8-13|15-17"
def get_book_ids(library_id, page_size=20, orderings=("release_date", "id")):
    return list(
        Book.objects.filter(library_id=library_id)
        .order_by(*orderings)[:page_size]
        .values_list("id", flat=True)
    )

def get_queryset_from_ids(book_ids, orderings=("release_date", "id")):
    return (
        Book.objects.filter(id__in=book_ids)
        .select_related("author")
        .order_by(*orderings)
    )

def get_queryset_using_ids(library_id, page_size=20, orderings=("release_date", "id")):
    book_ids = get_book_ids(library_id, page_size=page_size, orderings=orderings)
    return get_queryset_from_ids(book_ids, orderings=orderings)
```

```{python}
# |echo: false
with perf_counter(time_sql=True, print_sql=True):
    qs = get_queryset_using_ids(alexandria_id)
    qs = annotate_books_aggregate(qs)

    result = serialize_books(qs)
```

## Using Subqueries

Another idea: Instead of using aggregation (with a `GROUP BY` clause), we can use subqueries.

```{python}
# |echo: false
from django.contrib.postgres.expressions import ArraySubquery
from django.db.models import Count, OuterRef, QuerySet
from django.db.models.functions import Coalesce

from books.models import Book, BookTag, Review
```

```{python}
#| code-line-numbers: "|4|4-7|3,8|10-11"
def annotate_books_subquery(book_qs: QuerySet[Book]):
    return book_qs.annotate(
        review_count=Coalesce(
            Review.objects.filter(book_id=OuterRef("id"))
            .values("book_id")
            .annotate(count=Count("id", distinct=True))
            .values("count"),
            0,
        ),
        tag_names=ArraySubquery(
            BookTag.objects.filter(book_id=OuterRef("id")).values("name")
        ),
    )
```

```{python}
# |echo: false
with perf_counter(time_sql=True, print_sql=True):
    qs = get_queryset_using_ids(alexandria_id)
    qs = annotate_books_subquery(qs)

    result = serialize_books(qs)
```

## Subquery plan

```{python}
# |echo: false
qs = get_queryset_using_ids(alexandria_id)
qs = annotate_books_subquery(qs)
result = qs.explain(analyze=True)

print(result)
```

1 thing to note: the `loops=20` on the `books_reviews` index. If the query becomes more complex, we'll have to execute this part once par item in the page

# Using Joins in Python

```{python}
#| code-line-numbers: "|4-9|10-15|16-19"
def annotate_books_python(books):
    book_ids = [book.id for book in books]

    review_count = dict(
        Review.objects.filter(book_id__in=book_ids)
        .values("book_id")
        .annotate(count=Count("id"))
        .values_list("book_id", "count")
    )
    tag_list = dict(
        BookTag.objects.filter(book_id__in=book_ids)
        .values("book_id")
        .annotate(tags=ArrayAgg("name"))
        .values_list("book_id", "tags")
    )
    for book in books:
        book.review_count = review_count[book.id]
        book.tag_names = tag_list.get(book.id, [])
    return books
```

---

Warning: I did not manage to produce an example where this is significantly better

```{python}
# | echo: false
with perf_counter(time_sql=True, print_sql=True):
    qs = get_queryset_using_ids(alexandria_id, page_size=20)
    qs = annotate_books_python(list(qs))

    result = serialize_books(qs)
```

---

## One last implementation

We can try to reduce the number of queries by using Common Table Expressions (CTE), i.e. queries having `WITH` clauses.

::::: {.columns}

::: {.column width=50% .fragment}

```{.sql code-line-numbers="1-11|12-22|23-29|"}
WITH review_count AS (
    SELECT
        books_review.book_id as id,
        count(books_review.id) as review_count
    FROM
        books_review
    WHERE
        books_review.book_id IN %(book_ids)s
        GROUP BY
            book_id
),
tag_names AS (
    SELECT
        books_booktag.book_id as id,
        array_agg(DISTINCT books_booktag.name) as tags_list
    FROM
        books_booktag
    WHERE
        books_booktag.book_id IN %(book_ids)s
        GROUP BY
            book_id
)
SELECT
    id,
    review_count,
    COALESCE(tags_list, '{}') as tags_list
FROM
    review_count
        FULL OUTER JOIN tag_names USING (id);
```

```{python}
# | echo: false
raw_query = """
WITH review_count AS (
    SELECT
        books_review.book_id as id,
        count(books_review.id) as review_count
    FROM
        books_review
    WHERE
        books_review.book_id IN %(book_ids)s
        GROUP BY
            book_id
),
tag_names AS (
    SELECT
        books_booktag.book_id as id,
        array_agg(DISTINCT books_booktag.name) as tags_list
    FROM
        books_booktag
    WHERE
        books_booktag.book_id IN %(book_ids)s
        GROUP BY
            book_id
)
SELECT
    id,
    review_count,
    COALESCE(tags_list, '{}') as tags_list
FROM
    review_count
        FULL OUTER JOIN tag_names USING (id);

"""
```

:::

:::: {.column width=50% }

::: {.fragment }

```{python}
#| code-line-numbers: "|2-13|14-18|"
def annotate_books_cte(books):
    book_ids = tuple(book.id for book in books)

    annotations_by_book_id = {
        raw_book.id: {
            "review_count": raw_book.review_count,
            "tag_names": raw_book.tags_list,
        }
        for raw_book in Book.objects.raw(
                raw_query,
                params={"book_ids": book_ids}
            )
    }

    for book in books:
        annotations = annotations_by_book_id[book.id]
        book.review_count = annotations["review_count"]
        book.tag_names = annotations["tag_names"]

    return books
```

:::

::: {.fragment }

We could even merge this query with the one retrieving all the `Book` and `Person` fields.

This would require to write more raw SQL code, though, which is not always a good idea.

We could try to resolve this problem by using the package [`django-cte` package](https://github.com/dimagi/django-cte)

:::

::::

:::::

## Summary

```{python}
# | echo: false
from rich.table import Table
import time
from utils.perf_display import format_duration

table = Table(show_lines=True)
table.add_column("name", style="bold green")
table.add_column("duration")


def add_row(annotate_func):
    start = time.perf_counter()
    qs = get_queryset_using_ids(7, page_size=20)
    qs = annotate_func(qs)
    serialize_books(qs)
    duration = time.perf_counter() - start
    table.add_row(annotate_func.__name__, format_duration(duration))


for annotate_func in [
    annotate_books_aggregate, annotate_books_subquery, annotate_books_python, annotate_books_cte
]:
    add_row(annotate_func)

console.print(table)
```