intermediateView Python source

Choose a Python UDF style

This example computes one score, sqrt(a² + b²) + ln(1 + c), over 100,000 rows in five ways and checks that all five agree with the built-in version.

Run it

Run all commands below from the repository root (sail-examples/):

uv run python/udf-styles/main.py

Output

The newest capture appears first. Older output stays here when a result changes.

Current output · Sail 0.7.3 (revision 2)

Captured on Sail 0.7.3. Latest capture: 2026-11-07, Python 3.14.10.

Scores (rounded to 12 decimal places):
+---+--------------+--------------+--------------+--------------+--------------+
| id|       builtin|       row_udf|    pandas_udf|     arrow_udf|  map_in_arrow|
+---+--------------+--------------+--------------+--------------+--------------+
|  0|           0.0|           0.0|           0.0|           0.0|           0.0|
|  1|0.457012005054|0.457012005054|0.457012005054|0.457012005054|0.457012005054|
|  2|0.849485488971|0.849485488971|0.849485488971|0.849485488971|0.849485488971|
|  3|1.201136978368|1.201136978368|1.201136978368|1.201136978368|1.201136978368|
|  4|1.524617590798|1.524617590798|1.524617590798|1.524617590798|1.524617590798|
+---+--------------+--------------+--------------+--------------+--------------+

Maximum absolute difference from built-in <= 1e-12:
+-------+----------+---------+------------+
|row_udf|pandas_udf|arrow_udf|map_in_arrow|
+-------+----------+---------+------------+
|   true|      true|     true|        true|
+-------+----------+---------+------------+

Earlier output · Sail 0.7.3 (revision 1)

Captured on Sail 0.7.3. Latest capture: 2026-11-05, Python 3.12.14.

+---+-------------------+-------------------+-------------------+-------------------+-------------------+
| id|            builtin|            row_udf|         pandas_udf|          arrow_udf|       map_in_arrow|
+---+-------------------+-------------------+-------------------+-------------------+-------------------+
|  0|                0.0|                0.0|                0.0|                0.0|                0.0|
|  1|0.45701200505444706|0.45701200505444706|0.45701200505444706|0.45701200505444706|0.45701200505444706|
|  2|  0.849485488971323|  0.849485488971323|  0.849485488971323|  0.849485488971323|  0.849485488971323|
|  3|  1.201136978367944|  1.201136978367944|  1.201136978367944|  1.201136978367944|  1.201136978367944|
|  4|  1.524617590797868| 1.5246175907978683| 1.5246175907978683| 1.5246175907978683| 1.5246175907978683|
+---+-------------------+-------------------+-------------------+-------------------+-------------------+

+--------------------+--------------------+--------------------+--------------------+
|             row_udf|          pandas_udf|           arrow_udf|        map_in_arrow|
+--------------------+--------------------+--------------------+--------------------+
|3.552713678800501...|3.552713678800501...|3.552713678800501...|3.552713678800501...|
+--------------------+--------------------+--------------------+--------------------+

Python code

main.py112 lines

View main.py on GitHub

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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "pysail==0.7.3",
#   "pyspark-client==4.1.3",
#   "pyarrow==26.0.0",
#   "pandas==3.0.6",
#   "numpy==2.5.4",
# ]
# ///
"""The same per-row calculation in every Python style Spark offers, checked to agree.

score = sqrt(a^2 + b^2) + ln(1 + c), computed with built-in functions, a row UDF, a pandas UDF,
an Arrow UDF and mapInArrow. The README explains when to reach for each one.
"""

import math
from collections.abc import Iterator

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql import functions as F
from pyspark.sql.functions import arrow_udf, pandas_udf


def inputs(spark: SparkSession) -> DataFrame:
    return spark.range(100_000).select(
        "id",
        ((F.col("id") % 97) / 7.0).alias("a"),
        ((F.col("id") % 89) / 11.0).alias("b"),
        ((F.col("id") % 83) / 3.0).alias("c"),
    )


# 1. Built-in functions: no Python runs per row.
def builtin_score(a, b, c):
    return F.sqrt(a * a + b * b) + F.log1p(c)


# Styles 2 to 4 are plain functions here and become UDFs inside run(): with PySpark 4, wrapping a
# function as a UDF needs an active Spark session, which does not exist yet when a module is imported.


# 2. Row UDF: Python is called once per row.
def row_score(a: float, b: float, c: float) -> float:
    return math.sqrt(a * a + b * b) + math.log1p(c)


# 3. pandas UDF: called once per batch, with pandas Series.
def pandas_score(a: pd.Series, b: pd.Series, c: pd.Series) -> pd.Series:
    return np.sqrt(a * a + b * b) + np.log1p(c)


# 4. Arrow UDF: called once per batch, with Arrow arrays and pyarrow.compute.
def arrow_score(a: pa.Array, b: pa.Array, c: pa.Array) -> pa.Array:
    return pc.add(pc.sqrt(pc.add(pc.multiply(a, a), pc.multiply(b, b))), pc.log1p(c))


# 5. mapInArrow: whole record batches in, whole record batches out. Setup before the loop
#    runs once per partition, which makes it the place to load a model or open a client.
def map_scores(batches: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
    for batch in batches:
        a, b, c = (batch.column(name) for name in ("a", "b", "c"))
        score = pc.add(pc.sqrt(pc.add(pc.multiply(a, a), pc.multiply(b, b))), pc.log1p(c))
        yield pa.RecordBatch.from_arrays([batch.column("id"), score], names=["id", "map_in_arrow"])


def run(spark: SparkSession) -> dict[str, DataFrame]:
    data = inputs(spark)
    a, b, c = F.col("a"), F.col("b"), F.col("c")
    row_udf = F.udf(row_score, "double")
    pandas_udf_score = pandas_udf(pandas_score, "double")
    arrow_udf_score = arrow_udf(arrow_score, "double")
    scores = data.select(
        "id",
        builtin_score(a, b, c).alias("builtin"),
        row_udf(a, b, c).alias("row_udf"),
        pandas_udf_score(a, b, c).alias("pandas_udf"),
        arrow_udf_score(a, b, c).alias("arrow_udf"),
    ).join(data.mapInArrow(map_scores, "id long, map_in_arrow double"), "id")

    styles = ["row_udf", "pandas_udf", "arrow_udf", "map_in_arrow"]
    differences = scores.select(*[F.max(F.abs(F.col(s) - F.col("builtin"))).alias(s) for s in styles])
    return {
        "first_rows": scores.orderBy("id").limit(5),
        "max_difference_from_builtin": differences,
        "rows": scores.agg(F.count("*").alias("rows")),
    }


if __name__ == "__main__":
    from pysail.spark import SparkConnectServer

    server = SparkConnectServer()
    server.start()
    _, port = server.listening_address
    spark = SparkSession.builder.remote(f"sc://localhost:{port}").getOrCreate()
    results = run(spark)
    # Keep the console comparison readable across math libraries and CPU instruction sets.
    print("Scores (rounded to 12 decimal places):")
    first_rows = results["first_rows"]
    first_rows.select(
        "id", *[F.round(F.col(name), 12).alias(name) for name in first_rows.columns if name != "id"]
    ).show()
    print("Maximum absolute difference from built-in <= 1e-12:")
    differences = results["max_difference_from_builtin"]
    differences.select(*[(F.col(name) <= 1e-12).alias(name) for name in differences.columns]).show()
    spark.stop()
    server.stop()

The five styles, and when to use each

Style	Python is called	Use it when
Built-in functions (F.sqrt, F.log1p)	never	a built-in function exists. Check first; most do.
Row UDF (@F.udf)	once per row	the logic is plain Python on single values and nothing above fits.
pandas UDF (@pandas_udf)	once per batch of rows, with pandas Series	the logic is pandas or NumPy code.
Arrow UDF (@arrow_udf)	once per batch of rows, with Arrow arrays	the logic is pyarrow.compute code, or you want to avoid pandas.
mapInArrow	once per partition, over its record batches	you need setup once per partition (a model, a client), or the output has different rows or columns than the input.

The console rounds scores to 12 decimal places and shows whether each style's maximum absolute difference from the built-in result is at most 1e-12. Smaller rounding differences can vary across CPUs and math libraries. The run() results, including max_difference_from_builtin, retain full precision for the Sail/Spark comparison tests.

On Sail

• Python UDFs run inside the Sail server process, in one interpreter. NumPy and PyArrow compute functions release the GIL while they work on a batch.
• In an iterator-style arrow_udf, code before the loop runs once per batch, not once per partition. Put per-partition setup in mapInArrow, mapInPandas or an iterator pandas_udf.
• Build mapInArrow batches with Arrow types that match the declared schema; Apache Spark rejects an int64 column declared as int.

Related Sail behavior

• Create Python UDFs after starting the session
• Python UDFs share an interpreter in local mode
• Match Arrow output types to the declared schema
• Load a model once per partition

Keep exploring

AI data preparationEnrich a table with a model →Run a model once per distinct value, then join its answers back. AI data preparationAlign robot camera frames with state and actions →Match robot camera frames to nearby state and action records.