diff --git a/docs/source/conf.py b/docs/source/conf.py index 14cdc58d8..efb547652 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -122,10 +122,7 @@ def autoapi_skip_member_fn(app, what, name, obj, skip, options) -> bool: # noqa # Re-exports ("class", "datafusion.DataFrame"), ("class", "datafusion.SessionContext"), - ("class", "datafusion.QueryPlannerExportable"), ("class", "datafusion.SessionExtensionComponents"), - ("class", "datafusion.SessionComponentsExportable"), - ("class", "datafusion.SessionPlannerExportable"), ("module", "datafusion.common"), # Duplicate modules (skip module-level docs to avoid duplication) ("module", "datafusion.col"), diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md index 75c32444c..bbf6b65e7 100644 --- a/docs/source/contributor-guide/ffi-internals.md +++ b/docs/source/contributor-guide/ffi-internals.md @@ -111,6 +111,54 @@ library would serialize, and would do it with the codecs it was imported with. The extension-facing consequence — install codecs before a layered planner, and prefer `with_extensions` — is documented at {ref}`planner_codec_rebinding`. +(ffi_internals_commit_order)= + +## Why `with_extensions` commits last + +`with_extensions` promises that a bundle which raises leaves the session as it +was, apart from anything a hook writes to the context it is handed +({ref}`extension_bundles_transaction`). Keeping that promise is an ordering constraint on the implementation, not +a property of any one step, because the planner is bound on the shared +`SessionState` rather than on the returned handle. + +A call therefore splits into four steps, of which only the last writes: + +1. **Collect.** Every `__datafusion_session_components__` runs and its codecs + are gathered. Nothing is installed yet, so a hook that raises here has + touched nothing. +2. **Chains.** The codecs are assembled into the returned handle. Codec chains + live on that handle rather than on the session, so this step writes nothing + to the session even though it can fail on a bad capsule or a duplicate id. +3. **Resolve.** Every `__datafusion_session_planner__` runs, in argument order, + against the completed chains, and each supplied planner is exported to a + capsule. Every hook that can raise has run by the end of this step. +4. **Commit.** The accumulated planner is re-imported from its capsule and + bound, in a single `SessionState` rebuild. The bind is skipped entirely when + the call installed nothing, so an empty call does not drag a planner sitting + on another handle's codecs onto this one's. + +Only step 4 touches the session, and it is not itself infallible: +`_install_extension_planner` runs `ffi_query_planner_from_pycapsule` before it +calls `set_session_query_planner`, which cannot fail. The property that keeps +the promise is therefore about order, not about any step being incapable of +raising — every fallible operation, including the ones inside the commit, +completes before the first write. + +That is the rule for the next field added to `SessionExtensionComponents`, not +only a description of the current code: a new kind of component must do its +fallible work — importing a capsule, resolving a name — before anything is +written, so no failure can leave the session half-updated. + +There would be nothing to roll back to if one did. The returned handle shares one +session with the receiver, so the damage is visible from every other handle; +and undoing a registration is not the same as restoring what it displaced, +because deregistering a function that shadowed a built-in removes the built-in +too. The split is cheaper than an undo log that cannot be written correctly. + +The extension-facing statement of this is +{ref}`extension_bundles_transaction`, which says only that declaring a +component is safe where registering one during the hook is not. + ## Two argument kinds for one convention `CapsuleGetterArg` in `crates/util/src/lib.rs` distinguishes three cases: no diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 791fdfc8e..437deb165 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -281,6 +281,8 @@ for direct ones. The wrapper travels with the codec; the bundle does not. The query planner is exempt — it carries no wire id, so it may be an object or a capsule. +(extension_bundles_transaction)= + ## Failure and rollback Nothing is written to the session until every factory has returned and every @@ -290,6 +292,12 @@ table, say — is **not** rolled back, which is why bundle objects must be configuration-only: create fresh components on each call, never cache bound components, and do not retain the context passed in. +Declaring a component is what buys you that guarantee. Anything you return from +your hook is validated while a failure still costs nothing, and is written only +after every bundle in the call has succeeded. Anything you register yourself is +written immediately, before the other bundles have even run. The ordering that +makes this hold is recorded at {ref}`ffi_internals_commit_order`. + Like every other derivation, the returned context is a handle on the *same* session as the receiver — see {ref}`extension_sessions`. Only the Python-side codec chains belong to the returned handle; the planner is installed on the diff --git a/docs/source/llms.txt b/docs/source/llms.txt index 47b303ff6..597ce059d 100644 --- a/docs/source/llms.txt +++ b/docs/source/llms.txt @@ -31,7 +31,7 @@ - [`datafusion.expr`](https://datafusion.apache.org/python/autoapi/datafusion/expr/index.html): expression tree nodes (`Expr`, `Window`, `WindowFrame`, `GroupingSet`). - [`datafusion.functions`](https://datafusion.apache.org/python/autoapi/datafusion/functions/index.html): 290+ scalar, aggregate, and window functions. - [`datafusion.context.SessionContext`](https://datafusion.apache.org/python/autoapi/datafusion/context/index.html): session entry point, data loading, SQL execution. -- [`datafusion.extensions`](https://datafusion.apache.org/python/autoapi/datafusion/extensions/index.html): `SessionExtensionComponents`, `SessionComponentsExportable`, `SessionPlannerExportable`, `QueryPlannerExportable` — the extension-bundle protocol. +- [`datafusion.extensions`](https://datafusion.apache.org/python/autoapi/datafusion/extensions/index.html): `SessionExtensionComponents`, `SessionComponentsExportable`, `SessionPlannerExportable`, `QueryPlannerExportable`, `PhysicalOptimizerRuleExportable` — the extension-bundle and capsule-getter protocols. - [`datafusion.ipc`](https://datafusion.apache.org/python/autoapi/datafusion/ipc/index.html): worker and sender context slots for shipping expressions between processes. ## Examples diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index d60653a21..0b8e94bb5 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -159,6 +159,34 @@ installed produces the same bytes as before, as do functions encoded by name. Regenerate any plan you serialized with an earlier release and stored for later use, if it was produced by a session with an extension codec installed. +### Capsule-getter protocols moved to `datafusion.extensions` + +`PhysicalOptimizerRuleExportable` now lives in `datafusion.extensions`, next to +the other protocols an extension library implements against. It was previously +importable from `datafusion.context`, and that path is gone. + +```python +from datafusion.context import PhysicalOptimizerRuleExportable # before +from datafusion.extensions import PhysicalOptimizerRuleExportable # after +``` + +This affects type annotations only. The protocol is structural and not +`@runtime_checkable`, so nothing imports it to call `isinstance`, and +`SessionContext.add_physical_optimizer_rule` is unchanged — a rule object that +worked before still works, whether or not its library names the protocol +anywhere. + +The bundle protocols added in this release — `QueryPlannerExportable`, +`SessionComponentsExportable`, and `SessionPlannerExportable` — are reached the +same way, through `datafusion.extensions` rather than the package root. They are +new in 55.0.0, so no earlier import path existed. `SessionExtensionComponents` +stays at the root, because a bundle constructs one rather than merely naming it: + +```python +from datafusion import SessionExtensionComponents +from datafusion.extensions import SessionComponentsExportable +``` + ### `SessionContext.execute` renamed its second parameter The parameter is a single partition index, not a count, and is now named diff --git a/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py b/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py index 0c877d78e..9271c290a 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py +++ b/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py @@ -17,8 +17,13 @@ from __future__ import annotations +import doctest +import inspect +import io + import pyarrow as pa from datafusion import SessionContext +from datafusion.extensions import PhysicalOptimizerRuleExportable from datafusion_ffi_example import MyPhysicalOptimizerRule @@ -43,3 +48,42 @@ def test_ffi_physical_optimizer_rule_runs_during_planning(): f"before={before} after={after}" ) assert result[0].column(0).to_pylist() == [1, 2, 3] + + +def test_physical_optimizer_rule_docstring_example_still_runs(): + """Run the ``PhysicalOptimizerRuleExportable`` docstring example verbatim. + + The example is marked ``+SKIP`` because the main suite has no built FFI + extension to import, which is exactly how such an example rots. Here the + statements are parsed out of the live docstring, the skip is dropped, and + each one is executed. + + Only the ``+SKIP`` statements are taken. The docstring also carries a + runnable example above them, which the main suite already executes under + ``--doctest-modules``; running it again here would prove nothing. + + Only ``ctx`` is supplied, because the skipped statements go on using the + context the runnable block opened. Everything else resolves for real: a + renamed class, a changed signature, or a wrong expected output in the + docstring fails here. + """ + examples = [ + example + for example in doctest.DocTestParser().get_examples( + inspect.getdoc(PhysicalOptimizerRuleExportable) + ) + if example.options.pop(doctest.SKIP, False) + ] + assert examples, "docstring has no skipped examples to check" + + test = doctest.DocTest( + examples, + {"ctx": SessionContext()}, + "PhysicalOptimizerRuleExportable", + None, + None, + None, + ) + output = io.StringIO() + results = doctest.DocTestRunner().run(test, out=output.write) + assert results.failed == 0, output.getvalue() diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 1b44f8a73..76f647796 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -92,12 +92,7 @@ ) from .dataframe_formatter import configure_formatter from .expr import Expr, WindowFrame -from .extensions import ( - QueryPlannerExportable, - SessionComponentsExportable, - SessionExtensionComponents, - SessionPlannerExportable, -) +from .extensions import SessionExtensionComponents from .io import read_avro, read_csv, read_json, read_parquet from .options import CsvReadOptions from .plan import ( @@ -140,17 +135,14 @@ "ParquetColumnOptions", "ParquetWriterOptions", "PhysicalPartitioning", - "QueryPlannerExportable", "RecordBatch", "RecordBatchStream", "RuntimeEnvBuilder", "SQLOptions", "ScalarUDF", - "SessionComponentsExportable", "SessionConfig", "SessionContext", "SessionExtensionComponents", - "SessionPlannerExportable", "Table", "TableFunction", "TableProviderFactory", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 63cdfd487..4043e357b 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -58,6 +58,7 @@ import pyarrow as pa +from datafusion import extensions as _extensions from datafusion.catalog import ( Catalog, CatalogList, @@ -69,12 +70,6 @@ ) from datafusion.dataframe import DataFrame from datafusion.expr import sort_list_to_raw_sort_list -from datafusion.extensions import ( - QueryPlannerExportable, - SessionComponentsExportable, - SessionExtensionComponents, - SessionPlannerExportable, -) from datafusion.options import ( DEFAULT_MAX_INFER_SCHEMA, CsvReadOptions, @@ -99,6 +94,18 @@ from datafusion.catalog import CatalogProvider, Table from datafusion.common import DFSchema from datafusion.expr import Expr, SortKey + + # Type-only on purpose. `datafusion.extensions` is the one home for the + # capsule-getter protocols; importing these at runtime would make them + # reachable as `datafusion.context.*`, and for + # `PhysicalOptimizerRuleExportable` would restore the 54.0.0 path that + # 55.0.0 drops. Runtime checks go through the private `_extensions` alias. + from datafusion.extensions import ( + PhysicalOptimizerRuleExportable, + QueryPlannerExportable, + SessionComponentsExportable, + SessionPlannerExportable, + ) from datafusion.plan import ExecutionPlan, LogicalPlan from datafusion.user_defined import ( AggregateUDF, @@ -151,16 +158,6 @@ class TableProviderExportable(Protocol): def __datafusion_table_provider__(self, session: Any) -> object: ... # noqa: D105 -class PhysicalOptimizerRuleExportable(Protocol): - """Type hint for object that has __datafusion_physical_optimizer_rule__ PyCapsule. - - The method returns a PyCapsule wrapping an ``FFI_PhysicalOptimizerRule``, - typically produced by a separate compiled extension. - """ - - def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 - - class SessionConfig: """Session configuration options.""" @@ -1830,7 +1827,8 @@ def add_physical_optimizer_rule( Args: rule: Object exposing ``__datafusion_physical_optimizer_rule__``, - a :class:`PhysicalOptimizerRuleExportable`. + a + :py:class:`~datafusion.extensions.PhysicalOptimizerRuleExportable`. Examples: >>> from datafusion import SessionContext @@ -1918,7 +1916,8 @@ def with_extensions( every capsule has been validated, so a hook that raises leaves the session as it was. A hook that *mutates* the context it is handed — registering a table, say — is not rolled back, which is why bundle - objects must be configuration-only. + objects must be configuration-only. See + :ref:`extension_bundles_transaction`. Shares its session with this context — see :py:class:`SessionContext`. @@ -1977,7 +1976,11 @@ def with_extensions( """ for extension in extensions: if not isinstance( - extension, (SessionComponentsExportable, SessionPlannerExportable) + extension, + ( + _extensions.SessionComponentsExportable, + _extensions.SessionPlannerExportable, + ), ): msg = ( "Extension implements neither " @@ -1993,10 +1996,10 @@ def with_extensions( logical_codecs: list[LogicalExtensionCodecExportable] = [] physical_codecs: list[PhysicalExtensionCodecExportable] = [] for extension in extensions: - if not isinstance(extension, SessionComponentsExportable): + if not isinstance(extension, _extensions.SessionComponentsExportable): continue components = extension.__datafusion_session_components__(self) - if not isinstance(components, SessionExtensionComponents): + if not isinstance(components, _extensions.SessionExtensionComponents): msg = ( "__datafusion_session_components__ must return " "SessionExtensionComponents, got " @@ -2018,7 +2021,7 @@ def with_extensions( # rather than wrapping the session's default in an FFI hop. planner: _PyCapsule | None = None for extension in extensions: - if not isinstance(extension, SessionPlannerExportable): + if not isinstance(extension, _extensions.SessionPlannerExportable): continue fallback = ( planner @@ -2030,6 +2033,14 @@ def with_extensions( continue planner = new.ctx._export_query_planner(supplied) + # The commit step, and the only one that writes to the session. It can + # still raise -- `_install_extension_planner` re-imports the capsule + # before binding it -- so what keeps the promise is the order, not any + # step being incapable of failing: every fallible operation finishes + # before the first write. See docs/source/contributor-guide/ + # ffi-internals.md, "Why `with_extensions` commits last", for what a new + # component kind has to do to keep that true. + # # Rebinding the session's planner is a side effect on state shared with # every other handle, so do not pay it for a call that installs nothing # -- the same guard `with_python_udf_inlining` carries. With no codec diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 77ae92fc2..edb9a513e 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -32,11 +32,12 @@ implements either hook or both. Bundle order is significant for planners, which nest, and irrelevant for codecs, which accumulate. -Of the four names here, only the two bundle hooks are ``@runtime_checkable``, +Of the five names here, only the two bundle hooks are ``@runtime_checkable``, because :py:meth:`~datafusion.context.SessionContext.with_extensions` -dispatches on them from Python. :py:class:`QueryPlannerExportable` is a type -hint only, matching the other capsule-getter protocols in -:py:mod:`datafusion.user_defined` and :py:mod:`datafusion.catalog`. +dispatches on them from Python. :py:class:`QueryPlannerExportable` and +:py:class:`PhysicalOptimizerRuleExportable` are type hints only, matching the +other capsule-getter protocols in :py:mod:`datafusion.user_defined` and +:py:mod:`datafusion.catalog`. See :ref:`extension_bundles` in the online documentation for why the phases are split and for a worked implementation. @@ -57,6 +58,7 @@ ) __all__ = [ + "PhysicalOptimizerRuleExportable", "QueryPlannerExportable", "SessionComponentsExportable", "SessionExtensionComponents", @@ -64,6 +66,42 @@ ] +class PhysicalOptimizerRuleExportable(Protocol): + """Type hint for object that has __datafusion_physical_optimizer_rule__ PyCapsule. + + The method returns a PyCapsule wrapping an ``FFI_PhysicalOptimizerRule``, + typically produced by a separate compiled extension. It takes **no + argument**: a rule needs neither a codec nor a task-context provider, so + there is nothing session-scoped to hand it. + + Rules accumulate rather than replace. Install one with + :py:meth:`~datafusion.context.SessionContext.add_physical_optimizer_rule` + — see :ref:`extension_other_hooks`. + + Examples: + The getter is the whole protocol, and a capsule is what it must return + — anything else is refused where it is installed rather than at plan + time: + + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> ctx.add_physical_optimizer_rule(object()) + Traceback (most recent call last): + ... + RuntimeError: "Invalid datafusion_physical_optimizer_rule... + + Real usage. Skipped here (needs a built extension library); parsed out + of this docstring and run for real by + ``test_physical_optimizer_rule_docstring_example_still_runs`` in + ``datafusion-ffi-example``. + + >>> from datafusion_ffi_example import MyPhysicalOptimizerRule # doctest: +SKIP + >>> ctx.add_physical_optimizer_rule(MyPhysicalOptimizerRule()) # doctest: +SKIP + """ + + def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 + + class QueryPlannerExportable(Protocol): """Type hint for object that has a __datafusion_query_planner__ PyCapsule. @@ -90,7 +128,7 @@ class QueryPlannerExportable(Protocol): The protocol itself is not runtime-checkable: - >>> from datafusion import QueryPlannerExportable + >>> from datafusion.extensions import QueryPlannerExportable >>> try: ... isinstance(ctx, QueryPlannerExportable) ... except TypeError as e: @@ -110,7 +148,7 @@ def _not_a_codec_iterable(field: str, value: object) -> str: ) -@dataclass(frozen=True) +@dataclass(frozen=True, kw_only=True) class SessionExtensionComponents: """Components an extension contributes to a session context. @@ -121,6 +159,9 @@ class SessionExtensionComponents: components bound to a different session hold a task-context provider for that other session and cannot be rebound. + Construction is keyword-only, so later releases can add component kinds + without changing what an existing call means. + Query planners are not listed here. They install in a second phase so each can wrap the one before it — see :py:class:`SessionPlannerExportable`. @@ -233,7 +274,8 @@ class SessionComponentsExportable(Protocol): retain that context or cache the components they bound to it, since the next call may install onto a different session. They should also avoid mutating the context they are handed — a registration made during binding - is not rolled back if a later extension fails. + is not rolled back if a later extension fails. See + :ref:`extension_bundles_transaction`. A bundle that also contributes a query planner implements :py:class:`SessionPlannerExportable` alongside this protocol. @@ -250,10 +292,8 @@ class SessionComponentsExportable(Protocol): The codecs this bundle contributes. Examples: - >>> from datafusion import ( - ... SessionExtensionComponents, - ... SessionComponentsExportable, - ... ) + >>> from datafusion import SessionExtensionComponents + >>> from datafusion.extensions import SessionComponentsExportable >>> class MyLibraryExtension: ... def __datafusion_session_components__(self, ctx): ... return SessionExtensionComponents() @@ -314,7 +354,8 @@ class SessionPlannerExportable(Protocol): worth contrasting, because both plan queries successfully and only one of them is the no-op: - >>> from datafusion import SessionContext, SessionPlannerExportable + >>> from datafusion import SessionContext + >>> from datafusion.extensions import SessionPlannerExportable >>> class Contributes: ... def __datafusion_session_planner__(self, ctx, fallback): ... return None # the no-op: session keeps its own planner diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 9fbc4744f..274f2a9b1 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1086,6 +1086,21 @@ def test_session_extension_components_rejects_a_single_codec(field): SessionExtensionComponents(**{field: codec}) +def test_session_extension_components_is_keyword_only(): + codecs = ( + _NamedCodec( + SessionContext().__datafusion_logical_extension_codec__(), + "my_library.logical", + ), + ) + + with pytest.raises(TypeError, match=r"positional argument"): + SessionExtensionComponents(codecs) + + components = SessionExtensionComponents(logical_extension_codecs=codecs) + assert components.logical_extension_codecs == codecs + + def test_session_extension_components_rejects_a_string(): """A str is iterable, so it needs refusing on its own. diff --git a/python/tests/test_imports.py b/python/tests/test_imports.py index 0e0a3965f..e1ef0d38c 100644 --- a/python/tests/test_imports.py +++ b/python/tests/test_imports.py @@ -94,24 +94,6 @@ def test_datafusion_python_version(): assert datafusion.__version__ is not None -def test_extension_protocols_are_exported_together(): - """The extension protocol family is reachable from the package root. - - ``QueryPlannerExportable`` types the planner a - ``__datafusion_session_planner__`` hook returns, so a bundle author needs - it exactly as much as the other three; leaving it in the submodule made - one member of one family import differently from the rest. - """ - for name in [ - "QueryPlannerExportable", - "SessionExtensionComponents", - "SessionComponentsExportable", - "SessionPlannerExportable", - ]: - assert name in datafusion.__all__, f"{name} missing from datafusion.__all__" - assert getattr(datafusion, name) is getattr(datafusion.extensions, name) - - def test_class_module_is_datafusion(): # context for klass in [ @@ -210,6 +192,21 @@ def test_import_from_functions_submodule(): from datafusion.functions import foobar # noqa: F401 +@pytest.mark.parametrize( + "name", + [ + "PhysicalOptimizerRuleExportable", + "QueryPlannerExportable", + "SessionComponentsExportable", + "SessionPlannerExportable", + ], +) +def test_extension_protocols_import_from_extensions(name): + import datafusion.extensions + + assert getattr(datafusion.extensions, name).__module__ == "datafusion.extensions" + + def test_classes_are_inheritable(): class MyExecContext(SessionContext): pass