Skip to content

DuckDB

Installation

pip install "sqlframe[duckdb]"

Enabling SQLFrame

SQLFrame can be used in two ways:

  • Directly importing the sqlframe.duckdb package
  • Using the activate function to allow for continuing to use pyspark.sql but have it use SQLFrame behind the scenes.

Import

If converting a PySpark pipeline, all pyspark.sql should be replaced with sqlframe.duckdb. In addition, many classes will have a DuckDB prefix. For example, DuckDBDataFrame instead of DataFrame.

# PySpark import
# from pyspark.sql import SparkSession
# from pyspark.sql import functions as F
# from pyspark.sql.dataframe import DataFrame
# SQLFrame import
from sqlframe.duckdb import DuckDBSession
from sqlframe.duckdb import functions as F
from sqlframe.duckdb import DuckDBDataFrame

Activate

If you would like to continue using pyspark.sql but have it use SQLFrame behind the scenes, you can use the activate function.

from sqlframe import activate
activate("duckdb")

from pyspark.sql import SparkSession

SparkSession will now be a SQLFrame DuckDBSession object and everything will be run on DuckDB directly.

See activate configuration for information on how to pass in a connection and config options.

Creating a Session

SQLFrame uses the duckdb package to connect to DuckDB. A DuckDBSession, which implements the PySpark Session API, can be created by passing in a duckdb.Connection object or by allowing SQLFrame to create a connection for you. By default, SQLFrame will create a connection to an in-memory database.

from sqlframe.duckdb import DuckDBSession

session = DuckDBSession()

import duckdb
from sqlframe.duckdb import DuckDBSession

conn = duckdb.connect(database=":memory:")
session = DuckDBSession(conn=conn)

from sqlframe import activate
activate("duckdb")

from pyspark.sql import SparkSession

session = SparkSession.builder.getOrCreate()

import duckdb
from sqlframe import activate
conn = duckdb.connect(database=":memory:")
activate("duckdb", conn=conn)

from pyspark.sql import SparkSession

session = SparkSession.builder.getOrCreate()

Using DuckDB Unique Functions

DuckDB may have a function that isn't represented within the PySpark API. If that is the case, you can call it directly using PySpark call_function function.

from sqlframe.duckdb import DuckDBSession
from sqlframe.duckdb import functions as F

session = DuckDBSession()
(
    session.table("example.table")
    .select(F.call_function("CURRENT_SETTING", F.lit("access_mode")).alias("access_mode_value"))
    .show()
)

Example Usage

from sqlframe.duckdb import DuckDBSession
from sqlframe.duckdb import functions as F

session = DuckDBSession()

df_employee = session.createDataFrame(
    [
        {"id": 1, "fname": "Jack", "lname": "Shephard", "age": 37, "store_id": 1},
        {"id": 2, "fname": "John", "lname": "Locke", "age": 65, "store_id": 2},
        {"id": 3, "fname": "Kate", "lname": "Austen", "age": 37, "store_id": 3},
        {"id": 4, "fname": "Claire", "lname": "Littleton", "age": 27, "store_id": 1},
        {"id": 5, "fname": "Hugo", "lname": "Reyes", "age": 29, "store_id": 3},
    ]
)
df_store = session.createDataFrame(
    [
        {"store_id": 1, "store_name": "The Hatch"},
        {"store_id": 2, "store_name": "The Pearl"},
        {"store_id": 3, "store_name": "The Swan"},
    ]
)

(
    df_employee
    .join(df_store, on="store_id")
    .groupBy("store_name")
    .agg(F.count("*").alias("total_employees"))
    .show()
)

Supported PySpark API Methods

See something that you would like to see supported? Open an issue!

Catalog Class

Column Class

DataFrame Class

Functions

GroupedData Class

DataFrameReader Class

DataFrameWriter Class

SparkSession Class

DataTypes

Window Class

WindowSpec Class

Extra Functionality not Present in PySpark

SQLFrame supports the following extra functionality not in PySpark

Table Class

SQLFrame provides a Table class that supports extra DML operations like update and delete. This class is returned when using the table function from the DataFrameReader class.

import duckdb
from sqlframe.duckdb import DuckDBSession

conn = duckdb.connect(database=":memory:")
session = DuckDBSession(conn=conn)

df_employee = session.createDataFrame(
    [
        {"id": 1, "fname": "Jack", "lname": "Shephard", "age": 37, "store_id": 1},
        {"id": 2, "fname": "John", "lname": "Locke", "age": 65, "store_id": 2},
        {"id": 3, "fname": "Kate", "lname": "Austen", "age": 37, "store_id": 3},
        {"id": 4, "fname": "Claire", "lname": "Littleton", "age": 27, "store_id": 1},
        {"id": 5, "fname": "Hugo", "lname": "Reyes", "age": 29, "store_id": 3},
    ]
)

df_employee.write.mode("overwrite").saveAsTable("employee")

table_employee = session.table("employee")  # This object is of Type DuckDBTable

Update Statement

The update method of the Table class is equivalent to the UPDATE table_name statement used in standard sql.

# Generates a `LazyExpression` object which can be executed using the `execute` method
update_expr = table_employee.update(
    set_={"age": table_employee["age"] + 1},
    where=table_employee["id"] == 1,
)

# Executes the update statement
update_expr.execute()

# Show the result
table_employee.show()

Output: 

+----+--------+-----------+-----+----------+
| id | fname  |   lname   | age | store_id | 
+----+--------+-----------+-----+----------+
| 1  |  Jack  |  Shephard |  38 |    1     |
| 2  |  John  |   Locke   |  65 |    2     |
| 3  |  Kate  |   Austen  |  37 |    3     |
| 4  | Claire | Littleton |  27 |    1     |
| 5  |  Hugo  |   Reyes   |  29 |    3     |
+----+--------+-----------+-----+----------+

Delete Statement

The delete method of the Table class is equivalent to the DELETE FROM table_name statement used in standard sql.

# Generates a `LazyExpression` object which can be executed using the `execute` method
delete_expr = table_employee.delete(
    where=table_employee["id"] == 1,
)

# Executes the delete statement
delete_expr.execute()

# Show the result
table_employee.show()

Output: 

+----+--------+-----------+-----+----------+
| id | fname  |   lname   | age | store_id | 
+----+--------+-----------+-----+----------+
| 2  |  John  |   Locke   |  65 |    2     |
| 3  |  Kate  |   Austen  |  37 |    3     |
| 4  | Claire | Littleton |  27 |    1     |
| 5  |  Hugo  |   Reyes   |  29 |    3     |
+----+--------+-----------+-----+----------+