)]}' { "commit": "baef8f00e1de6a086fa1814c98e65d216fb16277", "tree": "53d7aaca6eeebed0d0f6662343b9b04a4dba6222", "parents": [ "55870ff30c4a1f086e0b63de434ebf8b674ac110" ], "author": { "name": "Tim Saucer", "email": "timsaucer@gmail.com", "time": "Fri May 15 09:00:21 2026 -0400" }, "committer": { "name": "GitHub", "email": "noreply@github.com", "time": "Fri May 15 09:00:21 2026 -0400" }, "message": "Add support for logical and physical codecs (#1541)\n\n* feat: unify logical + physical proto codec stack via SessionContext\n\nIntroduces a single composable codec layer that every serializer reads\nfrom the session, replacing the hardcoded `DefaultLogicalExtensionCodec`\n/ `DefaultPhysicalExtensionCodec` calls scattered across PyLogicalPlan,\nPyExecutionPlan, and the Rust-wrapped Python provider plumbing.\n\nKey changes:\n\n* New `PythonLogicalCodec` and `PythonPhysicalCodec` (crates/core/src/codec.rs)\n wrap any inner `LogicalExtensionCodec` / `PhysicalExtensionCodec`. Both\n share a `DFPYUDF1` magic-prefix path for in-band cloudpickle encoding\n of Python scalar UDFs, so an `ExecutionPlan` / `PhysicalExpr`\n referencing a Python `ScalarUDF` round-trips through either layer.\n Magic-prefix registry table (DFPYUDF1 in use; DFPYUDA1 / DFPYUDW1 /\n DFPYPE1 reserved) documented in the module header.\n\n* `PySessionContext` stores `Arc\u003cPythonLogicalCodec\u003e` and\n `Arc\u003cPythonPhysicalCodec\u003e` directly. FFI wrappers are built on demand\n via `ffi_logical_codec()` / `ffi_physical_codec()` for capsule export\n and downstream `RustWrappedPy*` consumers. Adds\n `__datafusion_physical_extension_codec__` getter +\n `with_physical_extension_codec` setter (symmetric with the logical\n pair).\n\n* `PyLogicalPlan.to_proto` / `from_proto` renamed to `to_bytes` /\n `from_bytes`, now reading the session\u0027s logical codec. `to_proto` /\n `from_proto` survive as deprecated thin wrappers emitting\n `DeprecationWarning`.\n\n* `PyExecutionPlan` gains the same `to_bytes` / `from_bytes` rename +\n deprecated aliases, plus `__datafusion_execution_plan__` capsule\n getter and `from_pycapsule` (ported from poc_ffi_query_planner).\n\n* New `PyPhysicalExpr` class with `to_bytes` / `from_bytes` /\n `from_pycapsule` / `__datafusion_physical_expr__`. `from_bytes`\n takes an input pyarrow Schema for column-reference resolution.\n\n* `datafusion-python-util` gains `from_pycapsule!` /\n `try_from_pycapsule!` macros + `physical_codec_from_pycapsule`,\n `task_context_from_pycapsule`, `create_physical_extension_capsule`\n (ported from poc_ffi_query_planner).\n\n* `PythonFunctionScalarUDF` exposes `func()`, `input_fields()`,\n `return_field()`, `volatility()`, `from_parts()` accessors needed\n by the codec.\n\nPython wrapper updates: `LogicalPlan` / `ExecutionPlan` add\n`to_bytes` / `from_bytes` + deprecate `to_proto` / `from_proto`;\n`ExecutionPlan` adds capsule getter + `from_pycapsule`; new\n`PhysicalExpr` wrapper class exported from the top-level package;\n`SessionContext` exposes the physical codec capsule + setter.\n\nTest coverage in python/tests/test_plans.py: round-trip via new API,\ndeprecation warnings on old API, capsule protocol getters,\nsession-routed codec on both layers.\n\n`PyLogicalPlan` PyCapsule protocol is intentionally not added —\n`datafusion-ffi` does not expose `FFI_LogicalPlan`, so there is no\nstable cross-crate shape to publish. Round-tripping a `LogicalPlan`\ngoes through `to_bytes` / `from_bytes` only.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e\n\n* test: FFI-example integration tests for codec + plan capsule APIs\n\nAdds four downstream-crate fixtures in `datafusion-ffi-example` so the\nnew PR1 surface can be tested with the same FFI-handoff pattern used\nfor table providers, UDFs, etc. Existing tests prove the API exists;\nthese tests prove it composes with code that lives in another crate.\n\nNew Rust types in `examples/datafusion-ffi-example/src/`:\n\n* `MyLogicalExtensionCodec` — delegates to\n `DefaultLogicalExtensionCodec` and bumps atomic counters on the UDF\n encode/decode entry points. Exported via\n `__datafusion_logical_extension_codec__`. Installed onto a session\n with `ctx.with_logical_extension_codec(my_codec)`.\n* `MyPhysicalExtensionCodec` — mirror for `PhysicalExtensionCodec`.\n* `MyExecutionPlan` — wraps a one-column `EmptyExec`, exposes\n `__datafusion_execution_plan__`. Lets the receiver consume an\n `ExecutionPlan` capsule that did not originate in\n datafusion-python.\n* `MyPhysicalExpr` — wraps `Literal(Int32(42))`, exposes\n `__datafusion_physical_expr__`. Same FFI handoff for physical\n expressions.\n\nNew tests:\n\n* `_test_logical_extension_codec.py` — codec installs cleanly, the\n session re-exports its capsule, and `try_encode_udf` fires on the\n user codec when serializing a plan that references a `ScalarUDF`.\n The decode counterpart is a round-trip check rather than a counter\n assertion: when the UDF is in the receiver\u0027s function registry,\n `parse_expr` resolves by name before consulting the codec.\n* `_test_physical_extension_codec.py` — symmetric.\n* `_test_execution_plan.py` — parametrized over typed-class vs\n raw-capsule input; verifies `ExecutionPlan.from_pycapsule` consumes\n the downstream capsule.\n* `_test_physical_expr.py` — same for `PhysicalExpr.from_pycapsule`.\n\nAPI changes forced by the new tests:\n\n* `PyLogicalPlan.to_bytes`, `PyExecutionPlan.to_bytes`,\n `PyPhysicalExpr.to_bytes` now accept an optional `ctx` parameter.\n When supplied, encoding routes through the session\u0027s installed\n codec instead of a fresh default. `ctx\u003dNone` preserves the previous\n default-codec behavior used by the deprecated `to_proto` shims.\n* The util `from_pycapsule!` / `try_from_pycapsule!` macros now\n validate the capsule name via `pointer_checked(Some(c\"...\"))`\n rather than `pointer_checked(None)`. The latter rejects named\n capsules outright with CPython\u0027s \"incorrect name\" error.\n* `SessionContext.with_logical_extension_codec` and\n `with_physical_extension_codec` now wrap the returned internal\n context in `SessionContext` so the result has the full Python\n surface. The pre-existing logical setter was returning a raw\n internal object that lacked `sql()` and friends.\n\n`examples/datafusion-ffi-example/Cargo.toml` gains `datafusion` and\n`datafusion-proto` workspace dependencies for the new Rust impls.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e\n\n* refactor: tighten PR1 scope to codec plumbing only\n\nReview feedback pass. PR1 is now strictly the composable codec layer +\nsession routing + class-method serialization API. Anything that\ntouches actual Python UDF inline encoding or Python expression\nwrapping moves to PR2 alongside the pickle work.\n\nDropped:\n\n* `encode_python_scalar_udf` / `decode_python_scalar_udf` helpers\n from `crates/core/src/codec.rs`, along with cloudpickle and pyarrow\n imports. The wrapper codecs now delegate every method to `inner`.\n `DFPYUDF1` magic constant is kept (marked `dead_code` for now) as a\n reservation so PR2 has a single definition site.\n* `udf.rs` reverted to pre-PR1 shape. The codec no longer needs\n `func()` / `input_fields()` / `volatility()` / `from_parts()`\n accessors. Re-added by PR2 when scalar-UDF inlining lands.\n* `PyPhysicalExpr` class + Python wrapper + `__init__` export +\n `MyPhysicalExpr` FFI fixture + `_test_physical_expr.py`. No\n consumer in PR1 or PR2 plan documents; symmetry with\n `PyExecutionPlan` is not enough to justify the surface area.\n* Rust-side `PyLogicalPlan::to_proto` / `from_proto` and\n `PyExecutionPlan::to_proto` / `from_proto` deprecated wrappers.\n The deprecation lives entirely in the Python wrapper layer, which\n emits `DeprecationWarning` and forwards to `to_bytes` /\n `from_bytes`. Less Rust duplication.\n* `PythonLogicalCodec::with_default_inner` /\n `PythonPhysicalCodec::with_default_inner` — redundant with\n `impl Default`. Logic moved into `Default::default`.\n* `PySessionContext::default_logical_codec` /\n `default_physical_codec` helpers. Inlined as\n `Arc::new(PythonLogicalCodec::default())` at the three call sites.\n\nTests (root: 1076, FFI example: 36) all green after the cuts.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e\n\n* remove unuseful code comments\n\n* docs: rewrite codec module comments around purpose, not PR sequence\n\nThe previous doc-block framed PythonLogicalCodec / PythonPhysicalCodec\nin terms of \"PR1 delegates, PR2 will add encoding\" — useful for\nreview, useless for someone reading the code later.\n\nReframed in terms of what the codecs exist to do: encode Python-side\nplan references (pure-Python UDFs, etc.) into the proto wire format\nso plans can cross process boundaries without the receiver having to\npre-register every callable. The wrappers sit at the top of the\nsession\u0027s codec stack and delegate non-Python encoding to a\ncomposable inner codec.\n\nMagic-prefix registry table loses the \"reserved\" column. Doc still\nnotes that the in-module impls currently delegate and that\nencoder/decoder hooks land alongside the corresponding Python-side\nserialization work.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e\n\n* feat(codec): forward every LogicalExtensionCodec /\nPhysicalExtensionCodec method to inner\n\nPythonLogicalCodec previously only overrode the four required methods\non the trait plus the scalar UDF pair, so the default trait impls\n(returning \"LogicalExtensionCodec is not provided\") shadowed any\ndownstream FFI codec for file formats, aggregate UDFs, and window\nUDFs. A user installing their own codec via\n`SessionContext.with_logical_extension_codec(...)` would silently\nlose access to its `try_*_file_format`, `try_*_udaf`,\n`try_*_udwf` implementations.\n\nForward every trait method to `inner` so the user-installed codec is\nfully reachable. Same change on the physical side, including\n`try_*_expr`, `try_*_udaf`, `try_*_udwf` — the corresponding\nPython-aware paths can layer on later by intercepting before\ndelegation.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e\n\n* docs: tighten codec dispatch test docstrings\n\nThe previous docstrings claimed the tests verify \"PythonLogicalCodec\ndelegates non-Python UDFs to the inner codec.\" That\u0027s\nforward-looking — the codecs currently delegate every UDF\nunconditionally, so the test would behave identically for Python and\nnon-Python UDFs.\n\nRewrite to describe what the test actually proves: the dispatch chain\n`PyLogicalPlan.to_bytes -\u003e session.logical_codec -\u003e PythonLogicalCodec\n-\u003e FFI -\u003e user impl` (and the physical mirror) forwards correctly,\nobservable via the user codec\u0027s atomic counter incrementing after one\nencode pass.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e\n\n* refactor(ffi-example): MyExecutionPlan emits real data via\nMemorySourceConfig\n\nWas a one-column `EmptyExec` stub useful only as a capsule-handoff\ntarget. Promoted to a minimal reference impl that a downstream Rust\ncrate can copy when exposing a custom `ExecutionPlan` to\ndatafusion-python: configurable `num_rows`, produces a single batch\nof sequential `Int32` values under column `value`, wrapped in\n`DataSourceExec` via `MemorySourceConfig::try_new_exec`. Header\ncomment explains the typical use case (remote backend, streaming\nsource, synthetic data generator) and the\n`__datafusion_execution_plan__` capsule shape downstream crates\nshould follow.\n\nTest asserts the schema-bearing plan survives the FFI hop: a\n`DataSourceExec` arrives with the expected partitioning and no\nchildren. Schema details are not surfaced through the FFI display\npath (only the wrapping `ForeignExecutionPlan` name + inner plan\nname appear), so the test does not assert the column name.\n\n`to_bytes` round-trip of an FFI-imported plan is not exercised:\nencoding requires a physical codec that knows how to serialize\n`ForeignExecutionPlan`, which the default codec does not. A\ndownstream user round-tripping such a plan must install their own\ncodec via `with_physical_extension_codec`. Documented in the test\nfile rather than asserted on.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e\n\n* refactor: drop dormant ExecutionPlan PyCapsule round-trip\n\n`PyExecutionPlan::from_pycapsule` and the matching\n`__datafusion_execution_plan__` exporter have no consumer in this\nrepo, on the POC `poc_ffi_query_planner` branch, or on any sibling\nbranch (`testing/datafusion-distributed`, `testing/ffi-library-marker`,\n`tmp/ffi-with-codecs`). The pair was wired up speculatively for FFI\nplan handoff that no Python code path actually performs today.\n\nDrop the whole capsule round-trip for `ExecutionPlan`:\n\n* Rust `PyExecutionPlan::from_pycapsule` and\n `__datafusion_execution_plan__`.\n* Python `ExecutionPlan.from_pycapsule` and\n `__datafusion_execution_plan__` wrappers.\n* `MyExecutionPlan` FFI fixture + `_test_execution_plan.py` + lib.rs\n registration. Was solely a test fixture for the dropped path.\n* `test_execution_plan_pycapsule_protocol` in `python/tests/test_plans.py`.\n\n`PyExecutionPlan.to_bytes` / `from_bytes` survive — they encode\nthrough the session\u0027s physical codec and have real coverage.\nCapsule round-trip can be re-added when a concrete consumer\n(distributed worker, bridge library) lands.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e\n\n* feat: PyExpr.to_bytes / from_bytes via session logical codec\n\nMirrors PyLogicalPlan / PyExecutionPlan: encode through the session\u0027s\ninstalled `LogicalExtensionCodec` (or a default-inner\n`PythonLogicalCodec` when no `ctx` is supplied), decode against the\nsession\u0027s function registry + codec via `parse_expr`.\n\nRust side calls `datafusion_proto::logical_plan::to_proto::serialize_expr`\nand `from_proto::parse_expr`. Python wrapper threads an optional\n`SessionContext` through.\n\nTests cover the session-routed roundtrip and the no-ctx default-codec\nencode path. Adds a third consumer of `session.logical_codec()`\nalongside `PyLogicalPlan` and the codec dispatch tests in the FFI\nexample, broadening coverage of the codec stack.\n\nThis is the last piece of the PR1 codec surface — follow-up pickle\nwork (`Expr.__reduce__`, worker-scoped context, multiprocessing) can\nbuild on this without bundling the byte-level serialization API.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e\n\n* test(ffi-example): assert codec roundtrip restores plan output\n\nPR review feedback: weak `is not None` checks let regressions slip\npast. Mirror python/tests/test_plans.py — logical compares\n`df.collect() \u003d\u003d round_trip.collect()`; physical compares\n`str(original) \u003d\u003d str(restored)`.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e\n\n---------\n\nCo-authored-by: Claude Opus 4.7 (1M context) \u003cnoreply@anthropic.com\u003e", "tree_diff": [ { "type": "modify", "old_id": "70f09ec467605da3948c2ac8f03f3df73301a18b", "old_mode": 33188, "old_path": "Cargo.lock", "new_id": "1d148b0e19d61296f2aa4e6dd35f7c41ee99030f", "new_mode": 33188, "new_path": "Cargo.lock" }, { "type": "add", "old_id": "0000000000000000000000000000000000000000", "old_mode": 0, "old_path": "/dev/null", "new_id": "088532df2b39f8dfc9eeb787c271ddd76afc3f7b", "new_mode": 33188, "new_path": "crates/core/src/codec.rs" }, { "type": "modify", "old_id": "e46d359d6f5a3b57b7967093e8baee44a679771d", "old_mode": 33188, "old_path": "crates/core/src/context.rs", "new_id": "96de01889e0210e351cc0047a30640911ee70e39", "new_mode": 33188, "new_path": "crates/core/src/context.rs" }, { "type": "modify", "old_id": "c4f2a12da412095ccba099a910aac0c7426f0e0b", "old_mode": 33188, "old_path": "crates/core/src/expr.rs", "new_id": "2e633baeb43b5b4ad988b2f20126ea7387dd219f", "new_mode": 33188, "new_path": "crates/core/src/expr.rs" }, { "type": "modify", "old_id": "77d69911adda41ad886cd0ed8169c014795dcfc0", "old_mode": 33188, "old_path": "crates/core/src/lib.rs", "new_id": "e3551c937e1fc9aa7e7e64a003256cff850e29e7", "new_mode": 33188, "new_path": "crates/core/src/lib.rs" }, { "type": "modify", "old_id": "fac973884bd2bc7d441ae0c114c56c50a7f756e7", "old_mode": 33188, "old_path": "crates/core/src/physical_plan.rs", "new_id": "594655a6093c73e84a6b7fc65545ad12f93985a5", "new_mode": 33188, "new_path": "crates/core/src/physical_plan.rs" }, { "type": "modify", "old_id": "631aa9b0902d47a80510c83f9c7cbba6ad00ffcc", "old_mode": 33188, "old_path": "crates/core/src/sql/logical.rs", "new_id": "647c3fa7e925005c37e42e49ff9d53a88a9db44b", "new_mode": 33188, "new_path": "crates/core/src/sql/logical.rs" }, { "type": "modify", "old_id": "00d5946a5ef2de9635006e471a67d2ce42a74be8", "old_mode": 33188, "old_path": "crates/util/Cargo.toml", "new_id": "c23667b0f5587f2ba7d7625ba97e9d2f014bd2c5", "new_mode": 33188, "new_path": "crates/util/Cargo.toml" }, { "type": "modify", "old_id": "5b1c8993681c909ad6d57b15f544541ad8250414", "old_mode": 33188, "old_path": "crates/util/src/lib.rs", "new_id": "72dc9aafc5d3c269fbe9006ad4273201d39f12e7", "new_mode": 33188, "new_path": "crates/util/src/lib.rs" }, { "type": "modify", "old_id": "178dce9f97d7958e81240ed88d4ef9839fa5db09", "old_mode": 33188, "old_path": "examples/datafusion-ffi-example/Cargo.toml", "new_id": "ffc839d56341a43f051e23765d8c611b1304d61a", "new_mode": 33188, "new_path": "examples/datafusion-ffi-example/Cargo.toml" }, { "type": "add", "old_id": "0000000000000000000000000000000000000000", "old_mode": 0, "old_path": "/dev/null", "new_id": "cd0c5a61ae585a94bbecf75d680cdcba3b7ba1ee", "new_mode": 33188, "new_path": "examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py" }, { "type": "add", "old_id": "0000000000000000000000000000000000000000", "old_mode": 0, "old_path": "/dev/null", "new_id": "28eaaf2d92f8f293f91ed2c43a5f0b6d95b48310", "new_mode": 33188, "new_path": "examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py" }, { "type": "modify", "old_id": "e708c49cc810c20485b6cf851abfd247929302a8", "old_mode": 33188, "old_path": "examples/datafusion-ffi-example/src/lib.rs", "new_id": "3323ac982bd887c87cbd16d05d1971546ed7cac8", "new_mode": 33188, "new_path": "examples/datafusion-ffi-example/src/lib.rs" }, { "type": "add", "old_id": "0000000000000000000000000000000000000000", "old_mode": 0, "old_path": "/dev/null", "new_id": "da9efb297d5f475f8a49a10125c8844334b4666c", "new_mode": 33188, "new_path": "examples/datafusion-ffi-example/src/logical_extension_codec.rs" }, { "type": "add", "old_id": "0000000000000000000000000000000000000000", "old_mode": 0, "old_path": "/dev/null", "new_id": "b1a586d9e3db5662ccd3fc2ee80d4713dfcbe8c1", "new_mode": 33188, "new_path": "examples/datafusion-ffi-example/src/physical_extension_codec.rs" }, { "type": "modify", "old_id": "dd6790402cf36e22048b2d9b83f22127ca0e38c0", "old_mode": 33188, "old_path": "python/datafusion/context.py", "new_id": "5c3501941a70c8dc86ac36016078c62509ed3860", "new_mode": 33188, "new_path": "python/datafusion/context.py" }, { "type": "modify", "old_id": "0f7f3ab5aebd0974b9b28221bfdabec191b48732", "old_mode": 33188, "old_path": "python/datafusion/expr.py", "new_id": "e0135e3edf0ca4f59bd1bdfaa24337c4cce1baaa", "new_mode": 33188, "new_path": "python/datafusion/expr.py" }, { "type": "modify", "old_id": "c0cfd523fa88094f6073f19d17967ebdf3dee300", "old_mode": 33188, "old_path": "python/datafusion/plan.py", "new_id": "b2c6eab3e17a88a0251e63f234608223cee81342", "new_mode": 33188, "new_path": "python/datafusion/plan.py" }, { "type": "modify", "old_id": "8aa791ae16853fa4892362726ca5ba2e037eb548", "old_mode": 33188, "old_path": "python/tests/test_expr.py", "new_id": "6a466f6f22d1ae7a8f1b3717140643332246720c", "new_mode": 33188, "new_path": "python/tests/test_expr.py" }, { "type": "modify", "old_id": "3705fc7ef64e40dd6988d9c71da85fe24cec1000", "old_mode": 33188, "old_path": "python/tests/test_plans.py", "new_id": "11e709f6b7065d18412de7de8a0a4289d9986128", "new_mode": 33188, "new_path": "python/tests/test_plans.py" } ] }