PySpark

Overview

PySpark is the Python interface to Apache Spark. It lets you write distributed data pipelines in Python using DataFrames, SQL and the MLlib machine learning library. You write familiar-looking Python, and the work runs in parallel across a cluster. It is the most popular way to use Apache Spark because Python is the common language of data science.

PySpark does not run your Python loops on a cluster directly. It builds a plan of operations, which Spark optimises and executes on the JVM. This is why DataFrame code is fast, while custom Python functions can be slow. Knowing that split is the key to using PySpark well. For tabular AI work, PySpark is the bridge between large lake tables, such as Delta Lake tables, and Python machine learning libraries and models.

What it is

PySpark is a Python library that ships with Spark. It exposes the main parts of the engine:

• pyspark.sql: SparkSession, DataFrame, Column, functions, types and window specifications.
• pyspark.ml: Pipelines, transformers, estimators and common algorithms.
• pyspark.streaming and Structured Streaming: Stream processing through the DataFrame API. See Spark Streaming.
• pyspark.pandas: A pandas-like API that runs on Spark.
• pyspark.rdd: The low-level RDD API, rarely needed now.

An analogy helps. Imagine you are a manager who speaks one language and the factory floor speaks another. PySpark is the translator at your desk. You give instructions in Python, the translator converts them into the factory's language, and the factory runs them at scale. Instructions that the translator understands well, such as built-in functions, are carried out directly. Instructions that need the factory to call back to your desk for every item, such as a plain Python function on each row, are slow because every item must be handed back and forth.

How Python and the JVM talk

When you call a DataFrame method, PySpark sends a command to the Spark JVM process. The JVM builds and optimises the plan. Data processed by built-in functions never needs to enter Python. Only when you use Python UDFs, or collect results to the driver, does data move between the JVM and Python worker processes. Efficient transfer formats based on Apache Arrow reduce this cost.

Local mode and cluster mode

In local mode, PySpark runs everything on your machine, which is great for learning and tests. In cluster mode, the same code runs on many machines through a cluster manager. Spark Connect, available in recent versions, lets a thin Python client talk to a remote Spark server, which separates the client from the cluster.

Where it came from

Spark started with Scala as its main language, but data scientists mostly used Python. A Python API appeared early in Spark's life, and it grew steadily as Python became the default for data work. At first, PySpark focused on RDDs and was seen as slower than Scala. The arrival of DataFrames changed this, because DataFrame operations are compiled to the same plan whichever language you use, so Python and Scala perform alike for built-in operations.

Later improvements focused on the Python boundary: Arrow-based transfer, vectorised pandas UDFs, the pandas API on Spark, better type hints and error messages, and a client-server mode. The project has steadily shifted from "Python is a second-class API" to "Python is the primary way to use Spark".

The exact timeline and feature set depend on your Spark version, so check the documentation for the one you run.

How it works

Step 1: Create a session

Everything starts with a SparkSession. It is the entry point to read data, run SQL and configure the engine.

Step 2: Load data into a DataFrame

A DataFrame is a distributed table. Reading creates a DataFrame, but little work happens yet, because Spark only needs file listings and schemas at that point.

Step 3: Build transformations

You chain methods such as select, filter, withColumn, join and groupBy. Each call returns a new DataFrame that stores a plan. This is lazy evaluation, like writing a recipe before cooking.

Step 4: Spark optimises the plan

The Catalyst optimiser rewrites the plan: it pushes filters down to the data source, removes unused columns, simplifies expressions and picks join strategies. Your Python code only described the intent, so Spark is free to improve it.

Step 5: An action triggers execution

Calling show, count, collect, toPandas or write executes the plan. The driver breaks it into stages and tasks and sends them to executors.

Step 6: Executors process partitions

Each task works on one partition of data. For built-in operations, executors run optimised JVM code. For Python UDFs, the executor starts Python worker processes and passes data to them in batches.

Step 7: Results come back or are written

Results are written to storage in parallel, or small results come back to the driver. toPandas pulls the entire result into driver memory, so use it only on small data.

Why built-in functions beat UDFs

A built-in function is part of the plan and runs inside the engine. A regular Python UDF is a black box. The optimiser cannot look inside it, and data must be serialised to Python and back. A vectorised pandas UDF moves data in Arrow batches and works on whole columns at once, which is much faster than row-by-row UDFs. The best order of preference is: built-in function, then SQL expression, then pandas UDF, then plain Python UDF.

