Read independent source partitions concurrently

A Python data source's partitions are read as separate tasks, so read() calls for different partitions run at the same time. Return one partition per page, file or key range to fetch an API or a set of files concurrently.

Check targets: pysail==0.7.3 with pyspark-client==4.1.3

The Python assertions below check this behavior. Details attributed to external documentation are outside those assertions. Newer releases may behave differently.

Run the check

Follow the test setup, then run from the repository root (sail-examples/). The harness applies the configuration declared in this check.

Sail

uv run pytest 'tests/test_facts.py::test_fact[python-data-source-partitions-concurrent[sail]]'

Apache Spark (requires Java 17)

uv run pytest 'tests/test_facts.py::test_fact[python-data-source-partitions-concurrent[spark]]'

Python code

python-data-source-partitions-concurrent.py87 lines

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87FACT = {
    "id": "python-data-source-partitions-concurrent",
    "title": "Read independent source partitions concurrently",
    "category": "data-sources",
    "text": """A Python data source's partitions are read as separate tasks, so `read()` calls for
        different partitions run at the same time. Return one partition per page, file or key range to
        fetch an API or a set of files concurrently.""",
    "engines": ["sail", "spark"],
    "order": 2,
}

import json
import threading
import time
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

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

PARTITIONS = 8


class SlowService:
    """Answers every request after 200 ms and records the most requests in flight at once."""

    def __init__(self):
        self.active = self.peak = 0
        lock = threading.Lock()
        service = self

        class Handler(BaseHTTPRequestHandler):
            def do_GET(self):
                with lock:
                    service.active += 1
                    service.peak = max(service.peak, service.active)
                time.sleep(0.2)
                with lock:
                    service.active -= 1
                body = json.dumps({"value": 1}).encode()
                self.send_response(200)
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)

            def log_message(self, *args):
                pass

        self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        self.url = f"http://127.0.0.1:{self.server.server_port}/"
        threading.Thread(target=self.server.serve_forever, daemon=True).start()


class SlowReader(DataSourceReader):
    def __init__(self, url):
        self.url = url

    def partitions(self):
        return [InputPartition(i) for i in range(PARTITIONS)]

    def read(self, partition):
        with urllib.request.urlopen(self.url, timeout=30) as response:
            yield (partition.value, json.load(response)["value"])


class SlowSource(DataSource):
    @classmethod
    def name(cls):
        return "fact_slow_source"

    def schema(self):
        return "part int, value int"

    def reader(self, schema):
        return SlowReader(self.options["url"])


def check(spark, engine):
    service = SlowService()
    try:
        spark.dataSource.register(SlowSource)
        rows = spark.read.format("fact_slow_source").option("url", service.url).load().collect()
    finally:
        service.server.shutdown()
        service.server.server_close()
    assert len(rows) == PARTITIONS, f"expected {PARTITIONS} rows, got {len(rows)}"
    assert service.peak > 1, f"partitions were read one at a time (peak {service.peak})"
    return {"partitions": PARTITIONS, "peak_concurrent_reads": service.peak}

Read independent source partitions concurrently

Run the check

Sail

Apache Spark (requires Java 17)

Python code

Related examples