Load a model once per partition

To load a model or client once per partition, put the setup before the loop in mapInArrow, mapInPandas or an iterator pandas_udf. In an iterator arrow_udf, code before the loop runs once per batch instead.

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[udf-setup-per-partition[sail]]'

Python code

udf-setup-per-partition.py79 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
79FACT = {
    "id": "udf-setup-per-partition",
    "title": "Load a model once per partition",
    "category": "python-udfs",
    "core": True,
    "text": """To load a model or client once per partition, put the setup before the loop in
        `mapInArrow`, `mapInPandas` or an iterator `pandas_udf`. In an iterator `arrow_udf`, code before the
        loop runs once per batch instead.""",
    "engines": ["sail"],
    "order": 2,
}

import tempfile
from collections.abc import Iterator
from pathlib import Path

import pandas as pd
import pyarrow as pa
from pyspark.sql import functions as F
from pyspark.sql.functions import arrow_udf, pandas_udf


def check(spark, engine):
    log_dir = Path(tempfile.mkdtemp())
    log_path = str(log_dir)

    def mark(kind, what):
        with open(f"{log_path}/{kind}-{what}.log", "a") as f:
            f.write("x\n")

    def count(kind, what):
        path = log_dir / f"{kind}-{what}.log"
        return len(path.read_text().splitlines()) if path.exists() else 0

    @arrow_udf("long")
    def arrow_iter(it: Iterator[pa.Array]) -> Iterator[pa.Array]:
        mark("arrow_udf", "setup")
        for array in it:
            mark("arrow_udf", "batch")
            yield array

    @pandas_udf("long")
    def pandas_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
        mark("pandas_udf", "setup")
        for series in it:
            mark("pandas_udf", "batch")
            yield series

    def map_arrow(it):
        mark("mapInArrow", "setup")
        for batch in it:
            mark("mapInArrow", "batch")
            yield batch

    def map_pandas(it):
        mark("mapInPandas", "setup")
        for frame in it:
            mark("mapInPandas", "batch")
            yield frame

    partitions = 2
    base = spark.range(0, 8000).repartition(partitions)
    base.select(arrow_iter("id").alias("y")).agg(F.sum("y")).collect()
    base.select(pandas_iter("id").alias("y")).agg(F.sum("y")).collect()
    base.mapInArrow(map_arrow, "id long").agg(F.sum("id")).collect()
    base.mapInPandas(map_pandas, "id long").agg(F.sum("id")).collect()

    kinds = ("arrow_udf", "pandas_udf", "mapInArrow", "mapInPandas")
    observed = {k: {"setup": count(k, "setup"), "batches": count(k, "batch")} for k in kinds}
    for kind in ("pandas_udf", "mapInArrow", "mapInPandas"):
        assert observed[kind]["setup"] == partitions, (
            f"{kind} setup should run once per partition: {observed}"
        )
        assert observed[kind]["batches"] > partitions, f"{kind} should see several batches: {observed}"
    arrow = observed["arrow_udf"]
    assert arrow["setup"] == arrow["batches"] > partitions, (
        f"arrow_udf setup should run per batch: {observed}"
    )
    return observed

Load a model once per partition

Run the check

Sail

Python code

Related examples