Skip to content

Add typed Protocols for FFI capsule exports (part of #1577) - #1671

Open
stantheman0128 wants to merge 4 commits into
apache:mainfrom
stantheman0128:fix/1577-ffi-typing-protocols
Open

stantheman0128 wants to merge 4 commits into
apache:mainfrom
stantheman0128:fix/1577-ffi-typing-protocols

Conversation

@stantheman0128

Copy link
Copy Markdown

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 Protocol describing 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 Protocol classes, following the existing TableProviderExportable / PhysicalOptimizerRuleExportable pattern (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 via hasattr in TableFunction.__init__ / TableFunction.udtf (those hasattr checks are the actual runtime dispatch and are left untouched). The func parameter on TableFunction.__init__ is now typed as Callable[..., Any] | TableFunctionExportable, and a new udtf overload documents the FFI-capsule call shape.
  • datafusion.context.ExtensionOptionsExportable: describes __datafusion_extension_options__(self), required by SessionConfig.with_extension (see crates/core/src/context.rs, SessionConfig::with_extension). with_extension's extension parameter is now typed against it.
  • datafusion.context.TaskContextProviderExportable: describes __datafusion_task_context_provider__(self). SessionContext already exposes one of these in context.py, and this protocol documents the shape for other extensions that want to decode or produce one, matching the task_context_from_pycapsule helper already in crates/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:

Are there any user-facing changes?

No runtime behavior changes. This only adds type hints (new Protocol classes) 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 - passed
  • uvx ruff@0.15.1 format --check python/datafusion/user_defined.py python/datafusion/context.py - passed
  • git diff --check - passed
  • python -m py_compile on both files - passed

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.
@stantheman0128

Copy link
Copy Markdown
Author

This PR and #1684 both show zero CI runs. The check suites are sitting at action_required, which looks like the first-time contributor approval gate rather than a problem with either branch.

Could a maintainer approve the workflow runs when you get a chance? I would rather have CI results attached before anyone spends review time on these.

stantheman0128 and others added 3 commits September 7, 2026 22:16
…rotocols

# Conflicts:
#	python/datafusion/context.py
Upstream moved QueryPlannerExportable into datafusion/extensions.py and
context.py now imports it. This branch still carried the old class body
from an earlier merge, so the two collided and ruff reported F811.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@stantheman0128

Copy link
Copy Markdown
Author

Thanks for merging main in. That surfaced an F811, and it came from the merge rather than from the original change: upstream moved QueryPlannerExportable into datafusion/extensions.py and context.py now imports it, while this branch still carried the old class body. Removed in 1a329ca.

I kept ExtensionOptionsExportable and TaskContextProviderExportable in context.py rather than moving them to extensions.py, since ArrowStreamExportable, ArrowArrayExportable, TableProviderExportable and PhysicalOptimizerRuleExportable still live there. Happy to move them if you would rather the FFI protocols all sit in one module.

Verified against what lint-python actually runs: ruff check python/, ruff format --check python/, and codespell --toml pyproject.toml at the pinned 2.4.1. All clean.

The workflow runs need approving again for this commit. Every push from a first-time contributor lands at action_required.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants