Writing a PySpark Custom Data Source

What this is about

For most of my Spark work, reading from an unusual source meant one of two things. Either pull the data down with Python first and wrap it in createDataFrame, which skips the engine entirely. Or write Scala, learn the DataSource V2 interfaces, and ship a JAR.

Spark 4.0 added a third option, the Python Data Source API. You subclass two classes, register them, and spark.read.format("your-thing") works like any built-in source. No JVM code and no JAR.

On open source Spark it needs 4.0 or later. On Databricks the requirement is Runtime 15.4 LTS and above, or serverless environment version 2. Worth noting that DBR 15.4 LTS is built on Spark 3.5.0, so Databricks had this backported and available before OSS 4.0 shipped it.

I wanted to use it rather than just read about it, so I contributed a connector to pyspark-data-sources, a collection of custom PySpark connectors. My addition was a JSONPlaceholder data source, merged in August 2025.

JSONPlaceholder is a free fake REST API, the thing most people reach for when testing an HTTP client. It serves six endpoints: users, posts, albums, todos, comments and photos.

What makes it worth wrapping is that the endpoints reference each other. posts.userId points at users.id. comments.postId points at posts.id. photos.albumId points at albums.id, which points back at users.id. That gives you a small relational dataset behind a live API, with no database, no credentials and no fixture files.

The DataSource implementation

The contract is two classes. DataSource describes what the source is. DataSourceReader does the fetching.

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

class JSONPlaceholderDataSource(DataSource):
    @classmethod
    def name(cls) -> str:
        return "jsonplaceholder"          # the string in .format()

    def __init__(self, options=None):
        self.options = options or {}      # everything from .option()

    def schema(self) -> str:
        ...                               # DDL string or StructType

    def reader(self, schema) -> DataSourceReader:
        return JSONPlaceholderReader(self.options)


class JSONPlaceholderReader(DataSourceReader):
    def partitions(self) -> List[InputPartition]:
        return [InputPartition(0)]        # how to split the work

    def read(self, partition) -> Iterator[Row]:
        ...                               # yield Rows for one partition

partitions() runs on the driver. read() runs on the executors, once per partition, in parallel.

Six endpoints with six different shapes raises a design question: do you register six data sources, or one? Six formats would mean format("jsonplaceholder-posts"), format("jsonplaceholder-users") and so on, which is a lot of near-duplicate classes. I went with one format and an endpoint option, which works because schema() can read the options:

def schema(self) -> str:
    schemas = {
        "posts":    "userId INT, id INT, title STRING, body STRING",
        "todos":    "userId INT, id INT, title STRING, completed BOOLEAN",
        "comments": "postId INT, id INT, name STRING, email STRING, body STRING",
        "albums":   "userId INT, id INT, title STRING",
        "photos":   "albumId INT, id INT, title STRING, url STRING, thumbnailUrl STRING",
        "users":    "id INT, name STRING, ...",
    }
    endpoint = self.options.get("endpoint", "posts")
    return schemas.get(endpoint, schemas["posts"])

schema() is called during query planning, after __init__ has already received the options. So one registered format resolves to a different schema on every read, decided at plan time. .option("endpoint", "todos") changes the columns of the DataFrame. This is the part of the API I’d not have guessed from the docs, and it’s what turns one format into a family of related sources.

Two smaller details. Schemas are DDL strings like "userId INT, id INT" rather than StructType trees, which is a lot less code for flat schemas. And the default matters: with no endpoint you get posts, so spark.read.format("jsonplaceholder").load() works with no configuration at all.

Most of the remaining code went into turning JSON into columns. Most endpoints are already flat but users is not, with nested address, address.geo and company objects. You can model that with nested StructTypes. I flattened it instead, using underscores in the column names:

address_street, address_suite, address_city, address_zipcode,
address_geo_lat, address_geo_lng,
company_name, company_catchPhrase, company_bs

The connector exists to demonstrate joins, and users.address_city is easier to select than users.address.city. No struct navigation, it works inside SQL strings, and it survives df.toPandas().

Every field access in the row builders is .get() with a default. If a key is missing and you index with [] instead, the KeyError happens inside an executor task and comes back as a Py4J stack trace wrapped in a Spark job failure. That’s much harder to debug than an empty string in a column.

One trap worth naming: lat and lng come back from the API as strings, so they’re STRING in the schema. Typing them as DOUBLE would be nicer to use and wrong about the source. I left them as strings and let the caller cast.

Where this fits in Databricks testing

The connector is built into the package, so there’s no extra to install:

%pip install pyspark-data-sources

Then register it once per session and read:

from pyspark_datasources import JSONPlaceholderDataSource

spark.dataSource.register(JSONPlaceholderDataSource)

users = spark.read.format("jsonplaceholder").option("endpoint", "users").load()
posts = spark.read.format("jsonplaceholder").option("endpoint", "posts").load()

posts_with_authors = posts.join(users, posts.userId == users.id)

spark.dataSource.register() is session scoped, so it belongs in the notebook itself or in a shared setup cell rather than somewhere global.

Where I’ve found this genuinely useful on Databricks:

Testing transformation logic with no data access. Getting Unity Catalog grants sorted out in a dev workspace is often slower than writing the transformation you wanted to test. This gives you two related tables to join in the time it takes a cluster to start, with no grants and no secrets.

Smoke-testing a notebook end to end. If you’re building a bronze to silver skeleton, you can point the read at this, confirm the joins and the write path run, and swap the source later. Real referential integrity means a join that works here fails for real reasons, not because fake keys never matched.

Onboarding and demos. You can show somebody the Data Source API or a three-way join on a running cluster without provisioning any access first.

Use .option("limit", "5") when you want a small DataFrame, and .option("id", "1") to pull a single record.

Two limits to be aware of. The connector is single-partition, since partitions() returns [InputPartition(0)], so one executor does all the fetching. For 5000 rows over one HTTP round-trip that’s the right call, but it means the source is useless for benchmarking anything about parallelism.

The other is that it needs outbound internet from wherever the task runs, which is an executor rather than the driver. Plenty of production workspaces block that, so this works in a dev workspace and fails in a locked-down one. Any test you come to rely on needs to account for that, and it’s worth knowing before you put it in a job rather than after.

Packaging this for internal APIs

The testing use is the small version of why this matters. If your company has a lot of internal APIs, the pattern is worth more than any single connector.

Write your data sources once as a Python package, publish the wheel to Artifactory or whatever private index you already run, and install it from the libraries section of a DAB project or with %pip install in a notebook. From then on, reading an internal API is a format() call:

from company_datasources import register_all

register_all(spark)

orders = (spark.read.format("orders-api")
    .option("region", "eu")
    .load())

A register_all(spark) helper in the package is worth adding, because registration is session scoped. Every notebook and every job still has to register, so you may as well make that one import rather than a list of spark.dataSource.register() calls that drifts as the package grows.

What you get out of it is that the auth, the pagination, the retry behavior and the schema live in one versioned place instead of being pasted into every notebook that needs that API. The people consuming it write DataFrame code and never see requests. When the API adds a field or changes its paging, you bump a version instead of hunting through notebooks. And the schema becomes something you can actually review, since it’s declared in one file rather than implied by whatever dictionary parsing somebody wrote that afternoon.

This is the case where the Python Data Source API earns its keep. A single connector for a public API is a nice exercise. A shared package of connectors for the APIs your teams already depend on changes how much boilerplate exists in the organization.

What the engine does with it

Worth understanding what you’re actually getting, because a custom source does not behave like a file format. Here is the same query, filter("userId = 1").select("id", "title"), against the connector and then against the same rows written to Parquet. Both plans are from open source Spark 4.0.1, so exact node text can differ on a given Databricks runtime.

The Python data source:

== Physical Plan ==
*(1) Project [id#1, title#2]
+- *(1) Filter (isnotnull(userId#0) AND (userId#0 = 1))
   +- BatchScan jsonplaceholder[userId#0, id#1, title#2, body#3] (Python) RuntimeFilters: []

Parquet:

== Physical Plan ==
*(1) Project [id#17, title#18]
+- *(1) Filter (isnotnull(userId#16) AND (userId#16 = 1))
   +- *(1) ColumnarToRow
      +- FileScan parquet [userId#16,id#17,title#18] Batched: true,
         PushedFilters: [IsNotNull(userId), EqualTo(userId,1)],
         ReadSchema: struct<userId:int,id:int,title:string>

Three things differ.

The scan node. Python data sources go through DataSource V2, so the logical plan shows RelationV2 and the physical plan shows BatchScan tagged (Python). Parquet still uses the V1 FileScan path with a Location and an InMemoryFileIndex behind it. Before analysis the plan is just UnresolvedDataSource format: jsonplaceholder, which is where schema() gets called to resolve the columns.

No pushdown at all. Parquet’s scan reports PushedFilters and a ReadSchema listing three of the four columns. The predicate went into the reader and body is never read off disk. The BatchScan line lists all four columns including body, which this query never uses, and there is no PushedFilters entry. The filter and the projection run as ordinary Spark operators after every row has already been fetched and handed over.

That isn’t something I skipped implementing. In PySpark 4.0.1, DataSourceReader has exactly two methods, partitions and read. There’s no pushdown hook to implement, so Spark has no way to tell the source that only userId = 1 was wanted.

Spark 4.1 changed half of this. It added DataSourceReader.pushFilters(filters), which the planner calls once with the filters it would like the source to handle, and you return the ones Spark still needs to apply itself. It’s batch reads only, not streaming, and the docs suggest implementing it only for sources that filter natively, such as databases and GraphQL APIs. On Databricks that arrives with Runtime 18 LTS, which is the release built on Spark 4.1.0.

Column pruning is a different story. There’s still no hook for it as of the 4.2 API docs, so body gets fetched and shipped whether or not the query wants it. Even with pushFilters implemented, the source can’t be told which columns to skip.

For this connector on 4.0, the options are the pushdown. endpoint, limit and id are read at planning time and are the only way to make it fetch less. A .filter() never reaches the API.

Rows cross a process boundary. Parquet reads columnar batches in the JVM and ColumnarToRow converts them, which is why that operator appears. The Python plan has no ColumnarToRow because the rows don’t start in the JVM at all. Spark launches a Python worker per task, your read() yields Row objects, and they get batched into Arrow and shipped across.

Arrow isn’t optional here, even though the docs present yielding Arrow batches yourself as a performance option. pyspark/sql/worker/plan_data_source_read.py imports pyarrow at the top of the module, so the read path needs it regardless of what your read() yields. I only found out because a local job died with ModuleNotFoundError: No module named 'pyarrow'. You won’t hit this on Databricks, where pyarrow is preinstalled on every runtime, but you will hit it in a bare local venv.

The same boundary applies at planning time. The driver starts a Python worker to call partitions() while it builds the plan, so an expensive partitions() slows down .load() and explain(), not just execution.

One consequence is easy to miss: nothing is cached between actions. Against a local server that counts requests, count() produced one request, a following collect() produced a second, and a third action produced a third. Adding .cache() brought three actions back down to one request. Re-reading a Parquet file is cheap and returns identical bytes. Re-reading an API costs a round trip and may return different data, so two actions on the same uncached DataFrame aren’t guaranteed to agree with each other. Cache or persist as soon as you’ve loaded.

None of this makes the API less useful, but it does tell you where it fits. A custom source is the right tool for data that only exists behind an API. It isn’t a way to make that API behave like a table.

The API itself turned out to be small. The whole file is 220 lines including docstrings, and most of that is JSON flattening rather than Spark plumbing. If you can write a requests call and describe a schema, you can write a connector, and the harder part was never the interface. It was deciding how the shape of the source should look as columns.