[alxmrs]

🪭 Claude (einfold)

For the maintainer's request to integrate einfold with xarray-sql's demos. python/ is a Python package, einfold, that installs einfold in any session of datafusion-python (Apache DataFusion's Python bindings), and so in xarray-sql's Context and ddxdb's:

from datafusion import SessionContext
from einfold.datafusion import EinFold

einfold = EinFold()
ctx = SessionContext().with_extensions(einfold)   # every later query is planned by einfold
einfold.plans_with_einfold(), einfold.plans_with_pipelines()

The package is organized by engine: einfold.datafusion now, with datafusion an optional extra (pip install "einfold[datafusion]"), so that einfold.duckdb can sit beside it. The README is written for someone who first meets einfold on PyPI.

It's spike S29's bridge (docs/spikes/s29-python-planner), made fit to keep:

• Detection includes pipelines. The planner calls DetectFolds::detect, so demo 1's attention runs fused.
• The codecs no longer leak.
  • A table is held only from its encoding, on the session's side, to its decoding on einfold's, within one planning call. The spike kept every table forever, and holding them weakly doesn't work: the codec is handed a fresh FFI wrapper per encode.
  • Plans are held the same way, from encoding to decoding. A plan that has run keeps what its operators hold, such as a hash join's table, so keeping plans any longer kept every query's tables: in xarray-sql's GPT demo, about 2 GB per query.
• MAX is relinked as well as SUM, in aggregates and windows. Built-ins cross the FFI as foreign wrappers, and detection matches them by type, so the attention pipeline wasn't found until MAX was relinked too.
• The returned plan can be reset, for recursive queries, by resetting the plan inside it.
• No debug switches. The extension counts what it planned instead.
• Plans of any depth decode (fixes A logical plan deeper than 100 protobuf messages fails to decode, so einfold fails queries DataFusion answers (the GPT's loss) #77). The session hands each logical plan over as protobuf, and the decoder, prost, stopped at 100 nested messages, which a chain of 15 joins exceeds. The decoder is this library's copy of prost, inside DataFusion's FFI, so the crate turns on prost's no-recursion-limit feature. xarray-sql's GPT demo (in Add an ml extra for gradients in SQL, and two ML demos xarray-sql#269) now plans and trains with einfold installed, matching DataFusion's loss to 2e-16. Deeper plans then overflowed the stack, from about 250 joins, where DataFusion alone plans about 1,350. datafusion-proto decodes a plan, on the thread that asks for it, by recursing without growing its stack (DataFusion 55: datafusion-proto decodes a logical plan without growing its stack, so a deep plan crashes the process — draft for upstream #95). So installing einfold raises the process's soft stack limit to 64 MB, as ulimit -s would: Linux grows the main thread's stack on demand up to that limit. A chain of 1,300 joins now plans.

From Python (datafusion-python 55.0.0 from PyPI, i7-8700):

query	without einfold	with einfold
demo 1's causal attention, 1,024 tokens	2.07 s	17.5 ms
demo 1's causal attention, 4,096 tokens	109 s	46 ms
matrix product, 50,000 × 16 · 16 × 8	37.9 ms	20.9 ms
SELECT 1	0.23 ms	1.69 ms

The last row is the fixed cost of crossing the FFI, 1.5–3 ms per query. That's for datafusion-python and its FFI to fix (#58).

Packaging. The crate is outside the Cargo workspace, like the benchmarks. It needs unsafe to read datafusion-python's FFI capsules, which the workspace forbids, and pyo3 and DataFusion's FFI crates. Each unsafe block carries a SAFETY comment, and unsafe_op_in_unsafe_fn is denied. A new CI job builds the wheel with maturin, runs clippy on the crate, and runs the tests against datafusion-python 55.0.0 from PyPI.

Tests (python/tests/test_einfold.py): each query gives the same rows with einfold installed as without, to 1e-12, and einfold runs what it should:

• a matrix product (a fold), and demo 1's attention (a pipeline);
• a query with no fold, left alone;
• installing in one context and querying through the shared session;
• a query run again, with nothing left in the codec after each run, and a table replaced between runs;
• xarray-sql's Context with from_dataset, and ddxdb's grad, both skipped where the packages aren't installed;
• plans that are chains of 15 and 60 joins (A logical plan deeper than 100 protobuf messages fails to decode, so einfold fails queries DataFusion answers (the GPT's loss) #77), and of 350 in a subprocess, past the default stack. Another test pins the limit that remains: 350 joins planned on a thread with an 8 MB stack still crash, until datafusion-proto's decoding grows its stack (DataFusion 55: datafusion-proto decodes a logical plan without growing its stack, so a deep plan crashes the process — draft for upstream #95).

All 9 pass locally against xarray-sql's environment. CI has neither xarray-sql nor ddxdb, so it runs 7.

Next, in xarray-sql: an EinFold host in benchmarks/ml/_hosts.py, a DataFusion host whose context is xql.Context().with_extensions(EinFold()). That's a PR there, once xarray-sql#269 merges.

After CodeRabbit's review

• Each foreign FFI struct's version is checked before it's cloned: the session's query planner, and both codecs.
• The README has a Limits section. Very deep plans can overflow a thread's stack outside Linux's main thread. A function registered as sum or max is replaced by DataFusion's, since nothing crossing the FFI tells them apart (Keep a session's own sum and max in plans einfold plans from Python #111).
• The deep-plan test sets its core limit in the child, on Linux only, and is skipped elsewhere.

🤖 Generated with Claude Code

https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS

[coderabbitai]

Important

Review skipped

Review was skipped as selected files did not have any reviewable changes.

Files ignored due to path filters: 1

• python/Cargo.lock is excluded by !**/*.lock

Run configuration

• Source mode: Auto
• Repository: xqlsystems/einfold
• Branch: m2/python
• File: .coderabbit.yaml
• Commit: 90e78f4
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Advanced
• Run ID: d9c44e5c-db7c-4b81-9912-857418662adb

Commits

Reviewing files that changed from the base of the PR and between b6a7706 and 90e78f4.

Files ignored due to path filters: 1

• python/Cargo.lock is excluded by !**/*.lock

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

Walkthrough

Adds a Python package and native extension for DataFusion sessions. The extension provides einfold planning, FFI codecs, plan counters, and deep-plan stack handling. Python tests and CI cover wheel installation, query results, integrations, and deep join plans.

Changes: Python DataFusion integration

Layer / File(s)	Summary
Python package and session hooks python/Cargo.toml, python/pyproject.toml, python/einfold/*, python/src/lib.rs, python/README.md, python/Cargo.lock.license, .gitignore	Adds Python package configuration and native exports for DataFusion session components, planner hooks, codecs, plan counters, and the DataFusion version. The README documents supported query patterns and session integration.
FFI codecs and planner execution python/src/codec.rs, python/src/bridge.rs	Adds codecs that transfer table providers and physical plans by reference. The planner bridge relinks matching sum and max functions, tracks plan counts, plans through EinFoldPlanner, and wraps execution in the library runtime.
Deep-plan handling and validation python/src/stack.rs, python/tests/test_einfold.py, .github/workflows/ci.yml	Adds Linux stack-limit handling and tests for query results, repeated queries, optional integrations, and deep join plans. CI builds and installs a release wheel before running the Python tests.

Priority: Pending latest changes

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature · Severity of issue fixed: Medium

Sequence Diagram(s)

sequenceDiagram
  participant DataFusionSession
  participant EinFold
  participant Bridge
  participant EinFoldPlanner
  DataFusionSession->>EinFold: Request session planner
  EinFold->>Bridge: Create planner with shared counters
  DataFusionSession->>Bridge: Plan logical query
  Bridge->>EinFoldPlanner: Plan relinked query with session rules
  EinFoldPlanner-->>Bridge: Return physical plan
  Bridge-->>DataFusionSession: Return opaque execution plan

Loading

Possibly related issues

• xqlsystems/einfold#58 — Adds table and physical-plan codecs, function relinking, and plan handling for the DataFusion Python FFI gaps.
• xqlsystems/einfold#39 — Adds the Python package and session extension needed to use einfold with DataFusion.
• xqlsystems/einfold#111 — Relinks foreign sum and max functions to DataFusion’s built-ins.
• xqlsystems/einfold#115 — Adds plan counters to the Python planner extension.
• xqlsystems/einfold#95 — Removes Prost’s recursion limit and adds Linux stack-limit handling for deep plans.
• xqlsystems/einfold#84 — Concerns plan handling through the Python extension’s bridge.
• xqlsystems/einfold#77 — Removes Prost’s recursion limit to support deeper logical-plan decoding.

Possibly related PRs

• xqlsystems/einfold#40 — Pins the workspace to DataFusion 55.2.0, which this Python bridge also uses.
• xqlsystems/einfold#110 — Continues the Python package work using the same package, native modules, and bridge.
• xqlsystems/einfold#37 — Adds the fold-planning and execution types that the Python bridge plans and executes.

Merge Risk: Moderate · previous review c6c5e

Custom aggregates can return incorrect results, and failed queries can retain table data. These risks should be addressed before merging.

Security architecture risk: Moderate · up to c6c5e

Installing the extension makes query planning part of the host process’s availability boundary. Deep queries can terminate the process on some threads, and interrupted plan transfers lack a complete resource-release guarantee. Security exposure depends on whether the application accepts untrusted SQL and shares a process across users.

Retained concerns

• Medium · security · inferred: The new planner routes session queries through recursive FFI decoding while disabling the protobuf recursion limit. The documented worker-thread and macOS behavior includes process termination for sufficiently deep plans. Raising Linux’s soft stack limit only accommodates the main thread; it does not turn excessive depth into a contained query error. If a host accepts less-trusted SQL, one enabled session can therefore threaten availability of the entire process. The negative subprocess test does not independently attribute the crash to the enabled planner because it runs the baseline first.
• Medium · reliability · inferred: Table and plan references are stored in process-global registries and released only when their identifiers are decoded. There is no per-transfer owner or abort cleanup. Published DataFusion 55.2.0 logical serialization encodes a LIMIT input before rejecting a non-literal limit, providing an encode-before-error path that cannot be cleaned inside Bridge. Such aborted transfers can accumulate identifiers and keep table providers alive after session replacement or disposal, weakening process-level resource containment. Supported Python-query reproduction and exact wheel dependency resolution remain unresolved; successful-query cleanup does not establish failure cleanup.

Security review details

Security Blast Radius

• inferred — The availability boundary is the hosting process, not an individual query or session: decoding can terminate the process, and reference registries are shared across sessions. A malicious-SQL attack requires an application to grant query submission to an enabled session; external endpoint, tenant, credential, and deployment exposure were not established.

Security Findings and Attack Paths

• inferred — The supported security concern is failure containment: attacker-selected query depth can reach recursive decoding through the installed planner. A second conditional path is failed serialization after reference registration, which can retain resources. Neither cross-session data theft nor privilege escalation was demonstrated.

Trust Boundaries and Controls

• observed — Capsule names and FFI version checks constrain the supported native interface. Reference decoding rejects malformed and unknown identifiers and consumes valid entries once, but identifiers have no session namespace. No inspected SQL path gives callers control over these encoded identifiers, so their predictability alone does not establish an authorization bypass.

Resilience and Maintainability Implications

• observed — Mutex protection serializes registry mutations and successful decoding releases entries. It does not associate entries with a transfer’s terminal state. Opaque prevents outbound traversal of inner physical nodes, but cannot protect inbound logical serialization that fails before the planner callback.

Hardening Proposals

• proposed — Give each transfer an owner spanning outer serialization through terminal decoding, with error and cancellation cleanup limited to that transfer’s identifiers. Clearing shared registries globally would interfere with concurrent transfers.
• proposed — For hosts accepting untrusted SQL, enforce a plan-complexity bound before recursive decoding or use decoding with bounded or growable stack behavior. Process isolation is an alternative containment boundary; a larger main-thread stack alone is not one.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 55.56% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly and concisely describes the main change: adding a Python einfold extension for datafusion-python sessions.
Description check		The description is detailed and on topic. It explains the user-facing API, design decisions, packaging, tests, performance results, limitations, and follow-up work. It does not use every template head…
Linked Issues check		[#77] requires einfold to plan logical plans deeper than prost's default 100-message recursion limit. python/Cargo.toml enables prost's no-recursion-limit feature. python/tests/test_einfold.py c…
Out of Scope Changes check		The changes support the stated DataFusion Python package and [#77]. The prost configuration and stack handling support deep logical-plan decoding. The bridge, codecs, counters, packaging, CI, README, …

Full details: Description check

Explanation

The description is detailed and on topic. It explains the user-facing API, design decisions, packaging, tests, performance results, limitations, and follow-up work. It does not use every template heading or include a completed issue-closing line, but the required information is mostly present.

Full details: Linked Issues check

Explanation

[#77] requires einfold to plan logical plans deeper than prost's default 100-message recursion limit. python/Cargo.toml enables prost's no-recursion-limit feature. python/tests/test_einfold.py covers join depths 15 and 60 and a 350-join subprocess case. The PR also reports successful planning at greater depth. The issue's fallback suggestion is optional, so its absence is not a gap.

Full details: Out of Scope Changes check

Explanation

The changes support the stated DataFusion Python package and [#77]. The prost configuration and stack handling support deep logical-plan decoding. The bridge, codecs, counters, packaging, CI, README, and integration tests support installing and validating EinFold in datafusion-python. No unrelated change is demonstrated by the supplied whole-PR summary.

Full details: Docstring Coverage

Explanation

Docstring coverage is 55.56% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 27 functions across 7 files.

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches: 💡 1

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[alxmrs]

Initial thoughts.

[alxmrs]

LGTM.