PySpark tutorial: getting started — a step-by-step guide

Getting started with PySpark

If you've been working with pandas on single machines and hit the ceiling on data size, or you're building data pipelines that need to scale across a cluster, PySpark is where you move next. This PySpark tutorial: getting started — a step-by-step guide walks you through the essentials so you can run Spark jobs from Python without wrestling with Scala or Java.

We'll cover installation, session creation, and the difference between notebook exploration and production-grade scripts. By the end, you'll have a working understanding of how to read data, transform it, and submit jobs that process gigabytes or terabytes.

Why PySpark matters in your data stack

Spark runs distributed compute across multiple machines. When you have a CSV with millions of rows or need to join datasets that won't fit in RAM, Spark partitions the work across your cluster and brings results back to you. PySpark is the Python interface to Spark's core engine, so you write Python code but get the performance benefits of distributed processing.

The alternative—scaling pandas or SQL alone—often means rewriting jobs or hitting performance walls. With PySpark, you write once and scale from a laptop to a hundred-node cluster without changing the core logic.

Install PySpark via pip for local development

The easiest way to get started is pip install pyspark. This installs PySpark and a bundled version of Spark so you can develop and test locally before deploying to a real cluster.

pip install pyspark

After installation, verify it works:

python -c "import pyspark; print(pyspark.__version__)"

You should see a version number (typically 3.4.x or later at the time of writing). This confirms Spark is installed and ready to use.

For production environments, you'll often run against a Spark cluster managed by your platform (Databricks, EMR, local Kubernetes, etc.), but the pip install approach is perfect for development and testing on your laptop. It includes Java and Scala libraries, so you don't need to install those separately.

Create a SparkSession—your entry point

Every Spark application starts with a SparkSession. Think of it as your connection to Spark: it manages the distributed context, SQL execution, and DataFrame operations. You create one like this:

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("my-first-app") \
    .getOrCreate()

The builder pattern lets you configure options. appName is required and visible in the Spark UI (usually on http://localhost:4040 for local runs). getOrCreate() returns an existing session if one exists, or creates a new one.

Once you have a session, you read data:

df = spark.read.parquet("s3://bucket/orders/")

Here we're reading Parquet files from S3. Spark handles the distributed download and parsing. The result is a DataFrame—Spark's distributed table abstraction, similar to a pandas DataFrame but spread across multiple nodes.

Now you can query it:

df.groupBy("country").count().show()

This groups orders by country and counts them. show() prints the first 20 rows to console. Under the hood, Spark parallelizes the grouping and aggregation across your cluster (or just your local machine in development).

Configuring your session for common scenarios

For more control, you can add configuration options:

spark = SparkSession.builder \
    .appName("etl-job") \
    .config("spark.sql.shuffle.partitions", "200") \
    .config("spark.driver.memory", "4g") \
    .config("spark.executor.memory", "8g") \
    .getOrCreate()

shuffle.partitions affects how Spark repartitions data during group-by or join operations. More partitions mean more parallelism but also more overhead. driver.memory and executor.memory control heap size. For local development, these defaults are often fine, but in production you tune them based on data size and available hardware.

Use notebooks for exploration and experimentation

Jupyter notebooks and Databricks notebooks are the go-to for data exploration. They let you write Python, run cells interactively, and see results immediately. This is where you prototype, test ideas, and debug transformations before moving to production scripts.

Setting up Jupyter locally:

pip install jupyter

Then start a notebook server:

jupyter notebook

In a new notebook cell, create your session and load data:

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("exploration").getOrCreate()
df = spark.read.csv("data/users.csv", header=True, inferSchema=True)
df.printSchema()
df.show(5)

Run the cell and you see the schema and first 5 rows. This iterative feedback loop is invaluable. You can add filters, aggregations, or joins in the next cell and immediately see the results.

Databricks notebooks for team work

If your team uses Databricks (a managed Spark platform), you get notebooks with built-in Spark clusters, versioning, and collaboration. The code is almost identical—your SparkSession is already created for you—but you don't manage infrastructure. In a Databricks cell, you can write PySpark directly without setting up a session yourself.

Notebooks are not for production jobs, though. They're meant for exploration, reporting, and ad-hoc analysis. When you need a reliable, scheduled job, move to a script.

Production scripts with spark-submit

Once your logic is solid, wrap it in a Python script and submit it with spark-submit. This is how you run Spark jobs in production: scheduled daily, triggered by pipelines, or launched by orchestration tools like Airflow or Dagster.

Create a file called process_orders.py:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, count

def main():
    spark = SparkSession.builder.appName("order-processor").getOrCreate()
    
    # Read orders
    df = spark.read.parquet("s3://bucket/orders/")
    
    # Filter and aggregate
    result = df.filter(col("amount") > 100) \
        .groupBy("country") \
        .count() \
        .orderBy(col("count").desc())
    
    # Write results
    result.write.mode("overwrite").parquet("s3://bucket/output/order_counts/")
    
    spark.stop()

if __name__ == "__main__":
    main()

Submit it to Spark:

spark-submit process_orders.py

On a local machine, this runs immediately. On a cluster, spark-submit contacts the cluster manager (YARN, Kubernetes, Standalone), allocates executors, and runs your job. You can pass configuration flags:

spark-submit \
    --master spark://cluster:7077 \
    --executor-memory 8g \
    --num-executors 10 \
    process_orders.py

This tells Spark to use 10 executors with 8GB memory each on your Spark cluster. For AWS EMR or Databricks, configuration looks different but the principle is the same: you define resources and let the cluster manager handle distribution.

Structuring jobs for maintainability

As jobs grow, keep them organized. A common pattern:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

def read_data(spark, path):
    return spark.read.parquet(path)

def transform(df):
    return df.filter(col("amount") > 100) \
        .groupBy("country") \
        .count()

def write_data(df, path):
    df.write.mode("overwrite").parquet(path)

def main():
    spark = SparkSession.builder.appName("etl").getOrCreate()
    
    df = read_data(spark, "s3://bucket/orders/")
    result = transform(df)
    write_data(result, "s3://bucket/output/")
    
    spark.stop()

if __name__ == "__main__":
    main()

Functions are testable, reusable, and easy to understand. This scales better than dumping everything into one script.

Common mistakes when starting with PySpark

• Forgetting to call spark.stop(): In production scripts, always call spark.stop() at the end. This releases cluster resources and lets the driver JVM shut down gracefully. Notebooks handle this for you, but scripts need it.
• Using Pandas methods on Spark DataFrames: Spark DataFrames look similar to Pandas but have different APIs. You can't call .df.iloc[0] or .df.apply() directly. If you need Pandas functionality, collect the data first: pdf = df.toPandas() (only safe for small results).
• Not partitioning output correctly: When writing large datasets, Spark tries to write one file per partition. If you have 200 partitions and write Parquet, you get 200 files. This is usually good for scalability but can surprise teams expecting a single file. Use .coalesce(1) only if you really need one output file.
• Ignoring the Spark UI during development: Open http://localhost:4040 while your job runs. You see task distribution, shuffle sizes, and bottlenecks. It's invaluable for optimization.
• Running complex transformations in the driver: If you use Python loops or custom functions without broadcasting or UDFs, you serialize everything back to the driver. This defeats the purpose of distributed computing. Use Spark's built-in functions (groupBy, join, filter, etc.) whenever possible.

Key takeaways

• Install PySpark with pip install pyspark for local development; it includes a bundled Spark instance.
• Every Spark application starts with a SparkSession created via SparkSession.builder.appName(...).getOrCreate().
• Use Jupyter or Databricks notebooks for exploration and prototyping; they offer interactive feedback and are ideal for data exploration.
• Move to spark-submit scripts for production jobs; structure them with functions and always call spark.stop() at the end.
• Configure memory, partitions, and executors via .config() or command-line flags depending on your cluster setup.
• Check the Spark UI at http://localhost:4040 to understand task distribution and identify bottlenecks.
• Avoid serializing logic back to the driver; use Spark's native functions (groupBy, join, filter) for distributed work.