2024-10-25

Schnelle Top-N-Aggregation und -Filterung mit DuckDB

Alex Monahan

Einführung in Top N

Ein häufiges Muster bei der Datenanalyse ist, nach den Zeilen zu suchen, die in einer bestimmten Metrik am höchsten oder niedrigsten liegen. Wenn die höchsten oder niedrigsten N Zeilen eines gesamten Datensatzes interessieren, sortieren die Standardklauseln ORDER BY und LIMIT von SQL nach der gewünschten Metrik und liefern nur N Zeilen. Zum Beispiel mit dem Scale-Factor-1-Datensatz (SF1) des TPC-H-Benchmarks:

INSTALL tpch;
LOAD tpch;
-- Generate an example TPC-H dataset
CALL dbgen(sf = 1);

-- Return the most recent 3 rows by l_shipdate
FROM lineitem
ORDER BY
    l_shipdate DESC
LIMIT 3;

Das ist nützlich, um schnell die ältesten oder neuesten Werte in einem Datensatz zu holen oder Ausreißer in einer bestimmten Metrik zu finden.

Ein anderer gängiger Ansatz ist, die Min-/Max-Zusammenfassungsstatistiken einer oder mehrerer Spalten abzufragen. Das findet Ausreißer, aber die Zeile, die den Ausreißer enthält, kann für jede Spalte eine andere sein – die Frage ist also eine andere. DuckDBs hilfreicher COLUMNS-Ausdruck erlaubt uns, den Maximalwert für alle Spalten zu berechnen.

FROM lineitem
SELECT
    max(COLUMNS(*));

Die Queries in diesem Beitrag nutzen ausgiebig DuckDBs FROM-first-Syntax. Damit lassen sich die Klauseln FROM und SELECT tauschen, und Letztere kann sogar ganz weggelassen werden.

Diese beiden Ansätze können aber nur bestimmte Arten von Fragen beantworten. Es gibt viele Szenarien, in denen das Ziel ist, die Top-N-Werte innerhalb einer Gruppe zu verstehen. Im ersten Beispiel oben: Wie berechnen wir die letzten 10 Lieferungen jedes Lieferanten? Die LIMIT-Klausel von SQL kann das nicht abbilden. Nennen wir diese Art der Analyse Top N nach Gruppe.

Diese Art der Analyse ist ein gängiges Werkzeug, um neue Datensätze zu erkunden. Anwendungsfälle sind etwa, die neuesten paar Zeilen für jede Gruppe zu holen oder die extremsten paar Werte in einer Gruppe zu finden. Bleiben wir beim Lieferbeispiel: Wir könnten uns die letzten 10 Lieferungen jeder Teilenummer ansehen oder die 5 teuersten Bestellungen pro Kunde finden.

Klassisches Top N nach Gruppe

In den meisten Datenbanken filtert man auf das Top N innerhalb einer Gruppe mit einer Window-Funktion und einem Common Table Expression (CTE). Dieser Ansatz funktioniert auch in DuckDB. Diese Query liefert zum Beispiel die 3 neuesten Lieferungen für jeden Lieferanten:

WITH ranked_lineitem AS (
    FROM lineitem
    SELECT
        *,
        row_number() OVER
            (PARTITION BY l_suppkey ORDER BY l_shipdate DESC)
            AS my_ranking
)
FROM ranked_lineitem
WHERE
    my_ranking <= 3;

In DuckDB lässt sich das mit der QUALIFY-Klausel vereinfachen. QUALIFY wirkt wie eine WHERE-Klausel, operiert aber speziell auf den Ergebnissen von Window-Funktionen. Mit dieser Anpassung kann der CTE entfallen, die Ergebnisse bleiben dieselben.

FROM lineitem
SELECT
    *,
    row_number() OVER
        (PARTITION BY l_suppkey ORDER BY l_shipdate DESC)
        AS my_ranking
QUALIFY
    my_ranking <= 3;

Das ist durchaus ein gangbarer Weg! Wo liegen aber die Schwächen? Obwohl die Query nur an den 3 neuesten Lieferungen interessiert ist, muss sie jede Lieferung sortieren, nur um diese Top 3 zu holen. Sortieren in DuckDB hat eine Komplexität von O(kn) dank DuckDBs innovativer Radix-Sort-Implementierung, das ist aber immer noch höher als das O(n) von DuckDBs Hash-Aggregat, zum Beispiel. Sortieren ist außerdem speicherintensiv im Vergleich zur Aggregation.

Top N in DuckDB

DuckDB 1.1 hat eine neue Fähigkeit ergänzt, die Top-N-Berechnungen deutlich vereinfacht und beschleunigt. Nämlich akzeptieren die Funktionen min, max, min_by und max_by jetzt alle einen optionalen Parameter N. Ist N größer als 1 (der Default), liefern sie ein Array der Top-Werte.

Als einfaches Beispiel fragen wir die neuesten (Top-3-)Lieferdaten ab:

FROM lineitem
SELECT
    max(l_shipdate, 3) AS top_3_shipdates;

Top N nach Spalte in DuckDB

Die Top-N-Auswahl wird noch nützlicher dank erneut des COLUMNS-Ausdrucks – wir können die 3 Top-Werte in jeder Spalte holen. Nennen wir das eine Top-N-nach-Spalte-Analyse. Mit gewöhnlichem SQL ist das besonders unübersichtlich! Man bräuchte eine Subquery oder Window-Funktion für jede einzelne Spalte… In DuckDB einfach:

FROM lineitem
SELECT
    max(COLUMNS(*), 3) AS "top_3_\0";

Top N nach Gruppe in DuckDB

Ausgestattet mit dem neuen Parameter N: Wie können wir eine Top-N-nach-Gruppe-Analyse beschleunigen?

Wollen Sie gleich zum Endergebnis? Springen Sie gerne voraus!

Wir nutzen drei weitere DuckDB-SQL-Features, um das möglich zu machen:

• Die Funktion max_by (auch bekannt als arg_max)
• Die Funktion unnest
• Automatisches Packen einer ganzen Zeile in eine STRUCT-Spalte

Die Funktion max liefert das Maximum (oder jetzt das Max-N!) einer bestimmten Spalte. Im Gegensatz dazu findet max_by den Maximalwert in einer Spalte und holt dann einen Wert aus derselben Zeile, aber einer anderen Spalte. Diese Query liefert zum Beispiel die IDs der 3 zuletzt gelieferten Bestellungen für jeden Lieferanten:

FROM lineitem
SELECT
    l_suppkey,
    max_by(l_orderkey, l_shipdate, 3) AS recent_orders
GROUP BY
    l_suppkey;

max_by ist eine Aggregatfunktion und nutzt damit DuckDBs schnelle Hash-Aggregation statt Sortierung. Statt nach l_shipdate zu sortieren, scannt max_by den Datensatz nur einmal und merkt sich die N höchsten l_shipdate-Werte. Dann liefert sie die Bestell-ID, die zu jedem der neuesten Lieferdaten gehört. Der Radix-Sort in DuckDB muss den Datensatz einmal pro Byte scannen, ein einmaliger Scan bringt also einen deutlichen Speedup. Sortiert man zum Beispiel nach einer 64-Bit-Ganzzahl, muss der Sortieralgorithmus den Datensatz 8-mal durchlaufen, gegenüber 1-mal mit diesem Ansatz! Ein einfacher Mikro-Benchmark steht im Abschnitt Performance-Vergleiche.

Diese SQL-Query hat aber ein paar Lücken. Die Query liefert Ergebnisse als LIST statt als getrennte Zeilen. Glücklicherweise kann unnest eine LIST in getrennte Zeilen aufteilen:

FROM lineitem
SELECT
    l_suppkey,
    unnest(
        max_by(l_orderkey, l_shipdate, 3)
    ) AS recent_orders
GROUP BY
    l_suppkey;

Die nächste Lücke: Es gibt keine einfache Möglichkeit, das zu den zurückgegebenen l_orderkey-Werten gehörige l_shipdate zu sehen. Diese Query liefert nur eine einzelne Spalte, während eine Top-N-nach-Gruppe-Analyse typischerweise die ganze Zeile braucht.

Glücklicherweise erlaubt DuckDB uns, den gesamten Inhalt einer Zeile so zu referenzieren, als wäre sie nur eine einzelne Spalte! Indem wir den Namen der Tabelle selbst (hier lineitem) statt des Namens einer Spalte angeben, kann max_by alle Spalten holen.

FROM lineitem
SELECT
    l_suppkey,
    unnest(
        max_by(lineitem, l_shipdate, 3)
    ) AS recent_orders
GROUP BY
    l_suppkey;

Machen wir das etwas freundlicher, indem wir den STRUCT in getrennte Spalten aufteilen, damit er zu unserem ursprünglichen Datensatz passt.

Die finale Top-N-nach-Gruppe-Query

Ein weiteres Argument an UNNEST teilt das rekursiv in getrennte Spalten auf. In diesem Fall heißt das: UNNEST läuft zweimal: einmal, um jede LIST in getrennte Zeilen zu wandeln, und dann noch einmal, um jeden STRUCT in getrennte Spalten zu wandeln. Die Spalte l_suppkey kann auch weggelassen werden, da sie automatisch schon enthalten ist.

FROM lineitem
SELECT
    unnest(
        max_by(lineitem, l_shipdate, 3),
        recursive := 1
    ) AS recent_orders
GROUP BY
    l_suppkey;

Dieser Ansatz ist auch nützlich für die häufige Aufgabe, zu deduplizieren, indem man den neuesten Wert innerhalb einer Gruppe findet. Ein Muster ist, den aktuellen Zustand eines Datensatzes zu finden, indem man das neueste Event in einer Event-Tabelle zurückgibt. Einfach ein N von 1 verwenden!

Wir haben jetzt eine Möglichkeit, eine Aggregatfunktion zu nutzen, um die Top-N-Zeilen pro Gruppe zu berechnen! Wie viel effizienter ist das also?

Performance-Vergleiche

Wir vergleichen den QUALIFY-Ansatz mit dem max_by-Ansatz für das Top-N-nach-Gruppe-Problem. Wir haben beide Queries besprochen, wiederholen sie aber zur Referenz unten.

FROM lineitem
SELECT
    *,
    row_number() OVER
        (PARTITION BY l_suppkey ORDER BY l_shipdate DESC)
        AS my_ranking
QUALIFY
    my_ranking <= 3;

FROM lineitem
SELECT
    unnest(
        max_by(lineitem, l_shipdate, 3),
        recursive := 1
    )
GROUP BY
    l_suppkey;

Während die Hauptquery läuft, starten wir außerdem einen Hintergrund-Thread, der periodisch DuckDBs Speicherverbrauch misst. Das nutzt die eingebaute Tabellenfunktion duckdb_memory() und enthält Informationen über Speicherverbrauch sowie temporären Festplattenverbrauch. Das kleine Python-Skript für das Benchmarking steht unter den Ergebnissen. Die Maschine für das Benchmarking war ein M1 MacBook Pro mit 16 GB RAM.

Wir sehen, dass der max_by-Ansatz in jeder dieser Situationen schneller ist, in manchen Fällen fast 5× schneller! Wird der Datensatz aber größer, schwächt sich der max_by-Ansatz relativ zu QUALIFY ab.

In manchen Fällen ist der Speicherverbrauch mit max_by ebenfalls deutlich niedriger. Der Speicherverbrauch des max_by-Ansatzes wird aber mit steigendem Scale signifikanter, weil die Zahl der unterschiedlichen l_suppkey-Werte linear mit dem Scale Factor wächst. Dieser erhöhte Speicherverbrauch erklärt wahrscheinlich den Performance-Rückgang, da beide Algorithmen an die maximale RAM-Menge auf meiner Maschine herankamen und auf Platte auszulagern begannen.

Um den Speicherdruck zu senken, wiederholen wir den Scale-Factor-10-Benchmark (SF10) mit weniger Threads (4 Threads und 1 Thread). Wir nutzen weiter eine max_memory-Einstellung von 36 GB. Die vorherigen SF10-Ergebnisse mit allen 10 Threads sind zur Referenz enthalten.

Der max_by-Ansatz ist so recheneffizient, dass er selbst mit 1 Thread deutlich schneller ist als der QUALIFY-Ansatz, der alle 10 Threads nutzt! Die Thread-Zahl zu senken, hat den Speicherverbrauch ebenfalls sehr wirksam gesenkt (fast 10× Reduktion).

Wann sollten wir also welchen Ansatz nutzen? Wie bei allen Datenbankdingen: es kommt darauf an! Ist der Speicher knapp, kann max_by ebenfalls Vorteile bieten, besonders wenn die Thread-Zahl so abgestimmt ist, dass kein Spill auf Platte entsteht. Gibt es aber ungefähr so viele Gruppen wie Zeilen, kommt QUALIFY in Betracht, weil wir etwas von der Speichereffizienz des max_by-Ansatzes verlieren.

import duckdb
import pandas as pd
from threading import Thread
from time import sleep
from datetime import datetime
from os import remove

def check_memory(stop_function, filepath, sleep_seconds, results_dict):
    print("Starting background thread")
    background_con = duckdb.connect(filepath)
    max_memory = 0
    max_temporary_storage = 0
    while True:
        if stop_function():
            break
        # Profile the memory
        memory_profile = background_con.sql("""
            FROM duckdb_memory()
            SELECT
                tag,
                round(memory_usage_bytes / (1000000), 0)::bigint AS memory_usage_mb,
                round(temporary_storage_bytes / (1000000), 0)::bigint AS temporary_storage_mb;
            """).df()
        print(memory_profile)
        total_memory = background_con.sql("""
            FROM memory_profile
            select
                sum(memory_usage_mb) AS total_memory_usage_mb,
                sum(temporary_storage_mb) AS total_temporary_storage_mb
            """).fetchall()
        print('Current memory:', total_memory[0][0])
        print('Current temporary_storage:', total_memory[0][1])

        if total_memory[0][0] > max_memory:
            max_memory = total_memory[0][0]
        if total_memory[0][1] > max_temporary_storage:
            max_temporary_storage = total_memory[0][1]

        print('Maximum memory:', max_memory)
        print('Maximum temporary_storage:', max_temporary_storage)

        sleep(sleep_seconds)

    results_dict["max_memory"] = max_memory
    results_dict["max_temporary_storage"] = max_temporary_storage
    background_con.close()

    return

def query_and_profile(filepath, sql):
    con = duckdb.connect(filepath)
    con.sql("set max_memory='36GB'")

    results_dict = {}
    stop_threads = False
    background_memory_thread = Thread(target=check_memory,
                                      args=(lambda : stop_threads, filepath, 0.1, results_dict, ))
    background_memory_thread.start()

    print("Starting query:")
    start_time = datetime.now()
    results_df = con.sql(sql).df()
    results_dict["total_time_seconds"] = (datetime.now() - start_time).total_seconds()
    print(results_df.head(10))

    stop_threads = True
    background_memory_thread.join()
    con.close()

    return results_dict

filepath = './arg_max_check_duckdb_memory_v3.duckdb'

con = duckdb.connect(filepath)
print("Begin initial tpch load")
con.sql("""call dbgen(sf=1);""")
con.close()

sql = """
    FROM lineitem
    SELECT
        UNNEST(
            max_by(lineitem, l_shipdate, 3),
            recursive := 1
        )
    GROUP BY
        l_suppkey
;"""

max_by_results = query_and_profile(filepath, sql)

sql = """
    FROM lineitem
    SELECT
        *,
        row_number() OVER
            (PARTITION BY l_suppkey ORDER BY l_shipdate DESC)
            AS my_ranking
    QUALIFY
        my_ranking <= 3
;"""

qualify_results = query_and_profile(filepath, sql)

print('max_by_results:', max_by_results)
print('qualify_results:', qualify_results)

remove(filepath)

Fazit

DuckDB bietet jetzt eine bequeme Möglichkeit, die Top-N-Werte sowohl der Aggregatfunktionen min und max zu berechnen als auch ihrer fortgeschrittenen Verwandten min_by und max_by. Sie sind leicht zu nutzen und ermöglichen außerdem komplexere Analysen wie die Berechnung des Top N für alle Spalten oder des Top N nach Gruppe. Im Vergleich zu einem Window-Funktions-Ansatz sind auch Performance-Vorteile möglich.

Wir würden gerne hören, auf welche kreativen Weisen Sie dieses neue Feature nutzen!

Viel Spaß beim Analysieren!