Repository: coiled/benchmarks (language: python)

Issue:
Add tpch benchmark queries

Add pyspark tpch benchmarks

Add tpch benchmark queries



Discussion / hints:

Fix patch:
diff --git a/ci/environment.yml b/ci/environment.yml
index 5d09daebb1..770024a6c0 100644
--- a/ci/environment.yml
+++ b/ci/environment.yml
@@ -1,7 +1,12 @@
 channels:
   - conda-forge
 dependencies:
-  - python >=3.9
+  - conda
+  - python =3.11
+  - pyspark ==3.4.1
+  - openjdk ==20.0.2
+  - python-duckdb ==0.9.1
+  - polars ==0.19.8
   - pip
   - coiled >=0.2.54
   - numpy ==1.24.4
diff --git a/cluster_kwargs.yaml b/cluster_kwargs.yaml
index 39dc2b94d6..1d48fa267f 100644
--- a/cluster_kwargs.yaml
+++ b/cluster_kwargs.yaml
@@ -57,9 +57,30 @@ uber_lyft_large:
   n_workers: 50
   worker_vm_types: [m6i.xlarge] # 4 CPU, 16 GiB (preferred default instance)
 
+tpch_pyspark:
+  n_workers: 20
+  worker_vm_types: [m6i.xlarge] # 4 CPU, 16 GiB (preferred default instance)
+  # Keep the same, as spark.driver/executor.memory calculated from scheduler size
+  # pyspark doesn't automatically calculate available memory for executors, 
+  # so we do it in pyspark_queries.utils, but this takes place on the scheduler
+  scheduler_vm_types: [m6i.xlarge]
+  scheduler_options:
+    idle_timeout: null  # Scheduler won't get any Dask tasks.
+  backend_options:
+    ingress:
+      - ports: [8786, 8787, 7077, 8080, 4040, 9797]
+        cidr: "0.0.0.0/0"
+    spot: false
+    spot_on_demand_fallback: true
+    multizone: true
+
 tpch:
   n_workers: 20
   worker_vm_types: [m6i.xlarge] # 4 CPU, 16 GiB (preferred default instance)
+  backend_options:
+    spot: false
+    spot_on_demand_fallback: true
+    multizone: true
 
 # For tests/workflows/test_pytorch_optuna.py
 pytorch_optuna:


