decentriq_util.spark

Functions

spark_session

def spark_session(
    temp_dir: str = '/scratch',
    input_files: Optional[list[str]] = None,
    name: str = 'local_spark_session',
    config: Optional[list[typing.Tuple[str, typing.Any]]] = None,
    parallelism: int = 8,
)

Create a spark session and configure it according to the enclave environment.

Parameters:

• temp_dir: Where to store temporary data such as persisted data frames, shuffle data and the SQL warehouse. Leave this at its default: /scratch is the fast, encrypted disk that is wiped when the computation ends, and any other location is both slower and not guaranteed to be writable.
• name: An optional name for this spark session.
• config: Extra settings to pass to the Spark session builder. These always take precedence over the settings determined automatically.
• parallelism: The size of the executor pool computing internal spark tasks. This should be left at its default value.
• input_files: Deprecated and ignored.

Example:

import decentriq_util as dq

# Path to a potentially very large file
input_csv_path = "/input/my_file.csv"

# Automatically create and configure a spark session and
# make sure it's being stopped at the end.
with dq.spark.spark_session() as ss:
    # Read from a CSV file
    df = ss.read.csv(input_csv_path, header=False)

    # Perform any pyspark transformations
    print(f"Original number of rows: {df.count()}")
    result_df = df.limit(100)

    # Write the result to an output file
    result_df.write.parquet("/output/my_file.parquet")

spark.sql.shuffle.partitions controls how many partitions a shuffle (a join, a groupBy, a dropDuplicates) produces. It defaults to 200. Each partition is processed as one task, so raising the count makes every task's working set smaller — worth doing on large inputs, where the default leaves partitions big enough to exhaust the memory available to a task:

with dq.spark.spark_session(
    config=[("spark.sql.shuffle.partitions", "1000")],
) as ss:
    ...