Commit ee2e011
committed
Add Substrait bridge to apply grad() rewrite in the Python context
Add the `grad` marker UDF and a plan-level rewriter (`rewrite_grad_calls`)
to the autograd engine, plus a `grad_rewrite` PyO3 function that bridges the
differentiation engine into the datafusion-python SessionContext.
Because the native extension links its own copy of DataFusion, expressions
cross the Python<->Rust boundary as Substrait protobuf. Python produces the
logical plan as Substrait; `grad_rewrite` consumes it into a DataFusion
LogicalPlan, rewrites every `grad(expr, column)` ScalarFunction into the
symbolic derivative via `differentiate`, and re-produces Substrait bytes for
Python to consume and execute. The custom xarray table provider round-trips
because Substrait serializes table scans by name (resolved against the
registry on consume), so the rewrite context only needs empty tables with
matching schemas.
`grad` is registered as a marker ScalarUDF that carries the differentiation
request intact through parsing, planning, and serialization; it is always
rewritten away before execution and errors if it ever reaches invoke.
Deps: datafusion-substrait 52 and prost 0.14 (matching the substrait crate).
Building now requires `protoc` (the substrait crate codegens from .proto).
Verified end to end (produce -> grad_rewrite -> consume -> execute) against
analytic derivatives for cos, the product rule, and exp with 0.0 error.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017mDoFJgsm9kS7SicGoCVF61 parent 6838859 commit ee2e011
4 files changed
Lines changed: 498 additions & 27 deletions
0 commit comments