Test patch:
diff --git a/tests/benchmarks/tpch/__init__.py b/tests/benchmarks/tpch/__init__.py
new file mode 100644
index 0000000000..d1698b094e
--- /dev/null
+++ b/tests/benchmarks/tpch/__init__.py
@@ -0,0 +1,1 @@
+from . import test_dask, test_duckdb, test_polars, test_pyspark  # noqa: F401
diff --git a/tests/benchmarks/tpch/conftest.py b/tests/benchmarks/tpch/conftest.py
new file mode 100644
index 0000000000..5fac0bfc33
--- /dev/null
+++ b/tests/benchmarks/tpch/conftest.py
@@ -0,0 +1,71 @@
+import functools
+import os
+
+import coiled
+import pytest
+
+DATASETS = {
+    "local": "./tpch-data/scale10/",
+    "scale 100": "s3://coiled-runtime-ci/tpch_scale_100/",
+    "scale 1000": "s3://coiled-runtime-ci/tpch-scale-1000/",
+}
+
+ENABLED_DATASET = os.getenv("TPCH_SCALE")
+if ENABLED_DATASET is not None:
+    if ENABLED_DATASET not in DATASETS:
+        raise ValueError("Unknown tpch dataset: ", ENABLED_DATASET)
+else:
+    ENABLED_DATASET = "scale 100"
+
+machine = {
+    "memory": "256 GiB",
+}
+
+
+@pytest.fixture(scope="module")
+def warm_start():
+    @coiled.function(**machine)
+    def _():
+        pass
+
+    _()  # run once to give us a warm start
+
+
+@pytest.fixture(scope="function")
+def restart(warm_start):
+    @coiled.function(**machine)
+    def _():
+        pass
+
+    _.client.restart()
+    yield
+
+
+def coiled_function(**kwargs):
+    # Shouldn't be necessary
+    # See https://github.com/coiled/platform/issues/3519
+    def _(function):
+        return functools.wraps(function)(coiled.function(**kwargs, **machine)(function))
+
+    return _
+
+
+@pytest.fixture
+def tpch_dataset_name():
+    return ENABLED_DATASET
+
+
+@pytest.fixture
+def tpch_dataset_path(tpch_dataset_name):
+    return DATASETS[tpch_dataset_name]
+
+
+@pytest.fixture
+def vm_type():
+    return "m6i.16xlarge"
+
+
+@pytest.fixture
+def region():
+    # Region of the TPCH data
+    return "us-east-2"
diff --git a/tests/benchmarks/tpch/pyspark_queries/__init__.py b/tests/benchmarks/tpch/pyspark_queries/__init__.py
new file mode 100644
index 0000000000..f8af6039b8
--- /dev/null
+++ b/tests/benchmarks/tpch/pyspark_queries/__init__.py
@@ -0,0 +1,25 @@
+from . import (  # noqa: F401
+    q1,
+    q2,
+    q3,
+    q4,
+    q5,
+    q6,
+    q7,
+    q8,
+    q9,
+    q10,
+    q11,
+    q12,
+    q13,
+    q14,
+    q15,
+    q16,
+    q17,
+    q18,
+    q19,
+    q20,
+    q21,
+    q22,
+    utils,
+)
diff --git a/tests/benchmarks/tpch/pyspark_queries/q1.py b/tests/benchmarks/tpch/pyspark_queries/q1.py
new file mode 100644
index 0000000000..dce8e27b66
--- /dev/null
+++ b/tests/benchmarks/tpch/pyspark_queries/q1.py
@@ -0,0 +1,28 @@
+query = """select
+        l_returnflag,
+        l_linestatus,
+        sum(l_quantity) as sum_qty,
+        sum(l_extendedprice) as sum_base_price,
+        sum(l_extendedprice * (1 - l_discount)) as sum_disc_price,
+        sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge,
+        avg(l_quantity) as avg_qty,
+        avg(l_extendedprice) as avg_price,
+        avg(l_discount) as avg_disc,
+        count(*) as count_order
+    from
+        lineitem
+    where
+        l_shipdate <= date('1998-09-02')
+    group by
+        l_returnflag,
+        l_linestatus
+    order by
+        l_returnflag,
+        l_linestatus
+"""
+
+
+def setup(spark):
+    from .utils import read_parquet_spark
+
+    read_parquet_spark(spark, "lineitem", "lineitem")
diff --git a/tests/benchmarks/tpch/pyspark_queries/q2.py b/tests/benchmarks/tpch/pyspark_queries/q2.py
new file mode 100644
index 0000000000..be53c761d1
--- /dev/null
+++ b/tests/benchmarks/tpch/pyspark_queries/q2.py
@@ -0,0 +1,52 @@
+query = """select
+        s_acctbal,
+        s_name,
+        n_name,
+        p_partkey,
+        p_mfgr,
+        s_address,
+        s_phone,
+        s_comment
+    from
+        part,
+        supplier,
+        partsupp,
+        nation,
+        region
+    where
+        p_partkey = ps_partkey
+        and s_suppkey = ps_suppkey
+        and p_size = 15
+        and p_type like '%BRASS'
+        and s_nationkey = n_nationkey
+        and n_regionkey = r_regionkey
+        and r_name = 'EUROPE'
+        and ps_supplycost = (
+            select
+                min(ps_supplycost)
+            from
+                partsupp,
+                supplier,
+                nation,
+                region
+            where
+                p_partkey = ps_partkey
+                and s_suppkey = ps_suppkey
+                and s_nationkey = n_nationkey
+                and n_regionkey = r_regionkey
+                and r_name = 'EUROPE'
+        )
+    order by
+        s_acctbal desc,
+        n_name,
+        s_name,
+        p_partkey
+    limit 100
+"""
+
+
+def setup(spark):
+    from .utils import read_parquet_spark
+
+    for name in ("part", "supplier", "partsupp", "nation", "region"):
+        read_parquet_spark(spark, name, name)
diff --git a/tests/benchmarks/tpch/pyspark_queries/q22.py b/tests/benchmarks/tpch/pyspark_queries/q22.py
new file mode 100644
index 0000000000..a416bd6aae
--- /dev/null
+++ b/tests/benchmarks/tpch/pyspark_queries/q22.py
@@ -0,0 +1,45 @@
+query = """
+    select
+        cntrycode,
+        count(*) as numcust,
+        sum(c_acctbal) as totacctbal
+    from (
+        select
+            substring(c_phone from 1 for 2) as cntrycode,
+            c_acctbal
+        from
+            customer
+        where
+            substring(c_phone from 1 for 2) in
+                ("13","31","23", "29", "30", "18", "17")
+            and c_acctbal > (
+                select
+                    avg(c_acctbal)
+                from
+                    customer
+                where
+                    c_acctbal > 0.00
+                    and substring (c_phone from 1 for 2) in
+                        ("13","31","23", "29", "30", "18", "17")
+            )
+            and not exists (
+                select
+                    *
+                from
+                    orders
+                where
+                    o_custkey = c_custkey
+            )
+        ) as custsale
+    group by
+        cntrycode
+    order by
+        cntrycode
+"""
+
+
+def setup(spark):
+    from .utils import read_parquet_spark
+
+    for name in ("orders", "customer"):
+        read_parquet_spark(spark, name, name)
diff --git a/tests/benchmarks/tpch/pyspark_queries/q3.py b/tests/benchmarks/tpch/pyspark_queries/q3.py
new file mode 100644
index 0000000000..1fca8912c0
--- /dev/null
+++ b/tests/benchmarks/tpch/pyspark_queries/q3.py
@@ -0,0 +1,32 @@
+query = """
+    select
+        l_orderkey,
+        sum(l_extendedprice * (1 - l_discount)) as revenue,
+        o_orderdate,
+        o_shippriority
+    from
+        customer,
+        orders,
+        lineitem
+    where
+        c_mktsegment = 'BUILDING'
+        and c_custkey = o_custkey
+        and l_orderkey = o_orderkey
+        and o_orderdate < date '1995-03-15'
+        and l_shipdate > date '1995-03-15'
+    group by
+        l_orderkey,
+        o_orderdate,
+        o_shippriority
+    order by
+        revenue desc,
+        o_orderdate
+    limit 10
+"""
+
+
+def setup(spark):
+    from .utils import read_parquet_spark
+
+    for name in ("customer", "orders", "lineitem"):
+        read_parquet_spark(spark, name, name)
diff --git a/tests/benchmarks/tpch/pyspark_queries/q4.py b/tests/benchmarks/tpch/pyspark_queries/q4.py
new file mode 100644
index 0000000000..9d78898213
--- /dev/null
+++ b/tests/benchmarks/tpch/pyspark_queries/q4.py
@@ -0,0 +1,30 @@
+query = """
+    select
+        o_orderpriority,
+        count(*) as order_count
+    from
+        orders
+    where
+        o_orderdate >= date '1993-07-01'
+        and o_orderdate < date '1993-07-01' + interval '3' month
+        and exists (
+            select
+                *
+            from
+                lineitem
+            where
+                l_orderkey = o_orderkey
+                and l_commitdate < l_receiptdate
+        )
+    group by
+        o_orderpriority
+    order by
+        o_orderpriority
+"""
+
+
+def setup(spark):
+    from .utils import read_parquet_spark
+
+    for name in ("orders", "lineitem"):
+        read_parquet_spark(spark, name, name)
diff --git a/tests/benchmarks/tpch/pyspark_queries/q5.py b/tests/benchmarks/tpch/pyspark_queries/q5.py
new file mode 100644
index 0000000000..a9956666a7
--- /dev/null
+++ b/tests/benchmarks/tpch/pyspark_queries/q5.py
@@ -0,0 +1,40 @@
+query = """
+    select
+        n_name,
+        sum(l_extendedprice * (1 - l_discount)) as revenue
+    from
+        customer,
+        orders,
+        lineitem,
+        supplier,
+        nation,
+        region
+    where
+        c_custkey = o_custkey
+        and l_orderkey = o_orderkey
+        and l_suppkey = s_suppkey
+        and c_nationkey = s_nationkey
+        and s_nationkey = n_nationkey
+        and n_regionkey = r_regionkey
+        and r_name = 'ASIA'
+        and o_orderdate >= date '1994-01-01'
+        and o_orderdate < date '1994-01-01' + interval '1' year
+    group by
+        n_name
+    order by
+        revenue desc
+"""
+
+
+def setup(spark):
+    from .utils import read_parquet_spark
+
+    for name in (
+        "customer",
+        "orders",
+        "lineitem",
+        "supplier",
+        "nation",
+        "region",
+    ):
+        read_parquet_spark(spark, name, name)
diff --git a/tests/benchmarks/tpch/pyspark_queries/q6.py b/tests/benchmarks/tpch/pyspark_queries/q6.py
new file mode 100644
index 0000000000..4d1ff28ba2
--- /dev/null
+++ b/tests/benchmarks/tpch/pyspark_queries/q6.py
@@ -0,0 +1,17 @@
+query = """
+select
+        sum(l_extendedprice * l_discount) as revenue
+    from
+        lineitem
+    where
+        l_shipdate >= date '1994-01-01'
+        and l_shipdate < date '1994-01-01' + interval '1' year
+        and l_discount between .06 - 0.01 and .06 + 0.01
+        and l_quantity < 24
+"""
+
+
+def setup(spark):
+    from .utils import read_parquet_spark
+
+    read_parquet_spark(spark, "lineitem", "lineitem")
diff --git a/tests/benchmarks/tpch/pyspark_queries/q7.py b/tests/benchmarks/tpch/pyspark_queries/q7.py
new file mode 100644
index 0000000000..db6c65edd8
--- /dev/null
+++ b/tests/benchmarks/tpch/pyspark_queries/q7.py
@@ -0,0 +1,48 @@
+query = """
+select
+        supp_nation,
+        cust_nation,
+        l_year,
+        sum(volume) as revenue
+    from
+        (
+            select
+                n1.n_name as supp_nation,
+                n2.n_name as cust_nation,
+                year(l_shipdate) as l_year,
+                l_extendedprice * (1 - l_discount) as volume
+            from
+                supplier,
+                lineitem,
+                orders,
+                customer,
+                nation n1,
+                nation n2
+            where
+                s_suppkey = l_suppkey
+                and o_orderkey = l_orderkey
+                and c_custkey = o_custkey
+                and s_nationkey = n1.n_nationkey
+                and c_nationkey = n2.n_nationkey
+                and (
+                    (n1.n_name = 'FRANCE' and n2.n_name = 'GERMANY')
+                    or (n1.n_name = 'GERMANY' and n2.n_name = 'FRANCE')
+                )
+                and l_shipdate between date '1995-01-01' and date '1996-12-31'
+        ) as shipping
+    group by
+        supp_nation,
+        cust_nation,
+        l_year
+    order by
+        supp_nation,
+        cust_nation,
+        l_year
+"""
+
+
+def setup(spark):
+    from .utils import read_parquet_spark
+
+    for name in ("supplier", "lineitem", "orders", "customer", "nation"):
+        read_parquet_spark(spark, name, name)
diff --git a/tests/benchmarks/tpch/pyspark_queries/utils.py b/tests/benchmarks/tpch/pyspark_queries/utils.py
new file mode 100644
index 0000000000..447e545511
--- /dev/null
+++ b/tests/benchmarks/tpch/pyspark_queries/utils.py
@@ -0,0 +1,46 @@
+import psutil
+
+# Allow overriding in single VM testing
+MASTER = None
+
+
+def get_or_create_spark(name: str = "SparkTest"):
+    from distributed.utils import get_ip
+    from pyspark.sql import SparkSession
+
+    # Memory for executors/driver - ~80% of total GiB
+    memory = int((psutil.virtual_memory().total * 0.8) * 9.313226e-10)
+    n_cores = psutil.cpu_count()
+
+    master = f"spark://{get_ip()}:7077"
+
+    spark = (
+        SparkSession.builder.master(MASTER or master)
+        .appName(name)
+        .config(
+            "spark.hadoop.fs.s3a.aws.credentials.provider",
+            "com.amazonaws.auth.EnvironmentVariableCredentialsProvider",
+        )
+        # ref: https://issues.apache.org/jira/browse/SPARK-44988  (unresolved, v3.4.0 and v3.4.1 affected)
+        # Illegal Parquet type: INT64 (TIMESTAMP(NANOS,true))
+        .config("spark.sql.legacy.parquet.nanosAsLong", "true")
+        .config("spark.memory.fraction", "0.7")
+        # TODO: Can these not be set automagically by pyspark?!
+        # ref: https://spark.apache.org/docs/latest/configuration.html#memory-management
+        .config("spark.driver.cores", f"{n_cores}")  # defaults to 1
+        .config("spark.driver.memory", f"{memory}g")  # defaults to 1GB
+        .config("spark.executor.memory", f"{memory}g")  # defaults to 1GB
+        .getOrCreate()
+    )
+    return spark
+
+
+def read_parquet_spark(spark, filename: str, table_name: str):
+    from tests.benchmarks.tpch.test_pyspark import DATASETS, ENABLED_DATASET
+
+    path = DATASETS[ENABLED_DATASET] + filename + "/"
+    print(f"Read from {path=}")
+    path = path.replace("s3://", "s3a://")
+    df = spark.read.parquet(path)
+    df.createOrReplaceTempView(table_name)
+    return df
diff --git a/tests/benchmarks/test_tpch.py b/tests/benchmarks/tpch/test_dask.py
similarity index 96%
rename from tests/benchmarks/test_tpch.py
rename to tests/benchmarks/tpch/test_dask.py
index 19e7133cbe..7c16fd7af1 100644
--- a/tests/benchmarks/test_tpch.py
+++ b/tests/benchmarks/tpch/test_dask.py
@@ -1,23 +1,12 @@
-import os
 from datetime import datetime
 
 import dask_expr as dd
 
