# 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.

"""Driver- and worker-side setup for distributing DataFusion expressions.

When a :class:`Expr` is shipped to a worker process (e.g. through

:func:`multiprocessing.Pool` or a Ray actor), the worker reconstructs the

expression against a :class:`SessionContext`. If the expression references

UDFs imported via the FFI capsule protocol — or any UDF the worker would

otherwise resolve from its registered functions rather than from inside

the shipped expression — install a configured :class:`SessionContext`

once per worker:

.. code-block:: python

from datafusion import SessionContext

from datafusion.ipc import set_worker_ctx

def init_worker():

ctx = SessionContext()

ctx.register_udaf(my_ffi_aggregate)

set_worker_ctx(ctx)

Built-in functions and Python UDFs (scalar, aggregate, window) travel

inside the shipped expression itself and do not need pre-registration

on the worker.

.. note:: Serialization model

Expressions containing Python UDFs (scalar, aggregate, window) are

serialized using :mod:`cloudpickle`. The callable itself travels

**by value** (bytecode and closure cells inlined), but any names the

callable resolves via ``import`` are captured **by reference** and

must be importable on the receiving worker.

The serialized payload is stamped with the sender's Python

``(major, minor)`` version. Loading on a different minor version

raises :class:`ValueError` with an actionable message — cloudpickle

payloads are not portable across Python minor versions. See

:meth:`datafusion.Expr.to_bytes` for examples of what travels by

value vs. by reference.

On the driver side, call :func:`set_sender_ctx` to control how

:func:`pickle.dumps` encodes expressions — for example, to apply

:meth:`SessionContext.with_python_udf_inlining` to every pickled

expression on this thread:

>>> import pickle

>>> from datafusion import SessionContext, col, lit

>>> from datafusion.ipc import clear_sender_ctx, set_sender_ctx

>>> driver_ctx = SessionContext().with_python_udf_inlining(enabled=False)

>>> set_sender_ctx(driver_ctx)

>>> try:

... blob = pickle.dumps(col("a") + lit(1))

... finally:

... clear_sender_ctx()

>>> isinstance(blob, bytes)

True

Without a sender context the default codec is used (Python UDF

inlining on). The sender context only affects pickle / ``to_bytes``

encoding; explicit ``expr.to_bytes(ctx)`` calls still use the supplied

``ctx``.

The thread-local sender context holds a strong reference to the

installed :class:`SessionContext` until :func:`clear_sender_ctx` is

called or the thread exits. Long-running driver threads that install a sender

context once and never clear it will retain that session for the

lifetime of the thread; pair :func:`set_sender_ctx` with

:func:`clear_sender_ctx` (e.g. in a ``try``/``finally``) when the

sender context is only needed for a bounded scope.

"""

from __future__ import annotations

import threading

from typing import TYPE_CHECKING

if TYPE_CHECKING:

from datafusion.context import SessionContext

__all__ = [

"clear_sender_ctx",

"clear_worker_ctx",

"get_sender_ctx",

"get_worker_ctx",

"set_sender_ctx",

"set_worker_ctx",

]

_local = threading.local()

def set_worker_ctx(ctx: SessionContext) -> None:

"""Install this worker's :class:`SessionContext` for shipped expressions.

Call once per worker — typically from a ``multiprocessing.Pool``

initializer or a Ray actor ``__init__``. Idempotent: overwrites any

previous value. Stored in a thread-local slot, so each thread within a

worker may install its own context independently.

Examples:

>>> from datafusion import SessionContext

>>> from datafusion.ipc import set_worker_ctx, get_worker_ctx, clear_worker_ctx

>>> set_worker_ctx(SessionContext())

>>> get_worker_ctx() is not None

True

>>> clear_worker_ctx()

"""

_local.ctx = ctx

def clear_worker_ctx() -> None:

"""Remove this worker's installed :class:`SessionContext`.

After clearing, expressions reconstructed in this worker fall back to

the global :class:`SessionContext` — adequate for built-ins and Python

UDFs (scalar, aggregate, window), but anything imported via the FFI

capsule protocol must be registered on the global context to resolve.

Examples:

>>> from datafusion import SessionContext

>>> from datafusion.ipc import set_worker_ctx, clear_worker_ctx, get_worker_ctx

>>> set_worker_ctx(SessionContext())

>>> clear_worker_ctx()

>>> get_worker_ctx() is None

True

"""

if hasattr(_local, "ctx"):

del _local.ctx

def get_worker_ctx() -> SessionContext | None:

"""Return this worker's installed :class:`SessionContext`, or ``None``.

Examples:

>>> from datafusion.ipc import get_worker_ctx, clear_worker_ctx

>>> clear_worker_ctx()

>>> get_worker_ctx() is None

True

"""

return getattr(_local, "ctx", None)

def set_sender_ctx(ctx: SessionContext) -> None:

"""Install this driver's :class:`SessionContext` for outbound pickles.

Controls how :func:`pickle.dumps` encodes :class:`Expr` instances on

this thread. The most useful application is propagating a session

configured with

:meth:`SessionContext.with_python_udf_inlining` so the toggle takes

effect through pickle (which otherwise calls

:meth:`Expr.to_bytes` with no context and uses the default codec).

Idempotent: overwrites any previous value. Stored in a thread-local

slot, so worker threads on the driver may install their own contexts.

Does not affect :meth:`Expr.to_bytes` calls that pass an explicit

``ctx`` — those continue to use the supplied context.

Examples:

>>> from datafusion import SessionContext

>>> from datafusion.ipc import set_sender_ctx, get_sender_ctx

>>> driver = SessionContext().with_python_udf_inlining(enabled=False)

>>> set_sender_ctx(driver)

>>> get_sender_ctx() is driver

True

"""

_local.sender_ctx = ctx

def clear_sender_ctx() -> None:

"""Remove this driver's installed sender :class:`SessionContext`.

After clearing, pickled expressions fall back to the default codec

(Python UDF inlining on).

Examples:

>>> from datafusion import SessionContext

>>> from datafusion.ipc import (

... set_sender_ctx, clear_sender_ctx, get_sender_ctx,

... )

>>> set_sender_ctx(SessionContext())

>>> clear_sender_ctx()

>>> get_sender_ctx() is None

True

"""

if hasattr(_local, "sender_ctx"):

del _local.sender_ctx

def get_sender_ctx() -> SessionContext | None:

"""Return this driver's installed sender :class:`SessionContext`, or ``None``.

Examples:

>>> from datafusion.ipc import get_sender_ctx, clear_sender_ctx

>>> clear_sender_ctx()

>>> get_sender_ctx() is None

True

"""

return getattr(_local, "sender_ctx", None)

def _resolve_ctx(

explicit_ctx: SessionContext | None = None,

) -> SessionContext:

"""Resolve a context for Expr reconstruction.

Priority: explicit argument > worker context > global context.

Falling back to the global :class:`SessionContext` (instead of a

freshly constructed one) preserves any registrations the user has

installed on it.

Examples:

>>> from datafusion import SessionContext

>>> from datafusion.ipc import _resolve_ctx, clear_worker_ctx

>>> clear_worker_ctx()

>>> isinstance(_resolve_ctx(), SessionContext)

True

>>> ctx = SessionContext()

>>> _resolve_ctx(ctx) is ctx

True

"""

if explicit_ctx is not None:

return explicit_ctx

worker = get_worker_ctx()

if worker is not None:

return worker

# Lazy import: `datafusion/__init__.py` imports `datafusion.ipc`

# before `datafusion.context`, so a module-top import would force

# `datafusion.context` to load mid-init of `datafusion.ipc`. The

# cycle is benign today (context.py only pulls expr.py at module

# scope, neither pulls ipc.py back), but a single new import in

# context.py's transitive deps could turn it into a real cycle.

# Deferring keeps `datafusion.ipc` import-order-independent.

from datafusion.context import SessionContext # noqa: PLC0415

return SessionContext.global_ctx()