kosiew opened a new pull request, #1277: URL: https://github.com/apache/datafusion-python/pull/1277
## Which issue does this PR close? * Closes #1237 ## Rationale for this change The current `AggregateUDF.udaf` and `AggregateUDF.from_pycapsule` methods in the DataFusion Python API lack proper type hinting and handling for CPython `PyCapsule` objects. This omission causes static type checking tools (e.g., mypy) to fail when users register UDAFs originating from external providers such as `geodatafusion`, even though the runtime behavior functions correctly. This PR addresses the gap by explicitly supporting PyCapsule types both in type hints and runtime checks. By doing so, it improves type safety, developer experience, and code clarity while maintaining full backward compatibility. example from #1237 ```python from datafusion import SessionContext, udf, udaf from geodatafusion import native ctx = SessionContext() ctx.register_udaf(udaf(native.Extent()))``` ### Before ``` ❯ mypy examples/datafusion-ffi-example/python/tests/_test_type_checking.py examples/datafusion-ffi-example/python/tests/_test_type_checking.py:5: error: No overload variant matches argument type "Extent" [call-overload] ... ``` ### After ``` ❯ mypy examples/datafusion-ffi-example/python/tests/_test_type_checking.py Success: no issues found in 1 source file ``` ## What changes are included in this PR? * Added `TypeGuard` function `_is_pycapsule()` for lightweight PyCapsule type validation. * Introduced `_PyCapsule` proxy class for static typing compatibility in non-type-checking contexts. * Extended overloads in `AggregateUDF.__init__` and `AggregateUDF.udaf()` to include `AggregateUDFExportable | _PyCapsule` argument types. * Added stricter constructor argument validation for callable accumulators. * Updated `AggregateUDF.from_pycapsule()` to support direct PyCapsule initialization. * Refactored Rust `PyAggregateUDF::from_pycapsule()` logic to delegate PyCapsule validation to a new helper function `aggregate_udf_from_capsule()` for cleaner handling. ## Are these changes tested? Yes: * The existing UDAF registration and execution tests cover runtime functionality continue to pass ## Are there any user-facing changes? Yes, minor improvements: * Users can now register UDAFs directly using PyCapsule objects without encountering static type-checking errors. * Type hints and IDE autocompletion will now accurately reflect valid function signatures. These changes are fully backward-compatible and non-breaking for existing user code. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
