sparktimise 2.0.0

A Python wrapper that analyses DataFrames and applies optimisation techniques to maximise PySpark session performance.

sparktimise

A PySpark optimisation library that inspects DataFrames and applies targeted performance improvements with minimal user code changes.

What The Package Does

sparktimise provides two ways to optimise PySpark jobs:

1. A pipeline context-manager workflow through optimise / SparkPipelineAutoTuner.
2. A functional workflow through analyse_* and optimise_* functions for explicit control.

Capability	Description	Primary API
Partition optimisation	Estimates optimal shuffle partitions and low-cardinality partition candidates	optimise_partitions
Skew mitigation	Detects skewed keys and applies salting columns	optimise_skew
Cache strategy	Recommends and applies StorageLevel persistence	optimise_cache
Spark session tuning	Recommends and applies Spark SQL/session settings	SparkPipelineAutoTuner / optimise_context
Broadcast analysis	Profiles table sizes for join strategy advice and optional hints	analyse_broadcast / apply_broadcast_hints
Reporting	Summarises pipeline steps and metadata in text or dict form	OptimisationReport

Installation

pip install sparktimise

For local development:

pip install -e .[dev]
pip install pyspark

Quick Start

Context-manager usage

from pyspark.sql import SparkSession
from sparktimise import SparkPipelineAutoTuner

spark = SparkSession.builder.appName("orders-job").getOrCreate()


def run_pipeline():
    orders = spark.read.parquet("s3a://my-bucket/orders/")
    return orders.groupBy("customer_id").count()


with SparkPipelineAutoTuner(
    spark=spark,
    pipeline_name="orders_pipeline",
    watched_modules=["my_project.orders"],
) as tuner:
    tuner.execute("run_orders", run_pipeline)

Entry-point usage through optimise

from pyspark.sql import SparkSession
from sparktimise import optimise

spark = SparkSession.builder.appName("orders-job").getOrCreate()

with optimise(
    spark,
    "orders_pipeline",
    run_type="optimise",
    watched_modules=["my_project.orders"],
) as tuner:
    tuner.execute("run_orders", run_pipeline)

Functional usage

from sparktimise.optimisation import optimise_partitions, optimise_skew

step1 = optimise_partitions(df, target_partition_bytes=134_217_728)
step2 = optimise_skew(step1.df, columns=["customer_id"])

optimised_df = step2.df
print(step1.transformations_applied)
print(step2.transformations_applied)

Configuration And Runtime Controls

The primary runtime configuration surface is the optimise context manager.

Parameter	Type	Default	Effect
run_type	str	"optimise"	"optimise", "baseline", or "report"
watched_functions	list[str] | None	None	Exact or wildcard qualified function names to auto-assess
watched_modules	list[str] | None	None	Module prefixes to auto-assess
auto_capture	bool	True	Enables context-manager function-return capture
include_plan	bool	False	Stores full plan text in assessments
run_id	str | None	None	Groups results under a named folder; defaults to pipeline name
results_root	str | None	None	Root folder to write sparktimise_results under
spark	SparkSession	Required	Session used for safe SQL/session tuning

For full configuration details, including file-backed config loading and Spark recommendation settings, see docs/configuration.md.

Architecture And Process Flow

sparktimise follows a hybrid pattern:

1. Functional core: analyser and optimiser functions.
2. Imperative shell: context-manager orchestration.
3. OOP boundaries: adapters, config, and reporting.

Detailed architecture and sequence diagrams are documented in docs/architecture.md.

Documentation Map

Document	Purpose
docs/README.md	Documentation index
docs/usage.md	End-to-end usage patterns and examples
docs/configuration.md	Configuration variables and file formats
docs/architecture.md	Internal design and process flow
docs/troubleshooting.md	Common setup/runtime/CI issues and fixes
CHANGELOG.md	Versioned release and change history

Development Setup

Requirements

Dependency	Purpose
Python 3.10+	Runtime and tooling
Java (JRE/JDK)	Required for Spark tests
PySpark	Runtime dependency for Spark operations

Local setup

python -m pip install -U pip
python -m pip install -e .[dev]
python -m pip install pyspark

Quality checks

ruff check src/sparktimise/ tests/
ruff format --check src/sparktimise/ tests/
mypy src/sparktimise/

Tests

# Unit tests (default)
python -m pytest tests/unit/ --tb=short -q

# Integration tests (requires Java + Spark)
python -m pytest tests/integration/ --run-spark --spark-smoke-timeout 60 --tb=short -q

Build package artifacts

python -m pip install build
python -m build --sdist --wheel

Artifacts are created under dist/.

License

MIT