[dbtsai]

What changes were proposed in this pull request?

Share immutable logical-plan and expression nodes instead of recursively cloning
their ancestors while building DataFrames and Columns.

• Store DataFrame roots and logical-plan child edges as Arc<LogicalPlan>.
• Make Expression an immutable shared handle around ExpressionNode.
  Common builders use Expression::literal, Expression::column, and
  Expression::function; other nodes use the From<ExpressionNode> conversion.
• Remove redundant boxes around expression handles in DirectShufflePartitionId
  and CaseWhen::else_expr.
• Preserve structural equality (including NaN), Debug/render output, protobuf
  output, argument order, and execution-time plan-ID assignment.
• Keep focused sharing, branch-immutability, serialization, and self-join plan-ID
  tests. Add a Rust API migration note in docs/compatibility.md.

CI follow-ups type-check documentation snippets without linking an executable per
block, constrain Spark 4.2 tests and workers to pandas 2.x, and retain per-file
pytest diagnostics. Benchmark code lives in
pyspark-connect-benchmark.

Why are the changes needed?

Extending and reusing a plan or expression currently copies existing subtrees.
Sharing immutable nodes avoids that repeated work. Expression and Column clones
become O(1) with respect to subtree size.

Fresh benchmark rerun on 2026-10-27 UTC, after the review cleanups, using
separately built release wheels from base c41ff82 and PR head ff6fddd3.
The separate
pyspark-connect-benchmark
harness is pinned to
1bbdec9.
These results supersede the earlier measurements of b45e5d6; old samples are
not pooled with this rerun.

The breadth stage covered 39 cases with six attempts per client/case. A separate
confirmation stage repeated eight selected cases with 18 attempts per client/case.
Each stage used two rounds with the Rust baseline/candidate execution slots
swapped. The tables compare the two Rust builds, not Rust against official Python.

Confirmation wall-time medians (18 samples per Rust version):

Case	Before	After	Result
chain-withColumn-256	7.969 ms	0.326 ms	24.46x faster
linear-expression-64	0.664 ms	0.084 ms	7.95x faster
increment-filter-depth256	40.728 us	0.241 us	168.97x faster
core-outbound-rpc-expression-10	15.044 ms	10.149 ms	32.5% less time
core-outbound-rpc-deep-10	5.557 ms	5.372 ms	3.3% less time
core-outbound-rpc-deep-100	6.329 ms	5.614 ms	11.3% less time
core-outbound-rpc-expression-5	6.488 ms	6.480 ms	About unchanged
prebuilt-schema-depth256	2.826 ms	2.857 ms	1.1% more time

The earlier 15.3% wall-time slowdown in core-outbound-rpc-deep-10 was not
reproduced: both confirmation rounds favor the PR, but their ranges overlap
widely, so this is not evidence of a statistically significant speedup.
The schema-depth256 timing remains mixed (-0.9% and +2.2% across rounds), rather
than a convincing timing regression. Expression-5 outbound also reverses
direction between rounds. These are not direct old-head/current-head comparisons,
so differences from the earlier run cannot be attributed solely to review cleanups.

Two cold-construction slowdown signals remain in the breadth stage:

Case	Before	After	Wall-time change	CPU-time change
core-construct-deep-10	5.189 ms	6.048 ms	+16.6%	About unchanged
core-construct-expression-5	4.987 ms	6.220 ms	+24.7%	+7.2%

Both have higher wall-time medians in both breadth rounds. Each has only six
samples per Rust version and overlapping ranges; neither was included in the
confirmation stage. They need targeted confirmation, and should not be dismissed
as noise or interpreted as established regressions.

Memory measurements (before -> after, MiB):

Case / stage	Absolute sampled peak RSS	Additional peak RSS
core-construct-expression-10 / breadth	105.148 -> 101.242	5.039 -> 1.289 (-74.4%)
core-outbound-rpc-expression-10 / confirmation	116.609 -> 112.695	14.219 -> 10.164 (-28.5%)
prebuilt-schema-depth256 / confirmation	104.727 -> 104.258	1.688 -> 2.031 (+352 KiB)

Additional peak RSS is the median of each sample's sampled peak minus its own
post-setup baseline, not the difference of two medians or total process memory.
Memory improvement is not universal: the schema case uses more incremental RSS
despite its slightly lower absolute process peak. RSS sampling targets 2 ms and
can miss short spikes; zero observed growth does not mean zero allocation.

These are client-side measurements against canned localhost gRPC responses on
an Apple M3 Max (64 GiB), macOS arm64, Python 3.12.14 and Rust 1.99.1, using
identical release build settings and dependencies for the two Rust versions.
Core one-shot cases include imports/first-use work inside timing; outbound cases
also include plan construction, encoding, transport, and response handling.
Matched warmups and RPC counters were checked for all Rust pairs, and source,
wheel, native-extension, dependency, and harness identities were audited.

Results measure plan and expression sharing together, not either Arrow-handoff
experiment or isolated protobuf encoding. They do not establish real-server
performance, statistical significance for small changes, API parity, or universal
gains. Breadth and confirmation samples are reported separately. Raw reports have
not been published in the harness repository.

Does this PR introduce any user-facing change?

Yes, for direct Rust users of the public representation:

• Expression becomes a shared handle; its enum variants move to
  ExpressionNode. Construct nodes using the common constructors or .into(),
  and match through expression.node().
• LogicalPlan child fields change from Box<LogicalPlan> to
  Arc<LogicalPlan>. Plan helpers accept owned/shared nodes through
  Into<Arc<LogicalPlan>>.
• DirectShufflePartitionId takes Expression and CaseWhen::else_expr takes
  Option<Expression>.

The 4.3 migration note in docs/compatibility.md includes a construction and
matching example. Ordinary Column, DataFrame, and Python APIs are unchanged.
The wire protocol is unchanged: serialization still expands shared occurrences
into independent protobuf trees. Recursive serialization/destruction limits
remain, and flat CASE/window containers can still copy collections of handles.

How was this patch tested?

Benchmark rerun at ff6fddd3:

• 1,134 fresh measured attempts: 1,122 passed and 12 failed.
• All 756 Rust samples passed. The 12 failures were exclusively the official
  PySpark 4.2.0 reference's known Exceeded upb_DecodeOptions_MaxDepth failures
  in plan-expression-32-schema-first and plan-expression-32-collect-empty.
  The benchmark launcher retained exit status 1; the full run is not reported
  as green.
• All 11 harness unit tests and all nine pilot attempts passed. Pilot samples
  are excluded from the measured results.
• Warmup/RPC matching, execution-order balance, and provenance audits passed.
  No additional live-server or official parity suite was run for this benchmark.

After the review cleanup:

• 424 Rust tests passed: 368 core unit tests and 56 selected integration/golden
  tests, including seven expression-sharing tests. The selected golden target
  also reports 12 explicitly ignored tests; these are not counted as passes.
• All core integration-test targets compiled, with and without wasm-udf.
• The native Python extension built in release mode. A before/after snapshot of
  its 108 exported classes showed no changed members or text signatures.
• All 52 Rust documentation snippets type-checked, including the migration
  example; Rust formatting and git diff --check passed.

Full official parity and live-server suites were not rerun locally for this
cleanup. Local workspace Clippy is blocked by the missing
wasm32-unknown-unknown target; focused core/binding Clippy also reports existing
lint failures.

CI follow-up at 092c199:

• The strict MkDocs build, all 52 Rust documentation snippets, five tooling tests,
  Ruff checks, Rust formatting, and git diff --check passed.
• All 368 Rust core unit tests passed. Every integration-test target compiled in
  release mode, both with and without wasm-udf. The 108-class Python API snapshot
  is unchanged. Workspace Clippy still reports existing advisory lint errors.
• The official Python transport with pandas 3.1.0rc0 reproduced CI's exact utility
  test profile: 61 failed and 20 passed. Spark 4.2 imports DateOffset from
  pandas.tseries.frequencies, which pandas 3.1 removed. With pandas 2.3.3, all
  81 utility tests passed through both the official and Rust transports.
• A local official-suite attempt completed 15 files with 696 passed tests and no
  failures before the next stateful-streaming file exceeded its 360-second cap.
  The attempt was stopped during its retry; this is not a full green parity run.
  In the completed Arrow Python UDF file, all 319 tests passed, versus 205 failures
  in the previous CI environment.
• Fresh GitHub checks are running on the CI follow-up commit; their results are
  pending. Local checks used macOS arm64, Python 3.12, and Java 21, rather than
  CI's Ubuntu/Python 3.11/Java 17 environment.

Was this patch authored or co-authored using generative AI tooling?

Yes.

This pull request and its description were written by Isaac.

[dbtsai]

cc @HyukjinKwon @viirya @dongjoon-hyun to take a look. Thanks.

[viirya]

Thanks for this. Sharing immutable plan/expression nodes is the right direction, and the change looks safe: immutability is enforced by the type system (there's no get_mut/make_mut path, and Expression's Arc is private), plan IDs are still assigned on the freshly built proto so shared self-joins keep distinct IDs, and none of the golden expectations change. The O(n²) → O(n) gain for chained transformations matches the benchmarks.

Most comments are inline: scope, constructor boilerplate, the two remaining boxed handles, and some redundant tests. One question that doesn't fit on a line: this changes the public Rust API (Expression becomes a handle around ExpressionNode, and LogicalPlan children become Arc) within the 4.3.0 line. Could we add a short migration note for the release notes? Keeping the plan::* helpers source-compatible through Into<Arc<LogicalPlan>> is a nice touch.

Optional: the benchmark section of the description is quite long. The main table plus a line on the unexplained core-outbound-rpc-deep-10 result would be enough for reviewers.

[viirya]

This impl is only used in tests; production code spells out Expression::new(ExpressionNode::X(...)) about 200 times (plus about 80 in tests). Could we either use .into() at call sites or add a few constructors for the common cases (e.g. Expression::literal, Expression::column, Expression::function(name, args))? That would shrink this diff a lot and keep future call sites readable.

[dbtsai]

Added Expression::literal, Expression::column, and Expression::function(name, args) and updated the common production and test call sites. Less common variants use the existing From<ExpressionNode> conversion through .into() or Expression::from. Distinct-function flags, metadata, and bound-column plan IDs are preserved.

[viirya]

Expression is now pointer-sized, so the two remaining Box<Expression>s, DirectShufflePartitionId(Box<Expression>) (L68) and CaseWhen::else_expr: Option<Box<Expression>> (L749), are just an extra allocation and indirection. Since this PR already changes the public representation, could we make these DirectShufflePartitionId(Expression) and Option<Expression> now rather than breaking the API a second time later? The boxed struct variants (Alias(Box<Alias>) etc.) still make sense for keeping the node size down; these two don't.

[dbtsai]

Removed both extra boxes: DirectShufflePartitionId now takes an Expression, and CaseWhen::else_expr is Option<Expression>. Updated the CASE and shuffle construction/serialization paths; the selected golden tests pass. The migration note also calls out these representation changes.

[viirya]

The change itself looks fine, since metadata emission still runs type and borrow checking. But this, the new lint.yml step, and scripts/tests/test_check_doc_snippets.py aren't related to plan/expression sharing. Could they go in a separate PR/JIRA so this one stays focused and the CI fix gets its own history entry?

[dbtsai]

Removed the docs-checker change, lint workflow change, and tooling tests from this PR's aggregate diff. The extracted CI fix is prepared as a separate follow-up and is being kept locally.

[viirya]

This pins the exact rustc argv, so it fails on any harmless flag change but can't catch a snippet that should fail to compile. test_checks_all_rust_snippets_even_after_a_failure covers behavior that predates this PR. If the tooling tests stay, I'd keep only what guards the new behavior. Either way, this belongs with the separate CI change.

[dbtsai]

The tooling tests are now outside this PR. In the local CI follow-up, replaced the argv assertion with compiler-backed checks: a valid snippet with an unresolved foreign symbol type-checks without native linking, while type-error and moved-value snippets fail. Removed the test for the pre-existing continue-after-failure behavior.

[viirya]

This and udf_clones_share_command_bytes_and_argument_lifetimes (L168) follow directly from node sharing: if the clone shares the node (already asserted by clones_and_arithmetic_chains_share_existing_nodes), whatever the node owns is shared too. Could we drop them, or fold the UDF proto encoding check into an existing test?

[dbtsai]

Dropped both redundant ownership/lifetime tests. The existing UDF serialization test already checks command bytes and eval type, so that coverage remains without another node-sharing test.

[viirya]

to_proto() returns an owned proto tree, so mutating one argument can't affect the other or the source expression; this part can't fail. The equality checks above it already show that the shared occurrence expands into two full subtrees.

[dbtsai]

Removed the owned-proto mutation assertions. Kept the equality checks that verify a shared expression occurrence expands into two complete protobuf subtrees.

[viirya]

This repeats the single-step Arc::ptr_eq assertion from transformations_and_clones_share_ancestors 256 times, and local_relation_payload_is_not_cloned_by_transformations (L2702) follows from that same pointer equality. Similarly, the per-builder sharing tests in readwriter.rs, merge.rs, streaming.rs and table_arg.rs could be one test. The plan-ID uniqueness test for shared self-joins is the important one here; I'd keep that and trim the rest.

[dbtsai]

Removed the repeated deep-plan and local-payload tests. Consolidated writer, V2 writer, streaming, merge, and table-argument sharing into related_builders_share_the_source_plan in dataframe.rs. Kept the single-step ancestor-sharing checks and the shared self-join plan-ID uniqueness test.

[dbtsai]

Thanks for the review. Addressed the feedback in ff6fddd: added common expression constructors, removed the two redundant expression boxes, trimmed/consolidated sharing tests, and added the Rust 4.3 migration note in docs/compatibility.md. I also shortened the description while retaining the benchmark caveats. The unrelated docs-checker CI changes have been extracted and are kept in a local follow-up.

Validation: 424 Rust tests passed; all core integration targets compiled with and without wasm-udf; the native Python extension built in release mode; all 52 Rust documentation snippets type-checked. A before/after snapshot of the 108 exported native classes found no changed members or text signatures. Formatting and diff checks passed.

Live-server/official parity suites and benchmarks were not rerun for this cleanup. Workspace Clippy is blocked locally by the missing WASM target, and focused core/binding Clippy still reports existing lint failures, as recorded in the updated description.

[viirya]

Thanks for the updates. All the earlier points are addressed: the constructors removed the boilerplate cleanly (distinct/plan-id/metadata flags are preserved, golden expectations unchanged), the two extra boxes are gone, the migration note reads well, and the trimmed test set keeps the important self-join plan-ID check. The compiler-backed snippet tests are a nice improvement over the argv assertion.

One thing: 092c199 brings the docs-snippet change, the lint.yml step and the tooling tests back into this PR, along with the pandas <3.0 pin and the parity-log retention in run_official_tests.py. They look right to me, and I assume the pandas pin is needed to get parity CI green at all, but none of it is about plan/expression sharing. Could they land as their own JIRA/PR first, with this one rebased on top? Otherwise they get squashed under this title. Not blocking if that's impractical.

Minor, optional: related_builders_share_the_source_plan asserts exact Arc::strong_count values, so an extra internal clone in any builder would fail it without being a bug; checking Arc::ptr_eq against each builder's plan would be sturdier.