[cromwellian]

🛠️ Codex

Repeated training calls can prepare SQL grad/VJP/JVP statements once and write new Dataset values into existing input buffers. Fine-chunk registration counts partitions arithmetically and registers Arrow batches directly; extension handles retain Dataset metadata and invalidation state. Context(threads=...) exposes target parallel partitions, and physical reuse declines when dynamic filters retain value-dependent bounds.

This independently useful change starts the NanoGPT optimization stack; shared/compact execution, resident training and the evidence report follow separately.

Apple M5 Ultra measurements: one-cell-chunk registration 31.372→0.405 ms (77.44× for that specific case), update-and-sum 2.769→1.027 ms (2.70×), repeated-CTE preparation 1.709→1.307 ms (1.31×). Tiny-model warm time improved 1.11×, but its five-step total regressed due to setup; this does not claim universally faster startup.

Validation: 47 focused AD/preparation tests passed, covering current inputs, changing join-key ranges, old-result snapshots, coordinate rejection and shared Context invalidation. The complete downstream stack passed 714 non-integration tests. Changed Python files pass Ruff.

Review order: this PR → #302 (shared kernels) → #303 (compact arrays) → #305 (resident training). #304 contains the separate consolidated report and evidence.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

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

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn

[alxmrs]

🌀 Claude: Thank you. Most of this is ready, and some of it closes our own open work.

Strengths

• prepare() for queries that take a gradient resolves the TODO(#294) we left in Context.prepare.
• Counting chunk partitions with math.prod is simpler and much faster (31.4 → 0.4 ms for one-cell chunks).
• Registering record batches directly, Context(threads=...), and with_extensions keeping the Dataset metadata are all good. The tests for them check values and metadata through public APIs.

A regression to fix: plan reuse is off by default
_safe_physical_reuse turns off the physical-plan reuse from #287 unless all four dynamic-filter settings are off. DataFusion 55 turns all four on by default, so every default Context now plans again on each call. On an i7-8700 (DataFusion 55.0.0), a prepared join and sum over 128 × 256 rows took:

	ms per call
main	0.30
this PR	0.60

The new test test_join_key_data_columns_can_change_their_range passes on main without the guard, with and without einfold installed. With einfold on main, the plan can't be copied at all and is planned again on each call. So the stale bounds probably come from the codec state reset in einfold#156. I'd fix them there, by resetting the dynamic filters along with the rest, and drop the guard here. If the guard turns out to be needed, it should only apply to the session codecs that need it.

Scope

• UDF tracking. _custom_udfs, _trusted_math_udfs and the register_udf/register_udaf/register_udwf overrides are only used by Share SQL stages and dispatch validated numerical kernels #302 and later. I'd move them to the PR that uses them.
• ad.prepare. This adds a second public prepare API beside Context.prepare, with its own lock. I'd keep only Context.prepare, and have it accept several statements that share a loss if that's needed.

Housekeeping

• The TODO(#294) has gone; please put Closes #294 in the description.
• CI hasn't run here yet, since the PR comes from a fork. A maintainer needs to approve it.