# Licensed to the Apache Software Foundation (ASF) under one

# or more contributor license agreements. See the NOTICE file

# distributed with this work for additional information

# regarding copyright ownership. The ASF licenses this file

# to you under the Apache License, Version 2.0 (the

# "License"); you may not use this file except in compliance

# with the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing,

# software distributed under the License is distributed on an

# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

# KIND, either express or implied. See the License for the

# specific language governing permissions and limitations

# under the License.

"""Distribute DataFusion expressions to Ray actors.

For background — the shipped-expression model, what travels inline vs

by name, portability requirements, and the security threat model —

see https://datafusion.apache.org/python/user-guide/distributing-work/expressions.html.

Build an expression in the driver, ship it to a pool of Ray actors, and

have each actor evaluate it against its own slice of data. Python UDFs

travel with the shipped expression — no actor-side registration needed.

Prerequisites:

pip install ray

Run:

python examples/ray_pickle_expr.py

"""

import pyarrow as pa

import ray

from datafusion import Expr, SessionContext, col, lit, udf

def _build_double_udf():

"""Return the demo UDF used by the driver."""

return udf(

lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]),

[pa.int64()],

pa.int64(),

volatility="immutable",

name="double",

)

@ray.remote

class DataFusionWorker:

"""A Ray actor with a private :class:`SessionContext`."""

def __init__(self) -> None:

self._ctx = SessionContext()

def evaluate(self, expr: Expr, batch_pylist: list[int]) -> list[int]:

"""Run the expression against an in-memory batch."""

# `expr` arrived here via Ray's automatic argument serialization;

# the Python UDF inside it was reconstructed from the bytes — no

# pre-registration on this actor required.

df = self._ctx.from_pydict({"a": batch_pylist})

out = df.with_column("result", expr).select("result")

return out.to_pydict()["result"]

def main() -> None:

ray.init(ignore_reinit_error=True)

expr = _build_double_udf()(col("a")) + lit(1)

workers = [DataFusionWorker.remote() for _ in range(2)]

batches = [[1, 2, 3], [10, 20, 30], [100, 200, 300]]

futures = [

workers[i % len(workers)].evaluate.remote(expr, batch)

for i, batch in enumerate(batches)

]

for batch, result in zip(batches, ray.get(futures), strict=True):

print(f"input {batch} -> {result}")

ray.shutdown()

if __name__ == "__main__":

main()