-DATASETS = {
-    "scale 100": "s3://coiled-runtime-ci/tpch_scale_100/",
-    "scale 1000": "s3://coiled-runtime-ci/tpch-scale-1000/",
-}
-
-enabled_dataset = os.getenv("TPCH_SCALE")
-if enabled_dataset is not None:
-    if enabled_dataset not in DATASETS:
-        raise ValueError("Unknown tpch dataset: ", enabled_dataset)
-else:
-    enabled_dataset = "scale 100"
+from .conftest import DATASETS, ENABLED_DATASET
 
 
 def read_data(filename):
-    path = DATASETS[enabled_dataset] + filename + "/"
+    path = DATASETS[ENABLED_DATASET] + filename + "/"
     return dd.read_parquet(path, engine="pyarrow")
 
 
diff --git a/tests/benchmarks/tpch/test_duckdb.py b/tests/benchmarks/tpch/test_duckdb.py
new file mode 100644
index 0000000000..94890f4c5d
--- /dev/null
+++ b/tests/benchmarks/tpch/test_duckdb.py
@@ -0,0 +1,297 @@
+import botocore
+import duckdb
+import pytest
+
+from tests.benchmarks.tpch.conftest import DATASETS, ENABLED_DATASET, coiled_function
+
+
+@pytest.fixture
+def connection():
+    def _():
+        con = duckdb.connect()
+
+        if ENABLED_DATASET != "local":  # Setup s3 credentials
+            creds = botocore.session.get_session().get_credentials()
+            con.install_extension("httpfs")
+            con.load_extension("httpfs")
+            con.sql(
+                f"""
+                SET s3_region='us-east-2';
+                SET s3_access_key_id='{creds.access_key}';
+                SET s3_secret_access_key='{creds.secret_key}';
+                """
+            )
+        return con
+
+    return _
+
+
+@coiled_function()
+def test_query_1(connection):
+    connection().execute(
+        f"""
+        with lineitem as (
+            select * from read_parquet('{DATASETS[ENABLED_DATASET]}lineitem/*.parquet')
+        )
+
+        select
+            l_returnflag,
+            l_linestatus,
+            sum(l_quantity) as sum_qty,
+            sum(l_extendedprice) as sum_base_price,
+            sum(l_extendedprice * (1 - l_discount)) as sum_disc_price,
+            sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge,
+            avg(l_quantity) as avg_qty,
+            avg(l_extendedprice) as avg_price,
+            avg(l_discount) as avg_disc,
+            count(*) as count_order
+        from
+            lineitem
+        where
+            l_shipdate <= date '1998-09-02'
+        group by
+            l_returnflag,
+            l_linestatus
+        order by
+            l_returnflag,
+            l_linestatus
+        """
+    ).arrow()
+
+
+@coiled_function()
+def test_query_2(connection):
+    connection().execute(
+        f"""
+        with part as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}part/*.parquet')),
+             supplier as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}supplier/*.parquet')),
+             partsupp as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}partsupp/*.parquet')),
+             nation as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}nation/*.parquet')),
+             region as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}region/*.parquet'))
+
+        select
+            s_acctbal,
+            s_name,
+            n_name,
+            p_partkey,
+            p_mfgr,
+            s_address,
+            s_phone,
+            s_comment
+        from
+            part,
+            supplier,
+            partsupp,
+            nation,
+            region
+        where
+            p_partkey = ps_partkey
+            and s_suppkey = ps_suppkey
+            and p_size = 15
+            and p_type like '%BRASS'
+            and s_nationkey = n_nationkey
+            and n_regionkey = r_regionkey
+            and r_name = 'EUROPE'
+            and ps_supplycost = (
+                select
+                    min(ps_supplycost)
+                from
+                    partsupp,
+                    supplier,
+                    nation,
+                    region
+                where
+                    p_partkey = ps_partkey
+                    and s_suppkey = ps_suppkey
+                    and s_nationkey = n_nationkey
+                    and n_regionkey = r_regionkey
+                    and r_name = 'EUROPE'
+            )
+        order by
+            s_acctbal desc,
+            n_name,
+            s_name,
+            p_partkey
+        limit 100
+        """
+    ).arrow()
+
+
+@coiled_function()
+def test_query_3(connection):
+    connection().execute(
+        f"""
+        with customer as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}customer/*.parquet')),
+             orders as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}orders/*.parquet')),
+             lineitem as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}lineitem/*.parquet'))
+
+        select
+            l_orderkey,
+            sum(l_extendedprice * (1 - l_discount)) as revenue,
+            o_orderdate,
+            o_shippriority
+        from
+            customer,
+            orders,
+            lineitem
+        where
+            c_mktsegment = 'BUILDING'
+            and c_custkey = o_custkey
+            and l_orderkey = o_orderkey
+            and o_orderdate < date '1995-03-15'
+            and l_shipdate > date '1995-03-15'
+        group by
+            l_orderkey,
+            o_orderdate,
+            o_shippriority
+        order by
+            revenue desc,
+            o_orderdate
+        limit 10
+        """
+    ).arrow()
+
+
+@coiled_function()
+def test_query_4(connection):
+    connection().execute(
+        f"""
+        with orders as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}orders/*.parquet')),
+             lineitem as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}lineitem/*.parquet'))
+
+        select
+            o_orderpriority,
+            count(*) as order_count
+        from
+            orders
+        where
+            o_orderdate >= date '1993-07-01'
+            and o_orderdate < date '1993-07-01' + interval '3' month
+            and exists (
+                select
+                    *
+                from
+                    lineitem
+                where
+                    l_orderkey = o_orderkey
+                    and l_commitdate < l_receiptdate
+            )
+        group by
+            o_orderpriority
+        order by
+            o_orderpriority
+    """
+    ).arrow()
+
+
+@coiled_function()
+def test_query_5(connection):
+    connection().execute(
+        f"""
+        with customer as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}customer/*.parquet')),
+             orders as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}orders/*.parquet')),
+             lineitem as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}lineitem/*.parquet')),
+             supplier as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}supplier/*.parquet')),
+             nation as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}nation/*.parquet')),
+             region as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}region/*.parquet'))
+
+        select
+            n_name,
+            sum(l_extendedprice * (1 - l_discount)) as revenue
+        from
+            customer,
+            orders,
+            lineitem,
+            supplier,
+            nation,
+            region
+        where
+            c_custkey = o_custkey
+            and l_orderkey = o_orderkey
+            and l_suppkey = s_suppkey
+            and c_nationkey = s_nationkey
+            and s_nationkey = n_nationkey
+            and n_regionkey = r_regionkey
+            and r_name = 'ASIA'
+            and o_orderdate >= date '1994-01-01'
+            and o_orderdate < date '1994-01-01' + interval '1' year
+        group by
+            n_name
+        order by
+            revenue desc
+    """
+    ).arrow()
+
+
+@coiled_function()
+def test_query_6(connection):
+    connection().execute(
+        f"""
+        with lineitem as (
+            select * from read_parquet('{DATASETS[ENABLED_DATASET]}lineitem/*.parquet')
+        )
+
+        select
+            sum(l_extendedprice * l_discount) as revenue
+        from
+            lineitem
+        where
+            l_shipdate >= date '1994-01-01'
+            and l_shipdate < date '1994-01-01' + interval '1' year
+            and l_discount between .06 - 0.01 and .06 + 0.01
+            and l_quantity < 24
+        """
+    ).arrow()
+
+
+@coiled_function()
+def test_query_7(connection):
+    connection().execute(
+        f"""
+        with supplier as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}supplier/*.parquet')),
+             lineitem as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}lineitem/*.parquet')),
+             orders as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}orders/*.parquet')),
+             customer as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}customer/*.parquet')),
+             nation as (select * from read_parquet('{DATASETS[ENABLED_DATASET]}nation/*.parquet'))
+
+        select
+            supp_nation,
+            cust_nation,
+            l_year,
+            sum(volume) as revenue
+        from
+            (
+                select
+                    n1.n_name as supp_nation,
+                    n2.n_name as cust_nation,
+                    year(l_shipdate) as l_year,
+                    l_extendedprice * (1 - l_discount) as volume
+                from
+                    supplier,
+                    lineitem,
+                    orders,
+                    customer,
+                    nation n1,
+                    nation n2
+                where
+                    s_suppkey = l_suppkey
+                    and o_orderkey = l_orderkey
+                    and c_custkey = o_custkey
+                    and s_nationkey = n1.n_nationkey
+                    and c_nationkey = n2.n_nationkey
+                    and (
+                        (n1.n_name = 'FRANCE' and n2.n_name = 'GERMANY')
+                        or (n1.n_name = 'GERMANY' and n2.n_name = 'FRANCE')
+                    )
+                    and l_shipdate between date '1995-01-01' and date '1996-12-31'
+            ) as shipping
+        group by
+            supp_nation,
+            cust_nation,
+            l_year
+        order by
+            supp_nation,
+            cust_nation,
+            l_year
+        """
+    ).arrow()
diff --git a/tests/benchmarks/tpch/test_polars.py b/tests/benchmarks/tpch/test_polars.py
new file mode 100644
index 0000000000..31309db0c7
--- /dev/null
+++ b/tests/benchmarks/tpch/test_polars.py
@@ -0,0 +1,256 @@
+from datetime import datetime
+
+import polars as pl
+from pyarrow.dataset import dataset
+
+from tests.benchmarks.tpch.conftest import DATASETS, ENABLED_DATASET, coiled_function
+
+
+def read_data(name, source=None):
+    pyarrow_dataset = dataset(
+        source=source or f"{DATASETS[ENABLED_DATASET]}{name}/", format="parquet"
+    )
+    return pl.scan_pyarrow_dataset(pyarrow_dataset)
+
+
+@coiled_function()
+def test_query_1(restart):
+    var_1 = datetime(1998, 9, 2)
+    q = read_data("lineitem")
+    (
+        q.filter(pl.col("l_shipdate") <= var_1)
+        .group_by(["l_returnflag", "l_linestatus"])
+        .agg(
+            [
+                pl.sum("l_quantity").alias("sum_qty"),
+                pl.sum("l_extendedprice").alias("sum_base_price"),
+                (pl.col("l_extendedprice") * (1 - pl.col("l_discount")))
+                .sum()
+                .alias("sum_disc_price"),
+                (
+                    pl.col("l_extendedprice")
+                    * (1.0 - pl.col("l_discount"))
+                    * (1.0 + pl.col("l_tax"))
+                )
+                .sum()
+                .alias("sum_charge"),
+                pl.mean("l_quantity").alias("avg_qty"),
+                pl.mean("l_extendedprice").alias("avg_price"),
+                pl.mean("l_discount").alias("avg_disc"),
+                pl.count().alias("count_order"),
+            ],
+        )
+        .sort(["l_returnflag", "l_linestatus"])
+    ).collect(streaming=True)
+
+
+@coiled_function()
+def test_query_2(restart):
+    print("Hello")
+    var_1 = 15
+    var_2 = "BRASS"
+    var_3 = "EUROPE"
+
+    print(1)
+    region_ds = read_data("region")
+    nation_ds = read_data("nation")
+    supplier_ds = read_data("supplier")
+    part_ds = read_data("part")
+    part_supp_ds = read_data("partsupp")
+    print(2)
+
+    result_q1 = (
+        part_ds.join(part_supp_ds, left_on="p_partkey", right_on="ps_partkey")
+        .join(supplier_ds, left_on="ps_suppkey", right_on="s_suppkey")
+        .join(nation_ds, left_on="s_nationkey", right_on="n_nationkey")
+        .join(region_ds, left_on="n_regionkey", right_on="r_regionkey")
+        .filter(pl.col("p_size") == var_1)
+        .filter(pl.col("p_type").str.ends_with(var_2))
+        .filter(pl.col("r_name") == var_3)
+    ).cache()
+    print(3)
+
+    final_cols = [
+        "s_acctbal",
+        "s_name",
+        "n_name",
+        "p_partkey",
+        "p_mfgr",
+        "s_address",
+        "s_phone",
+        "s_comment",
+    ]
+
+    print(4)
+    (
+        result_q1.group_by("p_partkey")
+        .agg(pl.min("ps_supplycost").alias("ps_supplycost"))
+        .join(
+            result_q1,
+            left_on=["p_partkey", "ps_supplycost"],
+            right_on=["p_partkey", "ps_supplycost"],
+        )
+        .select(final_cols)
+        .sort(
+            by=["s_acctbal", "n_name", "s_name", "p_partkey"],
+            descending=[True, False, False, False],
+        )
+        .limit(100)
+        .with_columns(pl.col(pl.datatypes.Utf8).str.strip().keep_name())
+    ).collect(streaming=True)
+    print(5)
+
+
+@coiled_function()
+def test_query_3(restart):
+    var_1 = var_2 = datetime(1995, 3, 15)
+    var_3 = "BUILDING"
+
+    customer_ds = read_data("customer")
+    line_item_ds = read_data("lineitem")
+    orders_ds = read_data("orders")
+
+    (
+        customer_ds.filter(pl.col("c_mktsegment") == var_3)
+        .join(orders_ds, left_on="c_custkey", right_on="o_custkey")
+        .join(line_item_ds, left_on="o_orderkey", right_on="l_orderkey")
+        .filter(pl.col("o_orderdate") < var_2)
+        .filter(pl.col("l_shipdate") > var_1)
+        .with_columns(
+            (pl.col("l_extendedprice") * (1 - pl.col("l_discount"))).alias("revenue")
+        )
+        .group_by(["o_orderkey", "o_orderdate", "o_shippriority"])
+        .agg([pl.sum("revenue")])
+        .select(
+            [
+                pl.col("o_orderkey").alias("l_orderkey"),
+                "revenue",
+                "o_orderdate",
+                "o_shippriority",
+            ]
+        )
+        .sort(by=["revenue", "o_orderdate"], descending=[True, False])
+        .limit(10)
+    ).collect(streaming=True)
+
+
+@coiled_function()
+def test_query_4(restart):
+    var_1 = datetime(1993, 7, 1)
+    var_2 = datetime(1993, 10, 1)
+
+    line_item_ds = read_data("lineitem")
+    orders_ds = read_data("orders")
+
+    (
+        line_item_ds.join(orders_ds, left_on="l_orderkey", right_on="o_orderkey")
+        .filter(pl.col("o_orderdate").is_between(var_1, var_2, closed="left"))
+        .filter(pl.col("l_commitdate") < pl.col("l_receiptdate"))
+        .unique(subset=["o_orderpriority", "l_orderkey"])
+        .group_by("o_orderpriority")
+        .agg(pl.count().alias("order_count"))
+        .sort(by="o_orderpriority")
+        .with_columns(pl.col("order_count").cast(pl.datatypes.Int64))
+    ).collect(streaming=True)
+
+
+@coiled_function()
+def test_query_5(restart):
+    var_1 = "ASIA"
+    var_2 = datetime(1994, 1, 1)
+    var_3 = datetime(1995, 1, 1)
+
+    region_ds = read_data("region")
+    nation_ds = read_data("nation")
+    customer_ds = read_data("customer")
+    line_item_ds = read_data("lineitem")
+    orders_ds = read_data("orders")
+    supplier_ds = read_data("supplier")
+
+    (
+        region_ds.join(nation_ds, left_on="r_regionkey", right_on="n_regionkey")
+        .join(customer_ds, left_on="n_nationkey", right_on="c_nationkey")
+        .join(orders_ds, left_on="c_custkey", right_on="o_custkey")
+        .join(line_item_ds, left_on="o_orderkey", right_on="l_orderkey")
+        .join(
+            supplier_ds,
+            left_on=["l_suppkey", "n_nationkey"],
+            right_on=["s_suppkey", "s_nationkey"],
+        )
+        .filter(pl.col("r_name") == var_1)
+        .filter(pl.col("o_orderdate").is_between(var_2, var_3, closed="left"))
+        .with_columns(
+            (pl.col("l_extendedprice") * (1 - pl.col("l_discount"))).alias("revenue")
+        )
+        .group_by("n_name")
+        .agg([pl.sum("revenue")])
+        .sort(by="revenue", descending=True)
+    ).collect(streaming=True)
+
+
+@coiled_function()
+def test_query_6(restart):
+    var_1 = datetime(1994, 1, 1)
+    var_2 = datetime(1995, 1, 1)
+    var_3 = 24
+
+    line_item_ds = read_data("lineitem")
+
+    (
+        line_item_ds.filter(
+            pl.col("l_shipdate").is_between(var_1, var_2, closed="left")
+        )
+        .filter(pl.col("l_discount").is_between(0.05, 0.07))
+        .filter(pl.col("l_quantity") < var_3)
+        .with_columns(
+            (pl.col("l_extendedprice") * pl.col("l_discount")).alias("revenue")
+        )
+        .select(pl.sum("revenue").alias("revenue"))
+    ).collect(streaming=True)
+
+
+@coiled_function()
+def test_query_7(restart):
+    nation_ds = read_data("nation")
+    customer_ds = read_data("customer")
+    line_item_ds = read_data("lineitem")
+    orders_ds = read_data("orders")
+    supplier_ds = read_data("supplier")
+
+    n1 = nation_ds.filter(pl.col("n_name") == "FRANCE")
+    n2 = nation_ds.filter(pl.col("n_name") == "GERMANY")
+
+    var_1 = datetime(1995, 1, 1)
+    var_2 = datetime(1996, 12, 31)
+
+    df1 = (
+        customer_ds.join(n1, left_on="c_nationkey", right_on="n_nationkey")
+        .join(orders_ds, left_on="c_custkey", right_on="o_custkey")
+        .rename({"n_name": "cust_nation"})
+        .join(line_item_ds, left_on="o_orderkey", right_on="l_orderkey")
+        .join(supplier_ds, left_on="l_suppkey", right_on="s_suppkey")
+        .join(n2, left_on="s_nationkey", right_on="n_nationkey")
+        .rename({"n_name": "supp_nation"})
+    )
+
+    df2 = (
+        customer_ds.join(n2, left_on="c_nationkey", right_on="n_nationkey")
+        .join(orders_ds, left_on="c_custkey", right_on="o_custkey")
+        .rename({"n_name": "cust_nation"})
+        .join(line_item_ds, left_on="o_orderkey", right_on="l_orderkey")
+        .join(supplier_ds, left_on="l_suppkey", right_on="s_suppkey")
+        .join(n1, left_on="s_nationkey", right_on="n_nationkey")
+        .rename({"n_name": "supp_nation"})
+    )
+
+    (
+        pl.concat([df1, df2])
+        .filter(pl.col("l_shipdate").is_between(var_1, var_2))
+        .with_columns(
+            (pl.col("l_extendedprice") * (1 - pl.col("l_discount"))).alias("volume")
+        )
+        .with_columns(pl.col("l_shipdate").dt.year().alias("l_year"))
+        .group_by(["supp_nation", "cust_nation", "l_year"])
+        .agg([pl.sum("volume").alias("revenue")])
+        .sort(by=["supp_nation", "cust_nation", "l_year"])
+    ).collect(streaming=True)
diff --git a/tests/benchmarks/tpch/test_pyspark.py b/tests/benchmarks/tpch/test_pyspark.py
new file mode 100644
index 0000000000..392db6a821
--- /dev/null
+++ b/tests/benchmarks/tpch/test_pyspark.py
@@ -0,0 +1,229 @@
+import asyncio
+import logging
+import os
+import pathlib
+import re
+import shlex
+import subprocess
+import sys
+
+import pyspark
+from conda.cli.python_api import Commands, run_command
+from distributed.diagnostics.plugin import SchedulerPlugin, WorkerPlugin
+from distributed.utils import get_ip
+
+from tests.benchmarks.tpch.conftest import DATASETS, ENABLED_DATASET  # noqa: F401
+
+from . import pyspark_queries as queries
+
+# pyspark/hadoop/aws-java-* are extra sensitive to version changes it seems.
+# ref: https://hadoop.apache.org/docs/r3.3.4/hadoop-aws/dependency-analysis.html
+# for verifying aws-java-sdk-bundle version w/ hadoop version.
+assert pyspark.__version__ == "3.4.1"
+HADOOP_AWS_VERSION = "3.3.4"  # sc._jvm.org.apache.hadoop.util.VersionInfo.getVersion()
+AWS_JAVA_SDK_BUNDLE_VERSION = "1.12.262"
+
+
+logger = logging.getLogger("coiled")
+
+PACKAGES = (
+    f"org.apache.hadoop:hadoop-client:{HADOOP_AWS_VERSION}",
+    f"org.apache.hadoop:hadoop-common:{HADOOP_AWS_VERSION}",
+    f"org.apache.hadoop:hadoop-aws:{HADOOP_AWS_VERSION}",
+    f"com.amazonaws:aws-java-sdk-bundle:{AWS_JAVA_SDK_BUNDLE_VERSION}",
+)
+
+
+def test_query_1(tpch_pyspark_client):
+    return run_tpch_pyspark(tpch_pyspark_client, queries.q1)
+
+
+def test_query_2(tpch_pyspark_client):
+    return run_tpch_pyspark(tpch_pyspark_client, queries.q2)
+
+
+def test_query_3(tpch_pyspark_client):
+    return run_tpch_pyspark(tpch_pyspark_client, queries.q3)
+
+
+def test_query_4(tpch_pyspark_client):
+    return run_tpch_pyspark(tpch_pyspark_client, queries.q4)
+
+
+def test_query_5(tpch_pyspark_client):
+    return run_tpch_pyspark(tpch_pyspark_client, queries.q5)
+
+
+def test_query_6(tpch_pyspark_client):
+    return run_tpch_pyspark(tpch_pyspark_client, queries.q6)
+
+
+def test_query_7(tpch_pyspark_client):
+    return run_tpch_pyspark(tpch_pyspark_client, queries.q7)
+
+
+def make_pyspark_submit_args(spark_master):
+    return f"""
+        --packages {','.join(PACKAGES)}
+        --exclude-packages com.google.guava:guava
+        --master=spark://{spark_master}
+        pyspark-shell
+    """
+
+
+def fix_timestamp_ns_columns(query):
+    """
+    scale100 stores l_shipdate/o_orderdate as timestamp[us]
+    scale1000 stores l_shipdate/o_orderdate as timestamp[ns] which gives:
+        Illegal Parquet type: INT64 (TIMESTAMP(NANOS,true))
+    so we set spark.sql.legacy.parquet.nanosAsLong and then convert to timestamp.
+    """
+    for name in ("l_shipdate", "o_orderdate"):
+        query = re.sub(rf"\b{name}\b", f"to_timestamp(cast({name} as string))", query)
+    return query
+
+
+def run_tpch_pyspark(tpch_pyspark_client, module):
+    from .pyspark_queries.utils import get_or_create_spark
+
+    is_local = tpch_pyspark_client is None
+
+    async def _run_tpch(dask_scheduler):
+        # Connecting to the Spark Connect server doens't appear to work very well
+        # If we just submit this stuff to the scheduler and generate a SparkSession there
+        # that connects to the running master we can at least get some code running
+
+        if not is_local:
+            queries.utils.MASTER = f"spark://{get_ip()}:7077"
+        else:
+            queries.utils.MASTER = "local[*]"
+
+        os.environ["SPARK_HOME"] = str(pathlib.Path(pyspark.__file__).absolute().parent)
+        os.environ["PYSPARK_PYTHON"] = sys.executable
+        os.environ["PYSPARK_SUBMIT_ARGS"] = make_pyspark_submit_args(
+            queries.utils.MASTER
+        )
+
+        def _():
+            spark = get_or_create_spark(f"query{module.__name__}")
+
+            # scale1000 stored as timestamp[ns] which spark parquet
+            # can't use natively.
+            if ENABLED_DATASET == "scale 1000":
+                module.query = fix_timestamp_ns_columns(module.query)
+
+            module.setup(spark)  # read spark tables query will select from
+            if hasattr(module, "ddl"):
+                spark.sql(module.ddl)  # might create temp view
+            q_final = spark.sql(module.query)  # benchmark query
+            try:
+                # trigger materialization of df
+                return q_final.toJSON().collect()
+            finally:
+                spark.catalog.clearCache()
+                spark.sparkContext.stop()
+                spark.stop()
+
+        return await asyncio.to_thread(_)
+
+    if not is_local:
+        rows = tpch_pyspark_client.run_on_scheduler(_run_tpch)
+    else:
+        rows = asyncio.run(_run_tpch(None))  # running locally
+    print(f"Received {len(rows)} rows")
+
+
+class SparkMaster(SchedulerPlugin):
+    def start(self, scheduler):
+        self.scheduler = scheduler
+        self.scheduler.add_plugin(self)
+
+        print("Installing procps...", end="")
+        try:
+            # All those nasty scripts require ``ps`` unix CLI util to be installed
+            _, stderr, returncode = run_command(
+                Commands.INSTALL, ["-c", "conda-forge", "procps-ng"]
+            )
+        except Exception as e:
+            msg = "conda install failed"
+            raise RuntimeError(msg) from e
+        print("done")
+        path = pathlib.Path(pyspark.__file__).absolute()
+        module_loc = path.parent
+        script = module_loc / "sbin" / "spark-daemon.sh"  # noqa: F841
+
+        host = scheduler.address.split("//")[1].split(":")[0]
+        spark_master = f"{host}:7077"
+
+        os.environ["SPARK_HOME"] = str(module_loc)
+        os.environ["PYSPARK_PYTHON"] = sys.executable
+        os.environ["PYSPARK_SUBMIT_ARGS"] = make_pyspark_submit_args(spark_master)
+
+        cmd = "spark-class org.apache.spark.deploy.master.Master"
+        print(f"Executing\n{cmd}")
+        self.proc_master = subprocess.Popen(shlex.split(cmd))
+
+        print("Launched Spark Master")
+
+    def close(self):
+        self.proc_master.terminate()
+        self.proc_master.wait()
+        return super().close()
+
+
+class SparkConnect(SchedulerPlugin):
+    def start(self, scheduler):
+        print("Starting SparkConnect")
+        self.scheduler = scheduler
+        self.scheduler.add_plugin(self)
+        path = pathlib.Path(pyspark.__file__).absolute()
+        module_loc = path.parent
+        script = module_loc / "sbin" / "spark-daemon.sh"
+
+        os.environ["SPARK_HOME"] = str(module_loc)
+        os.environ["PYSPARK_PYTHON"] = sys.executable
+
+        class_ = "org.apache.spark.sql.connect.service.SparkConnectServer"
+        host = scheduler.address.split("//")[1].split(":")[0]
+        spark_master = f"{host}:7077"
+        cmd = f"""{script} submit {class_} 1 \
+            --name "Spark Connect server" \
+            --packages org.apache.spark:spark-connect_2.12:{pyspark.__version__} \
+            --packages org.apache.hadoop:hadoop-aws:{HADOOP_AWS_VERSION} \
+            --packages com.amazonaws:aws-java-sdk-bundle:{AWS_JAVA_SDK_BUNDLE_VERSION} \
+            --master=spark://{spark_master} \
+            --conf spark.driver.host={host} \
+            --conf spark.driver.port=9797 \
+            --conf spark.jars=/tmp/hadoop-aws-{HADOOP_AWS_VERSION}.jar,/tmp/aws-java-sdk-bundle-{AWS_JAVA_SDK_BUNDLE_VERSION}.jar
+            """  # noqa: E501
+        print(f"Executing\n{cmd}")
+        self.proc_connect = subprocess.Popen(shlex.split(cmd))
+
+    def close(self):
+        self.proc_connect.terminate()
+        self.proc_connect.wait()
+        return super().close()
+
+
+class SparkWorker(WorkerPlugin):
+    def setup(self, worker):
+        self.worker = worker
+
+        path = pathlib.Path(pyspark.__file__).absolute()
+        module_loc = path.parent
+        host = worker.scheduler.address.split("//")[1].split(":")[0]
+        spark_master = f"{host}:7077"
+
+        os.environ["SPARK_HOME"] = str(module_loc)
+        os.environ["PYSPARK_PYTHON"] = sys.executable
+        os.environ["PYSPARK_SUBMIT_ARGS"] = make_pyspark_submit_args(spark_master)
+
+        print(f"Launching Spark Worker connecting to {spark_master}")
+        cmd = f"spark-class org.apache.spark.deploy.worker.Worker {spark_master}"
+        self.proc = subprocess.Popen(shlex.split(cmd))
+        print("Launched Spark Worker")
+
+    def teardown(self, worker):
+        self.proc.terminate()
+        self.proc.wait()
+        return super().teardown(worker)
diff --git a/tests/conftest.py b/tests/conftest.py
index cc1f0ad20b..4fabd16a5e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -55,6 +55,11 @@ def pytest_addoption(parser):
         "--benchmark", action="store_true", help="Collect benchmarking data for tests"
     )
     parser.addoption("--run-workflows", action="store_true", help="Run workflow tests")
