Skip to content

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

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

Add typed Protocols for FFI capsule exports (part of #1577)#1671
stantheman0128 wants to merge 1 commit 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.
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.

1 participant