Enriching Data in PySpark Without Async I/O

Enriching rows with data from an external service is one of the most common things we do in a data pipeline. Call an HTTP API, attach the response, move on.

In Apache Flink this is a solved problem: there is an Async I/O operator. In Apache Spark there is nothing like it, and the obvious solutions are either slow or fragile.

This article goes through the troubles with enrichment in Spark and presents two small open-source libraries for PySpark: michalklempa/spark-async and michalklempa/spark-http-async.

TLDR

uv add "spark-http-async @ git+https://github.com/michalklempa/spark-http-async"

from spark_http_async import enrich

enriched = enrich(requests_df, max_in_flight=50)

requests_df has a url and optional method (POST by default), headers map and binary body columns. You get the same rows back with response_body, response_status_code, request_elapsed_ms and response_error attached.

Problem

Let us take the simplest example. We have a DataFrame of questions and an HTTP service that answers a batch of them:

POST http://localhost:4040/
[{"question_id": "a"}, {"question_id": "b"}]

The service is not slow by itself, say 15 ms per request. But 15 ms multiplied by a million requests, one after another, is more than four hours of pure waiting. Network latency, not CPU, is the bottleneck. The CPU sits idle while we wait for the answer.

The way out is to keep many requests in flight at the same time. How you do that in Spark is where it gets uncomfortable.

What Flink does

Flink ships the AsyncDataStream operator. You implement an AsyncFunction, which gets an element and a result future. The operator keeps a configurable capacity of requests in flight, handles timeouts and offers ordered or unordered output. Your code never touches threads or event loops.

AsyncDataStream.unorderedWait(stream, new MyAsyncFunction(), 1000,TimeUnit.MILLISECONDS,100);

Spark has no such operator. Not in the RDD API, not in the DataFrame API, not in Structured Streaming.

The options in Spark

UDF per row

The first thing everybody writes is a Python UDF, or rdd.map, that does a requests.post per row. It works, it is trivial, and it is as slow as it sounds: one request at a time per task, a new connection per row unless you are careful, and the executor core is blocked on network wait for most of its time.

mapPartitions with a client per partition

The next step is mapPartitions, one HTTP client per partition, so that connections are reused. This fixes the connection setup cost, but each partition still sends one request at a time. Concurrency equals the number of cores, and that is far too low for an I/O-bound job.

More partitions than cores

One can repartition into hundreds of partitions, so that more tasks are scheduled than there are cores. Spark will still run only as many tasks at once as there are task slots. Oversubscribing cores with spark.executor.cores helps, but it is a cluster-wide knob that affects every other job, and each task carries a full Python worker process.

Load the partition and asyncio.gather

The most popular workaround: inside mapPartitions, load the whole partition into a list, fire all requests with asyncio.gather under a semaphore, return the results.

def enrich_partition(rows):
    rows = list(rows)
    return asyncio.run(fetch_all(rows))

This is fast, but it has costs you pay later:

the entire partition and all its responses live in memory at once,
nothing is returned until the last request finishes, so downstream stages are idle meanwhile,
one failing request can take the whole gather down, if you did not wrap it,
partition size now controls your memory, and your concurrency story is a semaphore bolted on the side, which must be kept in sync with the HTTP client connection limits.

What we actually want

Flink’s operator gives us the right shape, and we can copy it:

open() once per partition, to create the clients,
an async map(row) per row,
a bounded number of calls in flight,
no collecting of the partition, rows are read only when there is room for them,
close() at the end, also when something fails.

spark-async

spark-async implements exactly this on top of mapPartitions. You implement an AsyncMapper:

import httpx
from spark_async import AsyncMapper, map_partitions_async


class Lookup(AsyncMapper[dict, dict]):
    def __init__(self, base_url: str) -> None:
        self.base_url = base_url

    async def open(self) -> None:
        self.client = httpx.AsyncClient(base_url=self.base_url)

    async def map(self, row: dict) -> dict:
        response = await self.client.get("/items", params={"id": row["id"]})
        return {**row, "status": response.status_code}

    async def close(self) -> None:
        await self.client.aclose()


enriched = map_partitions_async(rdd, Lookup("http://localhost:4040"), max_in_flight=50)

The result is an ordinary RDD. Call spark.createDataFrame(enriched, schema) if you want a DataFrame.

How it works

For each partition, the library starts max_in_flight worker tasks on a private asyncio event loop. The workers share the partition’s input iterator, so a row is read only when a worker is free. Each worker awaits mapper.map(row) and puts the result on a bounded queue of size max_in_flight.

Spark pulls results from a plain Python generator. Each pull runs the event loop until at least one result exists, hands it over, together with everything else already finished, and returns. A full queue blocks the workers, so a slow consumer slows down the input reads too.

This gives the properties we wanted:

at most max_in_flight rows are being mapped and at most max_in_flight finished results wait in the queue, nothing else is held in memory,
results are streamed out as they complete, so the next Spark stage can start working immediately,
if a map() call raises, the remaining calls are cancelled, close() runs and the exception reaches Spark as a task failure. Errors you want to keep as data should be handled inside map().

Two things to keep in mind. Results come back in completion order, not input order, the same as unorderedWait in Flink. And the event loop runs only while Spark waits for the next result, so requests make no progress while Spark is busy with a row it has just received.

spark-http-async

Most of the time the mapper is “send an HTTP request, keep the response”. For that case there is spark-http-async, a thin layer over spark-async. It is built in two layers.

The core is send(df, build_request, max_in_flight). The DataFrame can have any structure, build_request is a function (row) -> httpx.Request which runs on the executors for every row. The library sends the request, carries all input columns through and attaches the response.

On top of it sits enrich(df, max_in_flight), which is just send with a default builder. The builder reads a url column and optional method (POST when missing), headers and body columns:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

from spark_http_async import enrich

spark = SparkSession.builder.master("local[*]").getOrCreate()

requests = (
    spark.createDataFrame([("a",), ("b",), ("c",)], ["question_id"])
    .withColumn("url", F.lit("http://localhost:4040/"))
    .withColumn("headers", F.create_map(F.lit("Content-Type"), F.lit("application/json")))
    .withColumn("body", F.to_json(F.array(F.struct("question_id"))).cast("binary"))
)

enriched = enrich(requests, max_in_flight=50)

enriched.select(
    "question_id",
    "response_status_code",
    F.col("response_body").cast("string").alias("response"),
    "response_error",
).show(truncate=False)

Any other columns are carried through unchanged. A 4xx or 5xx response is a normal result with its status code, while transport-level failures such as connection refused or a timeout end up in the response_error column, with the other response columns null. A bad endpoint does not kill your job.

If the default columns do not fit, for example you need a GET with query parameters, pass your own builder. Define it at module level, Spark pickles it to the executors:

import httpx

from spark_http_async import send

def get_item(row) -> httpx.Request:
    return httpx.Request("GET", "http://localhost:4040/items", params={"id": row.id})

responses = send(df, get_item, max_in_flight=50)

Because the builder returns a real httpx.Request, everything httpx supports is available: params, auth, cookies, JSON bodies.

max_in_flight is per partition, so the real concurrency is the number of partitions running at once times max_in_flight. If the target service has rate limits, do the multiplication before you pick the number. If the work is spread unevenly, repartition the input first.

Benchmarks

The spark-http-async repository contains a benchmarks/ project which compares enrich() with the approaches above against a small Spring Boot service that sleeps 15 ms per request:

requests per row,
requests per partition,
httpx client per partition,
httpx with asyncio.gather over the whole partition,
enrich() itself.

cd spark-http-async/benchmarks
uv sync
uv run python benchmark.py --questions 5000 --batch-size 10 --partitions 4

A local run with 5000 questions in batches of 10 (500 requests) and 4 partitions gave the following total times:

Variant	total_ms
`requests` per row	10081
`requests` per partition	11518
`httpx` per partition	10088
`httpx` + `asyncio.gather`	1594
`enrich()`	1760

The async variants are about six times faster than the sequential ones. The difference between enrich() and the gather variant is within the noise, differences under about 15 % between runs are. This is a single run against a service sleeping 15 ms per request, so run it on your own data and your own service, the numbers depend heavily on latency and on how many partitions you run at once.

Running on a cluster

The executors’ Python processes must be able to import spark_http_async, spark_async and httpx. Either install the packages into the Python environment of every worker node, or ship a packed virtual environment with spark-submit --archives and point PYSPARK_PYTHON at it.

In local mode, make sure the workers use the same interpreter as the driver:

import os
import sys

os.environ["PYSPARK_PYTHON"] = sys.executable

Summary

Spark has no async I/O operator, so enrichment from an external service ends up as per-row calls, which are slow, or as whole-partition gather calls, which are memory hungry and bursty. A bounded, streaming mapPartitions over an asyncio loop gives you the same model Flink offers out of the box.

Both libraries are MIT licensed, require Python 3.13 and PySpark 4.2 or newer, and are available at:

michalklempa/spark-async, the generic AsyncMapper,
michalklempa/spark-http-async, the HTTP enrich() on top.

Pull requests and issues are welcome.