Skip to content

perf: improve stateful LAG/LEAD evaluation for timezone timestamps - #27446

Open
rgbuilds wants to merge 1 commit into
apache:mainfrom
rgbuilds:perf/stateful-lead-lag-array
Open

rgbuilds wants to merge 1 commit into
apache:mainfrom
rgbuilds:perf/stateful-lead-lag-array

Conversation

@rgbuilds

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

Stateful LAG and LEAD evaluation scales poorly across partitions for timezone-aware timestamps.

Profiling the workload from #27101 showed that the existing path creates a ScalarValue for each output row and later converts those scalars back into an array. For timezone-aware timestamps, this repeatedly clones and drops shared timezone metadata, creating reference-count contention when partitions are evaluated concurrently.

This PR is a focused, partial improvement toward #27101. It does not optimize running aggregates, partition-key comparisons, or RANGE frame-bound searches.

What changes are included in this PR?

This PR adds an optional array-oriented stateful evaluation method, PartitionEvaluator::evaluate_stateful_array.

At a high level:

  • partition_evaluator.rs defines the array-evaluation contract and provides a default implementation that selects the existing scalar fallback.
  • standard.rs calls the optional array path, validates its result, and advances window state only after validation succeeds.
  • lead_lag.rs implements the optimized path for timezone-aware timestamps, constructing typed Arrow arrays without creating a ScalarValue for every output row.

The evaluator contract defines:

  • how buffered inputs and the first unfinished output row are provided;
  • how returned array elements map to a contiguous range of output rows;
  • the required output length and exact Arrow data type;
  • that partition completion must produce all remaining output rows;
  • how the consumed input range is reported;
  • that returning None selects scalar fallback and must have no observable side effects;
  • that evaluation errors are propagated rather than converted into fallback.

The method takes &self, which restricts ordinary mutation but does not guarantee a side-effect-free fallback because implementations may use interior mutability. Leaving state unchanged when returning None is therefore an explicit part of the contract.

The shared execution path validates the output length, output type, partition-completion requirements, and reported input range before updating its state. This ensures each output row is recorded exactly once, including LEAD results delayed until a later batch or partition completion.

The implementation preserves existing behavior for:

  • timestamp units and timezone metadata;
  • null and non-null default values;
  • positive, negative, zero, and extreme offsets;
  • IGNORE NULLS and RESPECT NULLS;
  • empty inputs;
  • batch and partition boundaries;
  • delayed LEAD output.

Existing evaluators retain their previous behavior because the new method has a default implementation that returns None.

The shared API is needed because the stateful execution loop is in datafusion-physical-expr, while the private LAG/LEAD evaluator is in datafusion-functions-window. Keeping the optimization entirely inside lead_lag.rs would not avoid scalar construction in the shared execution loop.

What is the testing strategy for this PR?

Focused tests compare stateful evaluation across multiple buffered batches with existing whole-partition evaluation. They cover:

  • timezone-aware timestamps;
  • null and non-null defaults;
  • IGNORE NULLS and RESPECT NULLS;
  • offsets including i64::MIN, negative, zero, positive, and i64::MAX;
  • empty completed partitions;
  • non-timezone scalar fallback;
  • array-result contract validation.

The existing bounded-window tests and relevant SQLLogicTests also pass, including timezone-preserving Parquet window-function coverage.

The following checks pass:

cargo fmt --all -- --check

cargo clippy --all-targets --all-features -- -D warnings

RUST_BACKTRACE=1 cargo test --profile ci \
    --exclude datafusion-examples \
    --exclude datafusion-benchmarks \
    --exclude datafusion-cli \
    --workspace --lib --tests --bins \
    --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption

Performance results

The benchmark used the 2,000,000-row dataset from #27101. Baseline and candidate queries used identical release profiles, data, and settings. Runs were interleaved, with the first run discarded as warm-up and five measured runs used for each result.

For timezone-aware LAG with 10 target partitions:

  • Baseline median: 0.689 seconds
  • Candidate median: 0.525 seconds
  • Improvement: 23.8%

For the corresponding timezone-free control with 10 target partitions:

  • Baseline median: 0.337 seconds
  • Candidate median: 0.338 seconds
  • Difference: approximately 0.3%, within observed noise

Single-partition results were approximately 2% slower for the candidate. A substantial parallel timezone/non-timezone gap also remains, so this is a partial improvement rather than a complete resolution of #27101.

Are there any user-facing changes?

SQL behavior and timestamp semantics are unchanged.

Users running stateful LAG or LEAD over timezone-aware timestamps may see improved parallel performance.

This adds the public, defaulted PartitionEvaluator::evaluate_stateful_array method. Existing implementations continue using the previous scalar behavior without modification, so this is not a breaking API change.

@github-actions github-actions Bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates functions Changes to functions implementation labels Nov 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

functions Changes to functions implementation logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant