Skip to content

Commit dc7bd86

Browse files
timsaucerclaude
andcommitted
docs: state that a declared function may not be a bare capsule
udf(), udaf() and udwf() accept a bare PyCapsule, but a declaration in SessionExtensionComponents must be a wrapper or an object exposing the getter. Say so on the field, and pin it with an FFI example test that shows the same capsule wrapping fine through udf(). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
1 parent 1c3c72b commit dc7bd86

2 files changed

Lines changed: 30 additions & 4 deletions

File tree

‎examples/datafusion-ffi-example/python/tests/_test_session_extension.py‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@
2121

2222
import pyarrow as pa
2323
import pytest
24-
from datafusion import SessionContext, SessionExtensionComponents
25-
from datafusion_ffi_example import MyFunctionExtension
24+
from datafusion import SessionContext, SessionExtensionComponents, udf
25+
from datafusion_ffi_example import IsNullUDF, MyFunctionExtension
2626

2727

2828
def _session():
@@ -96,6 +96,30 @@ def test_installing_the_library_twice_is_refused():
9696
ctx.udf("my_custom_is_null")
9797

9898

99+
def test_a_bare_capsule_is_not_a_declaration():
100+
"""A declared function is an object, even though ``udf`` takes a capsule.
101+
102+
The capsule itself is fine -- ``udf`` wraps it -- so the refusal is the
103+
declaration rule, not a bad capsule.
104+
"""
105+
ctx = SessionContext()
106+
capsule = IsNullUDF().__datafusion_scalar_udf__()
107+
108+
class CapsuleExtension:
109+
def __datafusion_session_components__(
110+
self, ctx: SessionContext
111+
) -> SessionExtensionComponents:
112+
return SessionExtensionComponents(udfs=(capsule,))
113+
114+
with pytest.raises(TypeError, match=r"__datafusion_scalar_udf__"):
115+
ctx.with_extensions(CapsuleExtension())
116+
117+
with pytest.raises(KeyError):
118+
ctx.udf("my_custom_is_null")
119+
120+
assert udf(capsule).name == "my_custom_is_null"
121+
122+
99123
def test_a_failure_after_the_hook_registers_nothing():
100124
"""The transaction covers functions imported across the FFI boundary too."""
101125
ctx = SessionContext()

‎python/datafusion/extensions.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -276,8 +276,10 @@ class SessionExtensionComponents:
276276
277277
Either a :py:class:`~datafusion.user_defined.ScalarUDF` or an object
278278
exposing ``__datafusion_scalar_udf__``, which is wrapped with
279-
:py:func:`~datafusion.udf` on the way in. The registered name comes from
280-
the function itself, not from this field.
279+
:py:func:`~datafusion.udf` on the way in. A bare ``PyCapsule`` is refused,
280+
though :py:func:`~datafusion.udf` accepts one; declare the object that
281+
exports it. The registered name comes from the function itself, not from
282+
this field.
281283
282284
One :py:meth:`~datafusion.context.SessionContext.with_extensions` call may
283285
not declare a name twice, whether the two claims come from one extension or

0 commit comments

Comments
 (0)