Add typed Protocols for FFI capsule exports (part of #1577) - #1671
Open
stantheman0128 wants to merge 1 commit into
Open
Add typed Protocols for FFI capsule exports (part of #1577)#1671stantheman0128 wants to merge 1 commit into
stantheman0128 wants to merge 1 commit into
Conversation
Adds TableFunctionExportable, ExtensionOptionsExportable, and TaskContextProviderExportable Protocol classes describing the PyCapsule dunder methods DataFusion's Rust side already expects, following the existing TableProviderExportable pattern. No runtime behavior changes.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Part of #1577. This PR covers 3 of the 6 items in that umbrella issue (see "Not included" below for why the rest are left out).
Rationale for this change
The FFI-pipeline typing audit referenced in #1577 found several places where DataFusion's FFI types are already imported on the Rust side, but the Python surface has no typed
Protocoldescribing the expected PyCapsule dunder method. Callers are left to read the Rust source or guess the signature.What changes are included in this PR?
Adds three
Protocolclasses, following the existingTableProviderExportable/PhysicalOptimizerRuleExportablepattern (name, docstring, and single dunder method), and points the corresponding parameter/return type hints at them:datafusion.user_defined.TableFunctionExportable: describes the__datafusion_table_function__(self, session)PyCapsule method already duck-typed viahasattrinTableFunction.__init__/TableFunction.udtf(thosehasattrchecks are the actual runtime dispatch and are left untouched). Thefuncparameter onTableFunction.__init__is now typed asCallable[..., Any] | TableFunctionExportable, and a newudtfoverload documents the FFI-capsule call shape.datafusion.context.ExtensionOptionsExportable: describes__datafusion_extension_options__(self), required bySessionConfig.with_extension(seecrates/core/src/context.rs,SessionConfig::with_extension).with_extension'sextensionparameter is now typed against it.datafusion.context.TaskContextProviderExportable: describes__datafusion_task_context_provider__(self).SessionContextalready exposes one of these incontext.py, and this protocol documents the shape for other extensions that want to decode or produce one, matching thetask_context_from_pycapsulehelper already incrates/util/src/lib.rs.Not included from #1577's item list
The remaining 3 items are execution-engine-level questions rather than typing changes, so I split each into its own issue for design discussion:
FFI_TableProvider,TableProvider(ABC)(item 2): wiring a plain-Python subclass through to the RustTableProvidertrait needs a newRustWrappedPyTableProvider. Design discussion: Wire FFI_TableProvider / TableProvider(ABC) through to the Rust TableProvider trait #1668FFI_TableProviderFactory,from_pycapsulehelper (item 5): the call site this item points at passes an argument to the dunder method, which the existingfrom_pycapsule!/try_from_pycapsule!macros don't support (call0()only). Discussion: from_pycapsule! / try_from_pycapsule! macros don't support dunder methods that take arguments #1669WindowUDF,ABCequivalent toAccumulator(item 6):WindowEvaluatoralready provides this capability; it's just notABCMeta-shaped because the requiredevaluate*method depends on a flag matrix. Confirming whether the literal ABC form is still wanted: WindowEvaluator already provides the pure-Python window UDF base class #1577 (item 6) asks for #1670Are there any user-facing changes?
No runtime behavior changes. This only adds type hints (new
Protocolclasses) and updates existing parameter/return annotations to reference them.Validation run locally:
uvx ruff@0.15.1 check python/datafusion/user_defined.py python/datafusion/context.py- passeduvx ruff@0.15.1 format --check python/datafusion/user_defined.py python/datafusion/context.py- passedgit diff --check- passedpython -m py_compileon both files - passed