pyspark

python/pyspark/ is the Python binding to Spark. It covers RDDs, DataFrames, SQL, MLlib, the pandas API on Spark, Structured Streaming, the Connect client, declarative pipelines, and a test framework. PySpark is the most-downloaded Spark surface (over 30 million PyPI downloads per month).

Purpose

• Expose every Spark API to Python.
• Run user Python code (UDFs, UDAFs, Pandas UDFs, applyInPandas, Arrow conversions) inside worker processes spawned by executors.
• Provide pandas API parity (the pandas-on-Spark layer) for users coming from pandas.
• Support both classic mode (Py4J -> JVM) and Connect mode (gRPC -> server-side JVM).

Directory layout

python/
  pyspark/
    __init__.py
    sql/                   - DataFrame, SparkSession, types, functions, observation
      connect/             - Connect client mirror of the SQL API
      session.py, dataframe.py, column.py, functions.py, types.py, ...
    ml/                    - DataFrame-based ML (mirror of mllib/.../ml)
      connect/             - Connect ML client
    mllib/                 - Legacy RDD-based ML (mirror of mllib/.../mllib)
    streaming/             - DStream API (legacy)
    pandas/                - The pandas API on Spark (~100k LoC)
    pipelines/             - Declarative pipelines client
    resource/              - ResourceProfile, ExecutorResources
    errors/                - Pythonic error classes mapped from JVM exceptions
    testing/               - PySpark test utilities
    tests/                 - The test suite
    cloudpickle/           - Vendored cloudpickle
    serializers.py         - Pickle, Arrow, batched serializers
    worker.py              - The Python worker entry point (~150 KB)
    daemon.py              - The Python daemon that spawns worker processes
    java_gateway.py        - Py4J gateway used in classic mode
    profiler.py            - cProfile and memory_profiler integration
  run-tests.py             - Test runner used by dev/run-tests
  setup.py, MANIFEST.in    - Packaging metadata
  packaging/               - Metadata used by sdist/wheel builds
  benchmarks/              - PySpark benchmarks
  docs/                    - Sphinx documentation source
  test_support/            - Test fixtures

Key abstractions

Type / file	What it is
pyspark.SparkContext (python/pyspark/core/context.py)	The classic-mode driver entry point. Wraps a JVM SparkContext via Py4J.
pyspark.sql.SparkSession (python/pyspark/sql/session.py)	DataFrame entry point. Routes to classic or Connect mode based on the remote URL.
pyspark.sql.DataFrame (python/pyspark/sql/dataframe.py)	The DataFrame API.
pyspark.sql.connect.session.SparkSession	The Connect-mode entry point.
pyspark.sql.connect.dataframe.DataFrame	The Connect DataFrame mirror.
pyspark.sql.types	Python types matching Catalyst types (StringType, ...).
pyspark.pandas.frame.DataFrame	pandas-API DataFrame backed by Spark.
pyspark.ml.Pipeline	Mirror of the Scala ml.Pipeline.
pyspark.errors.exceptions.captured.CapturedException	Wraps JVM exceptions raised by Py4J.
pyspark.worker (python/pyspark/worker.py)	Python worker process - runs UDFs, applyInPandas, mapInArrow.
pyspark.daemon (python/pyspark/daemon.py)	Forks a pool of worker processes per executor.

Two execution modes

graph LR
    subgraph Classic mode
        Py1[Python driver] -- Py4J --> JVM1[JVM SparkContext]
        JVM1 --> Engine1[Spark engine]
    end

    subgraph Connect mode
        Py2[Python driver] -- gRPC --> Server[Connect server]
        Server -- in-process --> Engine2[Spark engine]
    end

pyspark.sql.session.SparkSession.builder decides at runtime:

• If remote("sc://host") was used or SPARK_REMOTE is set, instantiate the Connect-mode session in pyspark.sql.connect.session.
• Otherwise, start (or attach to) a JVM via Py4J and use the classic session in pyspark.sql.session.

The two implementations share the same public surface; users do not generally have to know which mode they are in.

Python UDF execution

sequenceDiagram
    participant E as Executor (JVM)
    participant D as pyspark.daemon
    participant W as pyspark.worker

    E->>D: connect to Python daemon (Unix socket)
    D->>W: fork worker
    E->>W: send serialized UDF + Arrow batches
    W->>W: invoke UDF
    W->>E: return Arrow batches

Daemon and worker live in python/pyspark/{daemon,worker}.py. They exchange data over a socket using either Pickle (the legacy path) or Arrow IPC (the modern, much faster path). The dispatch logic that decides what kind of UDF to call lives in worker.py's read_udfs and wrap_udf functions.

worker.py handles many UDF flavors: regular Python UDF, Pandas UDF (scalar/grouped/window), applyInPandas, applyInArrow, mapInPandas, mapInArrow, Python UDTF, and grouped-map UDFs (the legacy and the Arrow-optimized variants).

Pandas API on Spark

python/pyspark/pandas/ is the merged Project Koalas codebase. It re-implements pandas on top of Spark's DataFrame:

• DataFrame, Series, and Index types.
• Lazy execution mapped to Catalyst plans.
• Per-row access falls back to to_pandas() semantics with explicit warnings.
• Plotting, rolling, expanding, and resample APIs.

Tests live in python/pyspark/pandas/tests/.

Streaming

python/pyspark/sql/streaming/ exposes the Structured Streaming API (DataStreamReader, DataStreamWriter, StreamingQuery, StreamingQueryManager). The legacy python/pyspark/streaming/ exposes the DStream API.

Errors and types

python/pyspark/errors/ mirrors the JVM error-class system. It generates a Pythonic hierarchy (AnalysisException, IllegalArgumentException, PySparkValueError, ...) and maps Java exception classnames into them via pyspark.errors.exceptions.captured.

Type stubs (.pyi files) sit alongside .py files; the type system is enforced by mypy (configured in python/mypy.ini). Stub maintenance is part of any API change.

Test infrastructure

• python/pyspark/testing/ ships re-usable harnesses (ReusedSQLTestCase, PySparkErrorTestUtils, ...) for use in user code and Spark's own tests.
• python/run-tests.py is the runner; python/run-tests-with-coverage adds Coverage.py.
• python/conf_vscode/ and python/run-with-vscode-breakpoints configure breakpoint debugging in VS Code.

Integration points

• In classic mode, calls Py4J into a JVM SparkContext. The gateway script is python/pyspark/java_gateway.py.
• In Connect mode, calls gRPC stubs in python/pyspark/sql/connect/proto/. See connect.md.
• UDF execution depends on Arrow for fast data transfer (pyarrow is a runtime dep for the Arrow-optimized path).
• Cloudpickle is vendored under python/pyspark/cloudpickle/ to avoid version skew with upstream cloudpickle.

Entry points for modification

• Adding a public DataFrame method: edit both python/pyspark/sql/dataframe.py and python/pyspark/sql/connect/dataframe.py to keep parity.
• Adding a SQL function: add a Python wrapper in python/pyspark/sql/functions.py and a Connect mirror in python/pyspark/sql/connect/functions.py. Update the function docs generator.
• Adding a UDF flavor: edit python/pyspark/worker.py (Python side) and the corresponding JVM operator under sql/core/.../execution/python/.
• Adding a pandas-on-Spark feature: edit the relevant file in python/pyspark/pandas/. Add unit tests under python/pyspark/pandas/tests/.