Key concepts and terms

A worked example

This example shows a small but realistic PySpark job: read, clean, feature-engineer, and write. Paths are fake.

Start a session and read

from pyspark.sql import SparkSession, functions as F, types as T

spark = SparkSession.builder.appName("pyspark_demo").getOrCreate()

schema = T.StructType([
    T.StructField("order_id", T.LongType(), False),
    T.StructField("customer_id", T.LongType(), False),
    T.StructField("order_ts", T.TimestampType(), True),
    T.StructField("quantity", T.IntegerType(), True),
    T.StructField("unit_price", T.DoubleType(), True),
    T.StructField("country", T.StringType(), True),
])

orders = (
    spark.read.schema(schema)
    .option("header", True)
    .csv("s3://example.internal/lake/bronze/orders/")
)

Giving an explicit schema avoids a slow, error-prone inference pass.

Clean

clean = (
    orders
    .dropDuplicates(["order_id"])
    .filter(F.col("quantity") > 0)
    .withColumn("country", F.upper(F.trim("country")))
    .withColumn("revenue", F.col("quantity") * F.col("unit_price"))
    .withColumn("order_date", F.to_date("order_ts"))
)

Missing values

profile = clean.select(
    *[F.sum(F.col(c).isNull().cast("int")).alias(c) for c in clean.columns]
)
profile.show()

filled = (
    clean
    .withColumn("unit_price_missing", F.col("unit_price").isNull().cast("int"))
    .fillna({"country": "UNKNOWN"})
)

The profile counts nulls per column in one pass. The indicator column keeps the fact that a price was missing, instead of hiding it.

Features per customer

from pyspark.sql import Window

features = (
    filled.groupBy("customer_id")
    .agg(
        F.countDistinct("order_id").alias("order_count"),
        F.sum("revenue").alias("total_spend"),
        F.max("order_ts").alias("last_order_ts"),
        F.avg("unit_price").alias("avg_unit_price"),
    )
    .withColumn(
        "days_since_last_order",
        F.datediff(F.current_date(), F.to_date("last_order_ts")),
    )
)

SQL and DataFrames together

features.createOrReplaceTempView("customer_features")

top = spark.sql("""
    SELECT customer_id, total_spend
    FROM customer_features
    WHERE order_count >= 3
    ORDER BY total_spend DESC
    LIMIT 100
""")
top.show()

A vectorised pandas UDF

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def log_spend(s: pd.Series) -> pd.Series:
    import numpy as np
    return np.log1p(s.fillna(0.0))

features2 = features.withColumn("log_spend", log_spend("total_spend"))

In practice, F.log1p is built in and would be faster. The pandas UDF is shown to illustrate the pattern for logic that has no built-in equivalent.

Write

(features2.write
    .mode("overwrite")
    .format("delta")
    .save("s3://example.internal/lake/gold/customer_features"))

A tiny MLlib pipeline

from pyspark.ml import Pipeline
from pyspark.ml.feature import VectorAssembler
from pyspark.ml.classification import LogisticRegression

labelled = features2.join(labels, "customer_id")   # labels: customer_id, churned

assembler = VectorAssembler(
    inputCols=["order_count", "total_spend", "days_since_last_order"],
    outputCol="features",
    handleInvalid="keep",
)
lr = LogisticRegression(featuresCol="features", labelCol="churned")
model = Pipeline(stages=[assembler, lr]).fit(labelled)

The labels DataFrame is assumed to exist. This pattern trains a simple baseline on the cluster.

A numeric picture

Say features has 2 million customers and the data is 3 GB after aggregation. With 128 MB partitions you get about 24 partitions. A toPandas() on this table would bring 3 GB into the driver, which is risky. Instead, take features.sample(0.01) for about 20,000 rows, convert that, and send it to a local model or notebook.

Strengths and limitations

Strengths

• Python first. Works with the Python data science ecosystem.
• Same performance for DataFrame code. Built-in operations run in the JVM like Scala.
• Scale without rewriting. The same code runs locally and on a cluster.
• SQL and Python together. Mix them freely in one job.
• Rich built-ins. Hundreds of functions for strings, dates, arrays, maps and windows.
• MLlib pipelines. Reusable feature and model steps.
• Pandas interoperability. Arrow-based conversion and pandas UDFs.
• Active development. Python support is a priority in recent Spark versions.

Limitations

• UDF overhead. Row-by-row Python is slow.
• Debugging across processes. Errors can come from the JVM, Python workers or the cluster, and stack traces can be long.
• Driver memory limits. collect and toPandas can overload the driver.
• Dependency management. Python packages must be available on all workers.
• Version coupling. PySpark, Spark and Python versions must be compatible.
• Limited GPU story. Heavy deep learning is usually done outside core Spark.
• Overkill for small data. Local tools are simpler.

How it compares with alternatives

A common pattern is to use PySpark for heavy joins and aggregations, then hand a compact result to pandas or a model library for the last mile.

Role in tabular foundation models

Tabular foundation models are pre-trained to understand tables: rows as examples, columns as features, with mixed types and missing values. They can often predict from a small set of labelled rows given in context, without training. PySpark is how teams get large real-world tables into that shape, and how they run the results back across millions of rows.

Building model-ready tables

PySpark builds the single flat table a model wants. It joins sources, deduplicates, fixes types and creates features. Typical outputs are layered tables such as those in the medallion architecture, with the final result in the gold layer. Tabular models like TabPFN, TabICL and TabSTAR can then use a sample.

Row and column structure

A PySpark DataFrame has rows and typed columns, which is the same abstraction models use. Converting a sample to pandas or NumPy takes one call, after you have filtered and sampled. Because models have limits on context size, PySpark is used for stratified sampling, nearest-neighbour style context selection with joins, and for splitting by time or by entity to avoid leakage.

Mixed types and missing values

PySpark has typed columns and keeps nulls explicit. Profiling nulls across hundreds of columns takes a few lines. You can add indicator columns, fill with domain defaults, or leave the nulls for the model to handle. For imputation, a typical job reads rows with missing cells, builds model inputs, calls a model in batches through a pandas UDF, and writes imputed values to a new column or table, keeping the original untouched.

No column names

To test whether a model depends on names, PySpark can rename all columns to generic labels with toDF(*new_names) and keep a mapping table. That also helps when a data source has meaningless names and you want to compare how a model performs with and without semantic headers.

Pre-training data pipelines

Pre-training uses many tables, often a mix of real and synthetic. PySpark can ingest large collections, apply the same cleaning to each, assign table ids, create train and validation splits at the table level, and store the result in Parquet. For synthetic data, a job can generate many tables in parallel with different random seeds. Pinning inputs to versions of Delta Lake tables makes corpora reproducible. See model research and model architecture for how such training is described.

In-context learning at scale

To score many rows with a model that takes context, you give each batch a context set. In PySpark, you can broadcast a small context DataFrame, collect it into local arrays on each worker, and score batches inside a pandas UDF or mapInPandas. This keeps data movement small and parallelises the scoring step.

def predict_batches(batches):
    model = load_model()                   # placeholder loader, once per task
    context_x, context_y = get_context()  # placeholder, e.g. from a broadcast
    for pdf in batches:
        yield pdf.assign(pred=model.predict(context_x, context_y, pdf))

scored = features.mapInPandas(predict_batches, schema=features.schema.add("pred", "double"))

The helper functions are placeholders for your own code, and the output schema must match what you yield.

Continual learning

For continual updates, a scheduled PySpark job or a streaming job reads only new rows, for example through a table change feed, and refreshes features and the context used by the model. Recording the processed table version makes each run restartable. See Spark Streaming for the streaming side.

Anomaly detection

PySpark computes per-group statistics, window features and model scores in one job. A simple pattern is to compute rolling means and standard deviations with window functions, then flag rows far from them, and compare with scores from a model. See anomaly detection. Results can be aggregated for dashboards.

Sector and semantic classification

PySpark can prepare features for classifying records into sectors, run a classifier in batches, and store labels with confidence scores. Evaluation joins predictions with ground truth and computes per-class metrics with groupBy. For ideas about measuring such systems, see the model benchmarks page.

Limits to keep in mind

Foundation models often need GPUs and specific libraries, while PySpark workers may be CPU-only. Teams then run the model service separately and call it in batches, or use GPU-enabled clusters. For small tables the whole cluster is unnecessary. For classic gradient boosting, XGBoost offers distributed training that works with Spark DataFrames.

Practical guidance and best practices

Code style

• Prefer built-in functions in pyspark.sql.functions.
• Import functions as F and types as T for readability.
• Chain transformations in clear steps, and name intermediate DataFrames.
• Put logic in functions that take and return DataFrames.
• Avoid hidden state, so functions are easy to test.
• Use column expressions instead of string SQL snippets when you can, for better errors.

Performance

• Select the columns you need as early as possible.
• Filter early and filter on partition columns.
• Broadcast small tables with F.broadcast.
• Avoid collect, and limit toPandas to small, sampled data.
• Use pandas UDFs, not plain UDFs, when you must write Python logic.
• Cache only what you reuse, and unpersist afterwards.
• Look at df.explain() and the Spark UI before tuning.

Schemas and types

• Define schemas explicitly for CSV and JSON.
• Use decimal for money.
• Store timestamps in UTC.
• Keep column names lowercase with underscores.
• Cast early and check with printSchema().

Data quality

• Count nulls and duplicates after each major step.
• Assert expected row counts and key uniqueness.
• Write rejected rows to a quarantine table with a reason.
• Use constraints from your table format where available.

Environments

• Pin Python, PySpark and Spark versions.
• Package dependencies for workers, not just the driver.
• Keep configuration out of code.
• Use secrets managers for credentials, never hard-coded keys.
• Test in local mode with small fixtures.

Orchestration

• Make jobs idempotent and parameterised by date.
• Schedule with Apache Airflow or a similar tool.
• Log row counts and durations for each step.
• Alert on failures and on large changes in output size.

Working with nested and semi-structured data

Real event data often has nested fields and arrays. PySpark can reach inside them without leaving the engine.

events = spark.read.json("s3://example.internal/lake/bronze/events/")

flat = (
    events
    .select(
        "event_id",
        F.col("user.id").alias("user_id"),
        F.col("user.country").alias("country"),
        F.explode("items").alias("item"),
    )
    .select("event_id", "user_id", "country",
            F.col("item.sku").alias("sku"),
            F.col("item.price").alias("price"))
)

explode turns each array element into its own row. Flattening nested data into plain columns is usually the first step toward a table a tabular model can read, since such models expect one value per cell.

Joins in practice

Anti joins are especially handy for data quality, for example finding orders whose customer is missing from the customer table.

Reading the plan

df.explain() prints the physical plan. A few things to look for:

• FileScan lines show pushed filters and the columns read. If a filter is not pushed down, check the column type and expression.
• BroadcastHashJoin is good for a small side. SortMergeJoin means both sides are shuffled and sorted.
• Exchange marks a shuffle. Count them and ask whether each is needed.
• HashAggregate appears twice in many plans, a partial and a final step, which reduces shuffled data.
• BatchEvalPython or ArrowEvalPython marks Python UDF work, a hint that a built-in may be faster.

A debugging routine

When a job fails or gives a wrong result, a calm routine saves time.

1. Reproduce on a small sample in local mode.
2. Print the schema at each step and check types.
3. Count rows before and after each join and filter.
4. Look for null keys and duplicate keys, which cause surprise row counts.
5. Read the first error in the stack trace, not the last.
6. Check the Spark UI for failed tasks and their executors.
7. Fix the cause, then add a small test so it cannot return.

Configuration worth knowing

Always change one setting at a time and measure.

Common mistakes

1. Calling toPandas() on large data. The driver runs out of memory.
2. Using row-by-row UDFs for simple logic. A built-in function is faster.
3. Forgetting laziness. Timing or error checks around a transformation measure nothing until an action runs.
4. Reusing a DataFrame that triggers repeated expensive work. Without caching, each action recomputes the plan.
5. Relying on schema inference. It is slow and can guess wrong types.
6. Joining on columns with different types. Silent casts cause missed matches.
7. Ignoring nulls in joins and comparisons. Null never equals null in normal comparisons.
8. Too many small output files. Repartition or compact before writing.
9. Mismatched Python versions. Driver and workers must agree.
10. Hard-coding credentials. Use roles or secret stores.
11. Using Python loops over DataFrame rows. Express logic as column operations instead.
12. Not setting a random seed for samples. Results change on each run.

Frequently asked questions

What is PySpark?

PySpark is the Python API for Apache Spark. It lets you use Spark's distributed engine from Python through DataFrames, SQL, streaming and the MLlib library. You write Python code, but the heavy work runs in parallel on a cluster, so you can process data far larger than a single machine can hold.

Is PySpark slower than Scala Spark?

For DataFrame and SQL code, the speed is almost the same, because both languages build the same plan that runs in the JVM. The difference appears with Python UDFs, which move data between the JVM and Python. Using built-in functions or vectorised pandas UDFs keeps PySpark close to Scala performance.

What is the difference between PySpark and pandas?

Pandas runs on one machine and keeps all data in memory, while PySpark spreads data across many machines and runs lazily with an optimiser. Pandas is simpler for small data. PySpark is better when data is large or when a pipeline must scale. The pandas API on Spark offers pandas-like syntax on top of Spark.

When should I use a pandas UDF?

Use a pandas UDF when you need Python or library logic that has no built-in Spark equivalent, such as applying a trained model or a specialised calculation. It processes data in batches using Arrow, which is far faster than a plain Python UDF that handles one row at a time.

Why is my PySpark job slow?

Common causes are skewed joins, too many shuffles, row-by-row UDFs, many small files, missing filters and wrong partition counts. Start by reading the query plan and the Spark UI to find the slow stage, then fix the query before changing cluster size. Most large speed-ups come from better code.

How do I convert between Spark and pandas DataFrames?

Use toPandas() to bring a Spark DataFrame to pandas and spark.createDataFrame(pandas_df) to go the other way. Enabling Arrow makes the transfer faster. Only call toPandas() on small results, such as a filtered or sampled table, since all rows load into driver memory.

Can I use PySpark for machine learning?

Yes. PySpark includes MLlib for distributed algorithms and pipelines, and it can prepare data for other libraries. Many teams use PySpark for feature engineering and batch scoring, and use other tools for deep learning. Tabular foundation models can be applied in batches through pandas UDFs or mapInPandas.

How do I handle missing values in PySpark?

You can count nulls with isNull, fill them with fillna, drop rows with dropna, or keep them and add indicator columns. Choose based on meaning: a missing value may be unknown, not applicable or truly zero. For model work, it is often best to keep the null and add a flag, or to use a model-based imputation step.

What is Spark Connect?

Spark Connect is a client-server architecture that lets a lightweight client, such as a Python process, talk to a remote Spark server over a protocol. It decouples the client application from the cluster driver, which simplifies upgrades and allows thin clients in notebooks and applications.

Do I need a cluster to learn PySpark?

No. Local mode runs on a laptop and is enough to learn the API with small datasets. You only need a cluster when the data or the workload outgrows one machine. Code written in local mode usually moves to a cluster with little change.

Glossary

• Action: A call that triggers Spark to run the planned work.
• Arrow: A columnar in-memory format used to move data efficiently between processes.
• Broadcast: Sending a small dataset to every executor for use in a join.
• Column expression: An object describing a value computed from columns.
• DataFrame: A distributed table with named, typed columns.
• Driver: The process that runs your main program and plans the job.
• Executor: A worker that runs tasks on partitions.
• Lazy evaluation: Delaying execution until a result is requested.
• mapInPandas: A method that applies a function to batches as pandas DataFrames.
• MLlib: Spark's machine learning library.
• Pandas UDF: A vectorised Python function working on pandas batches.
• Partition: A piece of a dataset processed by one task.
• Pipeline: A chain of feature and model stages in MLlib.
• SparkSession: The entry point for working with Spark.
• Spark Connect: A client-server mode for remote Spark access.
• Temp view: A named reference to a DataFrame usable in SQL.
• Transformation: A lazy operation that returns a new DataFrame.
• UDF: A user-defined function.
• Window function: A calculation over a related set of rows.

Summary

PySpark brings the power of Apache Spark to Python. You describe data work with DataFrames and SQL, Spark optimises and runs it across a cluster, and results are written back to storage or handed to Python tools. The key habits are to use built-in functions, avoid row-by-row UDFs, keep large data out of the driver, and read query plans before tuning.

For tabular foundation models, PySpark prepares model-ready tables, handles missing values, selects context samples, scales imputation, anomaly scoring and classification, and supports reproducible pipelines on versioned Delta Lake tables. It is not where large neural models run best, but it is an excellent way to feed them and to apply their predictions across very large tables.

Related

It belongs to Batch and streaming processing in my stack, under Data engineering and lakehouse. Others in the same group: