Using PySpark

Python computations support PySpark for processing datasets that are too large to fit into memory. Spark splits large workloads into smaller pieces. This is useful even for datasets that would theoretically fit into memory, since libraries like pandas often copy data during transformations.

Spark inside a data clean room runs in a secure enclave with constrained resources. This differs from a typical Spark deployment in a few important ways, so Spark code should be written with these constraints in mind.

How Spark runs in the enclave

• Single machine, no cluster: Spark runs in local mode (local[N]) on a single confidential-computing VM. Computations aren't distributed across multiple machines.
• PySpark 3.5.x: See the Python computation page for the full list of available libraries.
• Limited memory: The standard Python computation environment provides 64 GB of memory and 8+ CPU cores, depending on cluster load. Around 58 GB is available to the Spark runtime; the rest is reserved for the Python process and system overhead.
• No distributed file system: Shuffle data and disk-persisted DataFrames are written to a local encrypted SSD (/scratch) with around 800 GB of usable space, instead of a distributed file system.

danger

Because there's no cluster to spread the load across, operations that cause a shuffle (groupBy, large joins) can fail if the total data volume is too big. Follow the best practices below when working with large datasets.

Storage locations

Where you read and write data has a big impact on performance:

Path	Speed	Use for
/input	Slow (remote-backed)	Reading input datasets — ideally in a single pass
/output	Slow (remote-backed)	Writing final results — once, in a single pass
/scratch	Fast (encrypted SSD)	Intermediate data, Spark shuffle, and spill files
Any other path	Memory-backed	Avoid — files written here count toward the memory limit

Computations that do a lot of random access on /input or /output will be slow. Copy the data to /scratch first if you need to read it repeatedly, and never update files under /output in place.

note

The /scratch disk is encrypted and wiped after every computation, so it's safe to use for sensitive intermediate data.

Setting up a Spark session

Always create the Spark session with the spark_session() context manager from the built-in decentriq_util library. It configures Spark for the enclave environment automatically: it detects the available memory, directs temporary data to /scratch, and cleans up when the session ends.

import decentriq_util as dq

with dq.spark.spark_session() as spark:
    df = spark.read.csv("/input/my_data.csv", header=True)
    # ... transformations ...
    df.write.parquet("/output/result.parquet")
# The session is stopped and temporary directories are cleaned up

You can override or extend the configuration via the config argument. Custom values take precedence over the auto-detected defaults. Refer to the Spark configuration documentation for the available options.

For example, spark.sql.shuffle.partitions controls how many partitions are created during shuffles (joins, groupBy). It defaults to 200 — consider increasing it for large datasets so individual partitions stay small:

with dq.spark.spark_session(
    config=[
        ("spark.sql.shuffle.partitions", "400"),
    ],
) as spark:
    # Your code here
    pass

A typical PySpark computation

The following example combines the pieces above: it reads two input datasets, joins and deduplicates them, materializes the intermediate result to /scratch as Parquet so both downstream aggregations can reuse it, and writes the final results directly to /output.

import decentriq_util as dq
import pyspark.sql.functions as F

with dq.spark.spark_session() as spark:
    # Read the two input datasets
    customers = spark.read.csv("/input/customers/dataset.csv", header=True)
    transactions = spark.read.csv("/input/transactions/dataset.csv", header=True)

    # Join and remove duplicated rows, then materialize the
    # intermediate result to /scratch as Parquet
    joined = (
        transactions
        .join(customers, on="customer_id", how="inner")
        .dropDuplicates(["transaction_id"])
    )
    joined.write.parquet("/scratch/joined.parquet")

    # Read the intermediate result back and run two aggregations on it
    joined = spark.read.parquet("/scratch/joined.parquet")

    revenue_per_country = joined.groupBy("country").agg(
        F.sum("amount").alias("total_amount")
    )
    transactions_per_customer = joined.groupBy("customer_id").agg(
        F.count("transaction_id").alias("transaction_count")
    )

    # Write the final results directly to /output
    revenue_per_country.write.parquet("/output/revenue_per_country.parquet")
    transactions_per_customer.write.parquet("/output/transactions_per_customer.parquet")

If you only run a single aggregation on the joined data, skip the /scratch step and chain the transformations directly into the final write.

Best practices

Read input data only once

Reading from /input is slow, so avoid scanning it repeatedly (for example, by running two aggregations on the same input DataFrame without materializing it in between). If you need multiple passes over the data, copy it to /scratch as Parquet first. For a single aggregation pass, read /input directly — copying first would only waste scratch space.

Use Parquet for intermediate data

Write intermediate results to /scratch as Parquet. It's columnar, compressed, and lets Spark read only the columns and row groups a downstream step actually needs. Delete intermediate datasets as soon as they're no longer needed to free up scratch space.

Chain transformations instead of caching

Write your pipeline as consecutive transformations ending in a single write. This lets the Spark optimizer prune columns and combine processing steps. Avoid .cache(): cached data isn't reclaimed predictably. If a DataFrame is genuinely reused by multiple steps, write it to /scratch as Parquet, read it back, and delete it when done. Alternatively, you can use .persist(StorageLevel.DISK_ONLY), as long as it's used sparingly — it fills up the scratch disk with data that isn't very highly compressed.

Avoid Python UDFs

Python UDFs move every row between the JVM and the Python runtime, which is slow and memory-intensive. Prefer Spark-native functions (F.sha2, F.regexp_replace, F.expr(...), and so on). If a UDF is unavoidable, apply it only to a small, bounded dataset (for example, the distinct values of a column) and join the result back.

Don't pull large data to the driver

.collect() and .toPandas() load the entire DataFrame into memory and can crash the computation if the data is too large. Use Spark aggregations to reduce the data first, or .limit() before collecting.

Keep shuffles small

Shuffle data is written to /scratch and can fill it up. To reduce shuffle size:

• Select only the columns you need before joins and aggregations.
• Combine multiple aggregations into a single pass where possible.
• Use an explicit broadcast join when one side is small: large_df.join(F.broadcast(small_df), on="key").

When you only need to reduce the partition count, use coalesce(n) instead of repartition(n) — it avoids a shuffle entirely.

Repartition gzip inputs

Gzip-compressed CSV files aren't splittable: Spark reads them as a single partition regardless of size. Call repartition(N) right after reading a .gz source.

Writing a single CSV file

If you need a single CSV file as output, don't use coalesce(1) — it creates one giant partition that can exhaust memory. Instead, write the parts to /scratch and stream-merge them into /output:

Helper function for writing a single CSV file

import os
import shutil
import tempfile

def write_df_as_single_file(df, path, temp_dir="/scratch", header=None):
    """Write df as a single CSV file at `path`."""
    with tempfile.TemporaryDirectory(dir=temp_dir) as d:
        parts_dir = os.path.join(d, "parts")
        df.write.csv(parts_dir, header=False)
        parts = sorted(
            os.path.join(parts_dir, f)
            for f in os.listdir(parts_dir)
            if f.endswith(".csv")
        )
        with open(path, "wb") as out:
            if header:
                out.write((",".join(header) + "\n").encode())
            for p in parts:
                with open(p, "rb") as f:
                    shutil.copyfileobj(f, out)
                os.remove(p)  # free scratch space as we go

Troubleshooting

Error	Cause	Fix
Computation killed without a clear error	Container hit the memory limit	Usually a UDF, .collect(), or .toPandas() on large data
java.io.IOException: No space left on device	/scratch is full	Delete intermediate data promptly, reduce shuffle size
Extreme CPU usage with little progress	Memory pressure near the limit	Reduce parallelism, remove cached data

tip

For smaller datasets, plain Python with pandas is often simpler and more reliable than Spark. Reach for Spark only when the data is genuinely larger than memory or when pandas memory management causes issues.