2023-07-07

Vom Watscheln zum Fliegen: DuckDBs Funktionalität schnell mit skalaren Python-UDFs erweitern

Pedro Holanda, Thijs Bruineman, Phillip Cloud

User Defined Functions (UDFs) erlauben es, die Funktionalität eines Datenbankmanagementsystems (DBMS) um fachspezifische Aufgaben zu erweitern, die nicht als eingebaute Funktionen vorliegen. Wer häufig private Daten exportiert, profitiert zum Beispiel von einer Anonymisierungsfunktion, die den lokalen Teil einer E-Mail maskiert und die Domain erhält. Ideal wäre, diese Funktion direkt im DBMS auszuführen. Das hat mehrere Vorteile:

1. Leistung. Die Funktion ließe sich im selben Ausführungsmodell ausführen (z. B. Streaming-Ergebnisse, Beyond-Memory/Out-of-Core) wie das DBMS, ohne unnötige Transformationen.

2. Einfache Nutzung. UDFs lassen sich nahtlos in SQL-Abfragen einbinden; Nutzer können SQL nutzen, um die Funktionen aufzurufen. Daten müssen nicht durch einen getrennten Datenbank-Connector und externen Code. Die Funktionen stehen in verschiedenen SQL-Kontexten zur Verfügung (z. B. Subqueries, Join-Bedingungen).

3. Sicherheit. Die sensiblen Daten verlassen den DBMS-Prozess nie.

Es gibt zwei Hauptgründe, warum Nutzer oft keine UDFs implementieren. 1) Sicherheitsbedenken. UDFs sind vom Nutzer geschriebener Code, der im DBMS-Prozess läuft; es besteht das Risiko, den Server zum Absturz zu bringen. Bei DuckDB, einer eingebetteten Datenbank, ist das weniger kritisch: Jeder Analyst führt seinen eigenen DuckDB-Prozess. Die Stabilität eines zentralen Servers steht nicht auf dem Spiel. 2) Die Schwierigkeit der Implementierung schreckt ab. Hochperformante UDFs gibt es typischerweise nur in Low-Level-Sprachen. UDFs in Higher-Level-Sprachen wie Python kosten erheblich Leistung. Viele Nutzer können ihre UDFs daher nicht schnell umsetzen, ohne viel Zeit in eine Low-Level-Sprache und die internen Details des DBMS zu investieren.

DuckDB ist einen ähnlichen Weg gegangen. Als DBMS für analytische Aufgaben steht Leistung im Vordergrund, der Kern ist daher in C++ implementiert. Entsprechend lag der erste Fokus der Erweiterbarkeit auf C++. Diese Ente kann aber nicht nur watscheln, sie kann auch fliegen. Deshalb freuen wir uns, die kürzliche Ergänzung skalarer Python-UDFs in DuckDB bekannt zu geben.

DuckDB unterstützt zwei Arten von Python-UDFs, die sich im Python-Objekt unterscheiden, das zwischen DuckDBs nativen Datentypen und dem Python-Prozess vermittelt. Diese Kommunikationsschichten umfassen eingebaute Python-Typen und PyArrow Tables.

Die beiden Ansätze unterscheiden sich in zwei Punkten:

1. Zero-Copy. PyArrow Tables nutzen unsere Zero-Copy-Integration mit Arrow und übersetzen Datentypen ohne Copy-Kosten nach Python.

2. Vektorisierung. PyArrow-Table-Funktionen arbeiten auf Chunk-Ebene und verarbeiten Chunks mit bis zu 2048 Zeilen. Das maximiert Cache-Lokalität und nutzt Vektorisierung. Die UDF-Implementierung mit eingebauten Typen arbeitet dagegen zeilenweise.

Dieser Beitrag zeigt, wie Sie DuckDB mit Python-UDFs erweitern, mit besonderem Augenmerk auf PyArrow-UDFs. Im Quick-Tour-Abschnitt nutzen wir PyArrow-UDF-Typen. Wer Benchmarks sucht, kann zum Benchmark-Abschnitt springen. Eine detaillierte Beschreibung der Python-UDF-API steht in unserer Dokumentation.

Python-UDFs

Dieser Abschnitt zeigt mehrere praktische Beispiele für Python-UDFs. Jedes Beispiel nutzt einen anderen UDF-Typ.

Quick-Tour

Zur Demonstration betrachten wir folgendes Beispiel. Wir haben ein Dictionary world_cup_titles, das Ländern die Zahl gewonnener Weltmeisterschaften zuordnet. Wir wollen eine Python-UDF, die einen Ländernamen entgegennimmt, den entsprechenden Wert im Dictionary sucht und die Zahl der Weltmeisterschaften zurückgibt. Ist das Land nicht im Dictionary, gibt die UDF NULL zurück.

Eine Beispielimplementierung:

import duckdb
from duckdb.typing import *

con = duckdb.connect()

# Dictionary that maps countries and world cups they won
world_cup_titles = {
    "Brazil": 5,
    "Germany": 4,
    "Italy": 4,
    "Argentina": 2,
    "Uruguay": 2,
    "France": 2,
    "England": 1,
    "Spain": 1
}

# Function that will be registered as an UDF, simply does a lookup in the python dictionary
def world_cups(x):
     return world_cup_titles.get(x)

# We register the function
con.create_function("wc_titles", world_cups, [VARCHAR], INTEGER)

Das war’s: Die Funktion ist registriert und kann über SQL aufgerufen werden.

# Let's create an example countries table with the countries we are interested in using
con.execute("CREATE TABLE countries (country VARCHAR)")
con.execute("INSERT INTO countries VALUES ('Brazil'), ('Germany'), ('Italy'), ('Argentina'), ('Uruguay'), ('France'), ('England'), ('Spain'), ('Netherlands')")
# We can simply call the function through SQL, and even use the function return to eliminate the countries that never won a world cup
con.sql("SELECT country, wc_titles(country) AS world_cups FROM countries").fetchall()
# [('Brazil', 5), ('Germany', 4), ('Italy', 4), ('Argentina', 2), ('Uruguay', 2), ('France', 2), ('England', 1), ('Spain', 1), ('Netherlands', None)]

Fake-Daten mit Faker erzeugen (UDF mit eingebauten Typen)

Hier ein Beispiel, das die Faker-Bibliothek nutzt, um in DuckDB eine skalare Funktion zu erzeugen, die zufällige Daten zurückgibt. Die Funktion random_date braucht keine Eingaben und liefert eine Spalte DATE. Da Faker eingebaute Python-Typen nutzt, gibt die Funktion diese direkt zurück. Wichtig: Eine Funktion, die bezogen auf ihre Eingabe nicht deterministisch ist, muss als mit side_effects markiert werden.

import duckdb

# By importing duckdb.typing we can specify DuckDB Types directly without using strings
from duckdb.typing import *

from faker import Faker

# Our Python UDF generates a random date every time it's called
def random_date():
     fake = Faker()
     return fake.date_between()

Die Python-Funktion registrieren wir in DuckDB mit create_function. Da unsere Funktion keine Eingaben braucht, übergeben wir eine leere Liste als argument_type_list. Als Rückgabetyp geben wir DATE aus duckdb.typing an. Weil random_date() einen eingebauten Python-Typ (datetime.date) zurückgibt, müssen wir den UDF-Typ nicht angeben.

# To exemplify the effect of side-effect, let's first run the function without marking it.
duckdb.create_function('random_date', random_date, [], DATE)

# After registration, we can use the function directly via SQL
# Notice that without side_effect=True, it's not guaranteed that the function will be re-evaluated.
res = duckdb.sql('SELECT random_date() FROM range (3)').fetchall()
# [(datetime.date(2003, 8, 3),), (datetime.date(2003, 8, 3),), (datetime.date(2003, 8, 3),)]

# Now let's re-add the function with side-effects marked as true.
duckdb.remove_function('random_date')
duckdb.create_function('random_date', random_date, [], DATE, side_effects=True)
res = duckdb.sql('SELECT random_date() FROM range (3)').fetchall()
# [(datetime.date(2020, 11, 29),), (datetime.date(2009, 5, 18),), (datetime.date(2018, 5, 24),)]

Groß-/Kleinschreibung von Strings tauschen (PyArrow-UDF)

Ein Problem eingebauter Typen: Sie profitieren nicht von Zero-Copy, Vektorisierung und Cache-Lokalität. PyArrow als UDF-Typ sollte bevorzugt werden, um diese Optimierungen zu nutzen.

Als PyArrow-Funktion ein einfaches Beispiel: Kleinbuchstaben nach groß, Großbuchstaben nach klein. Glücklicherweise hat PyArrow das schon in der Compute-Engine; es reicht pc.utf8_swapcase(x).

import duckdb

# By importing duckdb.typing we can specify DuckDB Types directly without using strings
from duckdb.typing import *

import pyarrow as pa
import pyarrow.compute as pc

def swap_case(x):
     # Swap the case of the 'column' using utf8_swapcase and return the result
     return pc.utf8_swapcase(x)

con = duckdb.connect()
# To register the function, we must define it's type to be 'arrow'
con.create_function('swap_case', swap_case, [VARCHAR], VARCHAR, type='arrow')

res = con.sql("SELECT swap_case('PEDRO HOLANDA')").fetchall()
# [('pedro holanda',)]

Taxitarife vorhersagen (Ibis + PyArrow-UDF)

Python-UDFs sind mächtig, weil sie das große Python-Ökosystem und Werkzeuge erschließen, einschließlich Bibliotheken wie PyTorch und Tensorflow, die Machine-Learning-Operationen effizient umsetzen.

Außerdem bietet das Ibis-Projekt eine DataFrame-API mit guter DuckDB-Integration und unterstützt sowohl DuckDBs native Python- als auch PyArrow-UDFs.

In diesem Beispiel nutzen wir ein vortrainiertes PyTorch-Modell, um Taxitarife anhand der zurückgelegten Strecke zu schätzen. Ein vollständiges Beispiel steht in diesem Blogbeitrag des Ibis-Teams.

import torch
import pyarrow as pa
import ibis
import ibis.expr.datatypes as dt

from ibis.expr.operations import udf


# The code to generate the model is not specified in this snippet, please refer to the provided link for more information
model = ...

# Function that uses the model and a traveled distance input tensor to predict values, please refer to the provided link for more information
def predict_linear_regression(model, tensor: torch.Tensor) -> torch.Tensor:
    ...


# Indicate to ibis that this is a scalar user-defined function whose input format is pyarrow
@udf.scalar.pyarrow
def predict_fare(x: dt.float64) -> dt.float32:
    # `x` is a pyarrow.ChunkedArray; the `dt.float64` annotation indicate the element type of the ChunkedArray.

    # Transform the data from PyArrow to the required torch tensor format and dimension.
    tensor = torch.from_numpy(x.to_numpy()[:, None]).float()

    # Call the actual prediction function, which also returns a torch tensor.
    predicted = predict_linear_regression(model, tensor).ravel()
    return pa.array(predicted.numpy())


# Execute a query on the NYC Taxi Parquet file to showcase our model's predictions, the actual fare amount, and the distance.
expr = (
    ibis.read_parquet('yellow_tripdata_2016-02.parquet')
    .mutate(
        "fare_amount",
        "trip_distance",
        predicted_fare=lambda t: predict_fare(t.trip_distance),
    )
)
df = expr.execute()

Mit Python-UDFs in DuckDB und Ibis lassen sich Machine-Learning-Modelle nahtlos einbinden und Vorhersagen direkt in Ibis-Code und SQL-Abfragen ausführen. Das Beispiel zeigt, wie Taxitarife anhand der Distanz mit einem PyTorch-Modell vorhergesagt werden – Machine Learning in DuckDBs SQL-Umgebung, angesteuert von Ibis.

Benchmarks

In diesem Abschnitt vergleichen wir die beiden Arten von Python-UDFs in einfachen Benchmarks. Gemessen werden Ausführungszeit und Peak-Speicherverbrauch. Die Benchmarks laufen 5 Mal, der Median zählt. Die Maschine ist ein Mac Apple M1 mit 16 GB RAM.

Eingebautes Python vs. PyArrow

Für diese UDF-Typen erzeugen wir UDFs, die eine ganzzahlige Spalte entgegennehmen, zu jedem Wert eins addieren und das Ergebnis zurückgeben. Der Code für diesen Abschnitt liegt in einem GitHub Gist.

import pyarrow.compute as pc
import duckdb
import pyarrow as pa

# Built-In UDF
def add_built_in_type(x):
    return x + 1

# Arrow UDF
def add_arrow_type(x):
    return pc.add(x, 1)

con = duckdb.connect()

# Registration
con.create_function('built_in_types', add_built_in_type, ['BIGINT'], 'BIGINT', type='native')
con.create_function('add_arrow_type', add_arrow_type, ['BIGINT'], 'BIGINT', type='arrow')

# Integer View with 10,000,000 elements.
con.sql("""
     SELECT i
     FROM range(10000000) tbl(i);
""").to_view("numbers")

# Calls for both UDFs
native_res = con.sql("SELECT sum(add_built_in_type(i)) FROM numbers").fetchall()
arrow_res = con.sql("SELECT sum(add_arrow_type(i)) FROM numbers").fetchall()

Zwischen den beiden UDFs liegt mehr als eine Größenordnung. Die Differenz hat vor allem drei Gründe:

1. In Python sind Objekterzeugung und allgemeine Nutzung eher langsam. Gründe sind unter anderem automatische Speicherverwaltung, Interpretation und dynamische Typisierung.
2. Die PyArrow-UDF braucht keine Datenkopie.
3. Die PyArrow-UDF läuft vektorisiert und verarbeitet Chunks statt einzelner Zeilen.

Python-UDFs vs. externe Funktionen

Hier vergleichen wir eine Python-UDF mit einer externen Funktion. Die Funktion berechnet die Summe der Längen aller Strings einer Spalte. Der Code liegt in einem GitHub Gist.

import duckdb
import pyarrow as pa

# Function used in UDF
def string_length_arrow(x):
     tuples = len(x)
     values = [len(i.as_py()) if i.as_py() != None else 0 for i in x]
     array = pa.array(values, type=pa.int32(), size=tuples)
     return array


# Same Function but external to the database
def exec_external(con):
     arrow_table = con.sql("SELECT i FROM strings tbl(i)").arrow()
     arrow_column = arrow_table['i']
     tuples = len(arrow_column)
     values = [len(i.as_py()) if i.as_py() != None else 0 for i in arrow_column]
     array = pa.array(values, type=pa.int32(), size=tuples)
     arrow_tbl = pa.Table.from_arrays([array], names=['i'])
     return con.sql("SELECT sum(i) FROM arrow_tbl").fetchall()


con = duckdb.connect()
con.create_function('strlen_arrow', string_length_arrow, ['VARCHAR'], int, type='arrow')

con.sql("""
     SELECT
          CASE WHEN i != 0 AND i % 42 = 0
          THEN
               NULL
          ELSE
               repeat(chr((65 + (i % 26))::INTEGER), (4 + (i % 12))) END
          FROM range(10000000) tbl(i);
""").to_view("strings")

con.sql("SELECT sum(strlen_arrow(i)) FROM strings tbl(i)").fetchall()

exec_external(con)

Bei der Leistung gibt es keine wesentliche Regression durch UDFs. Sie behalten aber die Vorteile sichererer Ausführung und der Nutzung von SQL. In unserem Beispiel materialisiert die externe Funktion die gesamte Abfrage; der Peak-Speicherverbrauch ist 5× höher als beim UDF-Ansatz.

Fazit und weitere Entwicklung

Skalare Python-UDFs werden jetzt in DuckDB unterstützt – ein wichtiger Schritt, um die Funktionalität der Datenbank zu erweitern. Nutzer können komplexe Berechnungen in einer High-Level-Sprache ausführen. Außerdem können Python-UDFs DuckDBs Zero-Copy-Integration mit Arrow nutzen, Datentransferkosten vermeiden und Abfragen effizient halten.

Die Einführung von Python-UDFs ist ein großer Schritt, die Arbeit geht aber weiter. Unsere Roadmap umfasst:

1. Aggregat-/tabellenerzeugende UDFs: Derzeit können Nutzer skalare UDFs anlegen; wir arbeiten an Aggregatfunktionen (Berechnungen über eine Menge von Werten, ein Ergebnis) und tabellenerzeugenden Funktionen (Tabellen ohne Beschränkung der Spalten- und Zeilenzahl).

2. Typen: Skalare Python-UDFs unterstützen derzeit die meisten DuckDB-Typen, ausgenommen ENUM und BIT. Wir erweitern die Typunterstützung, damit die Funktionalität vollständig wird.

Wenn Sie Probleme mit unseren Python-UDFs haben, öffnen Sie bitte ein Issue im DuckDB Issue Tracker.