Repository navigation
Fix CASE evaluation for custom column expressions - #24484
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24484 +/- ##
==========================================
- Coverage 81.59% 81.58% -0.01%
==========================================
Files 1123 1123
Lines 406811 406947 +136
Branches 406811 406947 +136
==========================================
+ Hits 331923 331999 +76
- Misses 55451 55489 +38
- Partials 19437 19459 +22 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
Thanks for working on this. I tested the CASE expression changes locally and the fix looks correct for the reported custom PhysicalExpr failure. The new regression test also fails as expected if the projection guards are removed, so it is exercising the actual bug.
I left a few non-blocking suggestions below. The main thing I would like us to be explicit about is that this is a conservative fix for the custom leaf case, rather than a complete solution to the broader scope-aware dependency problem discussed in #21231.
Overall, the fallback to evaluating the original CASE body against the full batch looks correct to me.
| expr.downcast_ref::<LambdaVariable>() | ||
| { | ||
| used_column_indices.insert(lambda_variable.index()); | ||
| } else if expr.downcast_ref::<Literal>().is_none() |
There was a problem hiding this comment.
I think this check is a little broader than the specific hazard we are trying to protect against. It will also disable CASE projection for some in-tree leaf expressions that cannot actually read the input batch positionally.
For example, a nullary ScalarFunctionExpr such as random() has no children, is not a Literal, and is not const-folded because it is volatile. That means an expression like CASE WHEN random() < 0.5 THEN a ELSE b END would now lose the projection fast path on a wide batch.
Would it make sense to keep known-safe DataFusion leaves projection-enabled, and reserve this fallback for expressions whose dependencies we genuinely cannot determine? If we prefer the conservative behavior for now, I think it would be worth calling out the performance tradeoff in the PR description.
There was a problem hiding this comment.
Addressed in 48b72e44c: known-safe ScalarFunctionExpr leaves remain projectable, while unknown leaves use the full-batch path. The regression suite also covers a nullary random() expression.
| } | ||
|
|
||
| #[test] | ||
| fn case_without_expr_with_custom_column() -> Result<()> { |
There was a problem hiding this comment.
Nice regression test. One thing I noticed is that this currently covers only the NoExpression evaluation path, while the PR adds the same projection guard to case_when_with_expr and expr_or_expr as well.
Could we add small tests for those two shapes too? A base-expression CASE using CustomColumn, plus a single-WHEN-with-ELSE case, would make sure all three guarded paths stay covered if this code is refactored later.
There was a problem hiding this comment.
Added coverage for the three guarded evaluation shapes: searched CASE, base-expression CASE, and the single-WHEN-with-ELSE path, each using a column-like custom expression.
| struct ProjectedCaseBody { | ||
| projection: Vec<usize>, | ||
| body: CaseBody, | ||
| supports_projection: bool, |
There was a problem hiding this comment.
Could we document the invariant on supports_projection here? When this is false, the derived projection and rewritten body must not be used, and evaluation needs to fall back to the original CASE body on the full batch.
That relationship is important enough that a short field comment, plus a sentence in the ProjectedCaseBody docs, would make the representation much easier to reason about.
There was a problem hiding this comment.
Documented the supports_projection invariant next to the representation and in the ProjectedCaseBody docs: when it is false, evaluation must use the original CASE body against the full batch.
| } else if expr.downcast_ref::<Literal>().is_none() | ||
| && expr.children().is_empty() | ||
| { | ||
| // Unknown leaves may read input columns without exposing a Column child. |
There was a problem hiding this comment.
This looks like a reasonable conservative fix for the reported crash. One process point though: the discussion in #21231 was also considering a PhysicalExpr dependency-reporting extension point, or a dedicated opt-in/opt-out mechanism for projection.
This PR takes a narrower approach by treating unknown leaves as unsafe. I think that is fine as a targeted fix, but it would be helpful to say so explicitly in the PR description and use wording like Part of #21231 rather than closing the issue entirely. The separate scope-aware traversal problem still remains.
There was a problem hiding this comment.
The PR description now calls this Part of #21231 and states the scope explicitly. The dependency-reporting extension discussed in the issue remains separate work.
There was a problem hiding this comment.
Small cleanup suggestion: the logic that filters the projection to valid batch indices, checks whether projection is supported, and decides whether projection is worthwhile is now repeated in three places.
A helper on ProjectedCaseBody, something like projection_for(&self, batch) -> Option<Vec<usize>>, could centralize that decision. That would also make it harder for a future CASE evaluation path to accidentally forget the supports_projection check.
d4a8b0f to
48b72e4
Compare
|
Addressed in 48b72e4 and rebased onto current
The branch is conflict-free after the rebase and the repository checks are running on the new head. |
|
The first post-rebase run caught two fixture typos and the clippy-required branch shape. I corrected the ELSE expectations to 40, reordered the imports, and collapsed the nested else-if in 6c1847f. Rust 1.97 validation in WSL passes rustfmt, strict physical-expr clippy, all three custom-column regressions, and the full physical-expr library suite (1,629 passed, 2 ignored). Fresh workflows are currently waiting for the repository approval gate. |
Signed-off-by: Hasnaat Hussain <[email protected]>
Signed-off-by: Hasnaat Hussain <[email protected]>
Signed-off-by: Hasnaat Hussain <[email protected]>
6c1847f to
f48c0e9
Compare
|
I refreshed the branch onto current The functional diff is still limited to |
There was a problem hiding this comment.
@Hasnaathussain, thanks for addressing the earlier review feedback. I took another pass through the follow-up changes and everything looks good to me.
The conservative CASE projection fallback is now limited appropriately, while known-safe Literal and ScalarFunctionExpr leaves keep the projection fast path. The nullary random() regression test confirms that behavior.
The added tests also cover all three evaluation shapes that use the projection decision: searched CASE, CASE with a base expression, and the single-WHEN-with-ELSE path. The supports_projection invariant is now documented, and ProjectedCaseBody::projection_for gives the evaluation paths a single place to make the projection decision.
I don't see any new correctness, schema, or regression issues in these changes. Looks good to me. Thanks!
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
|
🚀 |
Which issue does this PR address?
Rationale for this change
CaseBody::projectdiscovers input columns by downcasting expression-tree nodes to DataFusion's concreteColumnorLambdaVariable. A third-party column-likePhysicalExprcan evaluate an input column without exposing either node. The projection then omits that dependency and CASE evaluates the expression against a batch with the wrong schema.What changes are included in this PR?
CASE projection now falls back to the original input batch when it encounters an unknown leaf expression. Built-in
ScalarFunctionExprleaves remain eligible for projection because they receive the row count and do not read input columns through the batch. The projection decision is centralized inProjectedCaseBody::projection_for, with the invariant documented next to the representation.The regression coverage exercises searched CASE, base-expression CASE, and the single-WHEN expression path with a custom column-like expression. It also checks that a nullary
random()scalar function keeps the projection fast path.This is a conservative correctness fix for the custom-leaf case. The broader dependency-reporting design discussed in #21231 remains separate work.
Are these changes tested?
cargo fmt --all -- --checkcargo clippy -p datafusion-physical-expr --all-targets --all-features -- -D warningscargo test -p datafusion-physical-expr expressions::case::testscargo test -p datafusion-physical-exprAre there any user-facing changes?
Custom physical expressions can now be evaluated correctly inside searched and base-expression CASE expressions. There is no public API change.