Python Data Source API

Why it exists

Reading from a REST API, a proprietary file format or an internal service used to mean either writing a Scala data source or pulling data into the driver with Python and calling createDataFrame, which does not scale and runs on one machine. The Python Data Source API runs your Python reading code on executors, in parallel, and plugs into the standard reader and writer API.

A batch reader, step by step

from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition

class CountriesSource(DataSource):
    @classmethod
    def name(cls):
        return "countries"                      # used in spark.read.format(...)

    def schema(self):
        return "code string, name string, region string"

    def reader(self, schema):
        return CountriesReader(self.options)

class CountriesReader(DataSourceReader):
    def __init__(self, options):
        self.regions = options.get("regions", "Asia,Europe").split(",")

    def partitions(self):
        return [InputPartition(r) for r in self.regions]   # one task per region

    def read(self, partition):
        import requests                          # runs on the executor
        rows = requests.get(f"https://api.example.com/{partition.value}").json()
        for c in rows:
            yield (c["code"], c["name"], partition.value)

spark.dataSource.register(CountriesSource)
df = spark.read.format("countries").option("regions", "Asia,Africa").load()

1. 1
   name() is the short name used in format().
2. 2
   schema() returns the output schema, as a DDL string or a StructType. Users can also pass .schema(...) themselves.
3. 3
   reader() creates the reader; self.options holds the options passed with .option().
4. 4
   partitions() splits the work. Each partition becomes one Spark task, so this is what makes the source parallel. Without it, one task reads everything.
5. 5
   read(partition) runs on an executor and yields tuples (or Rows) matching the schema. Imports inside it run where it executes.

Writers

A source can also implement writer(schema, overwrite), returning a DataSourceWriter. Its write(iterator) method runs per partition on executors and returns a commit message; commit(messages) runs once on the driver after all tasks succeed, and abort(messages) if any fail. That two-phase shape lets a writer publish results atomically, for example by writing to staging locations and moving them on commit.

Streaming

For Structured Streaming, implement a stream reader with initialOffset(), latestOffset(), partitions(start, end), read(partition) and commit(end). Spark tracks offsets in the checkpoint, so the source gets exactly-once processing as long as reading a given offset range is repeatable. Stream writers follow the same pattern as batch writers, per micro-batch.

Practical advice

• Keep read() self-contained: open connections inside it, because the reader object is pickled and sent to executors.
• Design partitions so each is independent and similar in size: pages of an API, date ranges, file groups.
• Respect rate limits: many partitions mean many concurrent requests.
• Python sources do not get filter or column pushdown automatically; filter inside read() using your own options when the API supports it.

Common mistakes