+    parser.addoption(
+        "--tpch-non-dask",
+        action="store_true",
+        help="Run all (DuckDB / Polars / PySpark) TPCH benchmarks",
+    )
 
 
 def pytest_collection_modifyitems(config, items):
@@ -65,6 +70,15 @@ def pytest_collection_modifyitems(config, items):
         ):
             item.add_marker(skip_workflows)
 
+    skip_benchmarks = pytest.mark.skip(reason="need --tpch-non-dask option to run")
+    for item in items:
+        if not config.getoption("--tpch-non-dask") and not (
+            str(item.path).startswith(
+                str(TEST_DIR / "benchmarks" / "tpch" / "test_dask")
+            )
+        ):
+            item.add_marker(skip_benchmarks)
+
 
 dask.config.set(
     {
@@ -515,6 +529,51 @@ def small_client(
         client.run(lambda: None)
 
 
+@pytest.fixture(scope="module")
+def tpch_pyspark_cluster(
+    request, dask_env_variables, cluster_kwargs, github_cluster_tags
+):
+    from .benchmarks.tpch.test_pyspark import SparkMaster, SparkWorker
+
+    module = os.path.basename(request.fspath).split(".")[0]
+    module = module.replace("test_", "")
+    kwargs = dict(
+        name=f"tpch-{module}-{uuid.uuid4().hex[:8]}",
+        environ=dask_env_variables,
+        tags=github_cluster_tags,
+        **cluster_kwargs["tpch_pyspark"],
+    )
+    dump_cluster_kwargs(kwargs, f"tpch_pyspark.{module}")
+    with Cluster(**kwargs) as cluster:
+        client = cluster.get_client()
+        client.wait_for_workers(cluster_kwargs["tpch_pyspark"]["n_workers"])
+        client.register_plugin(SparkMaster(), name="spark-master")
+        client.register_plugin(SparkWorker())
+
+        yield cluster
+
+
+@pytest.fixture
+def tpch_pyspark_client(
+    request,
+    testrun_uid,
+    tpch_pyspark_cluster,
+    cluster_kwargs,
+    upload_cluster_dump,
+    benchmark_all,
+):
+    test_label = f"{request.node.name}, session_id={testrun_uid}"
+
+    with Client(tpch_pyspark_cluster) as client:
+        with upload_cluster_dump(client):
+            log_on_scheduler(client, f"Finished client setup of {test_label}")
+            with benchmark_all(client):
+                yield client
+            log_on_scheduler(client, f"Starting client teardown of {test_label}")
+        client.restart()
+        client.run(lambda: None)
+
+
 @pytest.fixture(scope="module")
 def tpch_cluster(request, dask_env_variables, cluster_kwargs, github_cluster_tags):
     module = os.path.basename(request.fspath).split(".")[0]
@@ -526,8 +585,9 @@ def tpch_cluster(request, dask_env_variables, cluster_kwargs, github_cluster_tag
         **cluster_kwargs["tpch"],
     )
     dump_cluster_kwargs(kwargs, f"tpch.{module}")
-    with Cluster(**kwargs) as cluster:
-        yield cluster
+    with dask.config.set({"distributed.scheduler.worker-saturation": "inf"}):
+        with Cluster(**kwargs) as cluster:
+            yield cluster
 
 
 @pytest.fixture
diff --git a/tests/workflows/test_pytorch_optuna.py b/tests/workflows/test_pytorch_optuna.py
index c3441bef9f..cb2e207654 100644
--- a/tests/workflows/test_pytorch_optuna.py
+++ b/tests/workflows/test_pytorch_optuna.py
@@ -11,7 +11,6 @@
 optuna = pytest.importorskip("optuna")
 
 
-@pytest.mark.xfail(reason="https://github.com/coiled/benchmarks/issues/969")
 @pytest.mark.client(
     "pytorch_optuna",
     worker_plugin=PipInstall(
@@ -19,16 +18,16 @@
             # FIXME https://github.com/coiled/platform/issues/1249
             #       package_sync doesn't seem to deduce GPU specific installs of
             #       libraries like torch
-            "torch==2.0.0",
+            "torch",
             # FIXME Windows package_sync doesn't like torchvision
-            "torchvision==0.15.1",
+            "torchvision",
             # FIXME https://github.com/boto/botocore/issues/2926
             #       urllib3 v2 removed openssl / ciphers which causes an error in
             #       botocore httpsession
             "urllib3<2.0.0",
         ],
         pip_options=["--force-reinstall"],
-        restart=True,
+        restart_workers=True,
     ),
 )
 def test_hpo(client):