Repository navigation
CaseWhen does not work with custom implemented column expression #21231
Description
Activity
Just for my understanding, could you explain the purpose of the custom column reference implementation?
In my case, lambda variable
My first thought was to add something like this
pub trait PhysicalExpr { fn column_index(&self) -> Option<usize> { None } }to avoid the downcast to
Column.For something like this (not real syntax, just making some educated guesses based on what you said)
case x@0 when y@1 then array_transform(z@2, l@3 -> l@3 + 1) endwhere the lambda variable is an extra index compared to the parent scope, or
case x@0 when y@1 then array_transform(z@2, l@0' -> l@0' + 1) endwhere
0'is a redefined index that's not going to be sufficient though. Without knowledge of the existence of lambda functions you can't really do the right thing. In the first alternative you would end up trying to project a column index that doesn't exist from the perspective of the record batch coming into the case expression. In the second alternative you would end up retaining an existing column index, but that's actually a misinterpretation of the index.One way around that might be to push the recursion into
PhysicalExprwith something likefn required_column_indices(&self) -> Vec<usize> { self.children() .iter() .map(|child| child.required_column_indices()) .flatten() .sorted() .unique() .collect() }but I don't think that's going to get us to a solution for lambda functions either. I only just started looking at the lambda PR. I'll continue with that tomorrow; maybe I can find some inspiration there.
Short term, we might just have to disable the 'project' code path. It could be turned into an optimisation instead with a dedicated 'project' wrapper expression. That way you could at least opt-out if it's getting in the way (or alternatively make it opt-in only).
case x@0 when y@1 then array_transform(z@2, l@3 -> l@3 + 1) endThat's how it will look when column capture support get's reintroduced.
Before projecting, can't we just remove the out of bounds indexes relative to the batch being projected? Since lambda variables are always pushed to the end of the batch, all of the out of bounds ones should refer lambda variables not yet available.fn case_when_with_expr( &self, batch: &RecordBatch, projected: &ProjectedCaseBody, ) -> Result<ColumnarValue> { let return_type = self.data_type(&batch.schema())?; // projected.projection may include indexes of lambda variables not available on this batch let projection = projected .projection .iter() .copied() .filter(|index| *index < batch.num_columns()) .collect::<Vec<_>>(); if projection.len() < batch.num_columns() { let projected_batch = batch.project(&projection)?; projected .body .case_when_with_expr(&projected_batch, &return_type) } else { self.body.case_when_with_expr(batch, &return_type) } }
There are cases where variables should be projected:
array_transform( x@0, a@2 -> case when y@1 > 2 then array_transform(a@2, b@3 -> b@3 +1) else then array_transform(a@2, c@3 -> c@3 * 2) end )Here, CaseWhen should project
a@2but notb@3norc@3
To be stricter, we can also modify the column indices collecting tree traversal to keep track of new lambdas variables introduced by lambda nodes, and skip projecting them. On the example above, CaseWhen can see the lambdasb -> b + 1andc -> c*2, store the name of it's parametersbandc, and skip collecting the indices of any variable with those names. Since it wouldn't have seen the lambdaa -> case .., it should project the index of variable namedanormallycase x@0 when y@1 then array_transform(z@2, l@0' -> l@0' + 1) endThat's how it looks like today at 1b7f4bf. Since there's no column capture for now and no regular column can exist within a lambda body, we just jump all lambda nodes, and none of the variables of it's body will be visited.
@gstvg yes, when the column extraction logic is “lambda aware” we can certainly implement this correctly. What Inwas trying to think about is how we can write it so that third-party code that introduces new constructs does not get broken. In other words, what extension points do we need to allow @rluvaton’s “custom column” implementation (which I’m using as a proxy for whatever kind of extension people might have built) to work correctly without the DataFusion crate being aware of its existence and semantics.
I think something similar to
childrenandwith_new_childrenthe case when that I have a problem with is inside the lambda function, like
transform(list, x => when x > col('other_col') then col('other_col') when x < col('another_col') then x else -2)
Before projecting, can't we just remove the out of bounds indexes relative to the batch being projected? Since lambda variables are always pushed to the end of the batch, all of the out of bounds ones should refer lambda variables not yet available.
No, because of the index is inbound? you still don't want to capture that and modify that subtree as this is also very prune to bugs
the problem beside the API for detecting and creating with different index custom column implementation is that you should traverse the expression through
childrenfunction, as it can go too deepLets say we have this expression:
array_transform( very_nested_list, nested_list => case when something then transform( concat(nested_list, external_array), list => len(list) + len(nested_list) ) )if you get all expression that implement
column_indexso in the subtree under the CASE WHEN you will get the following expressions:nested_list, external_array, list, nested_listthe
listhere is not bounded to the same schema as theCASE WHENcontext, it is in a different schema, so the referenced index might not exist in the schema that theCASE WHENis in.so you should not traverses/modify the tree through
childrenas it will lead you to that. instead you should traverse/modify through children that have the same input schema as you, and then the innermosttransformlambda physical expression will not return expression under the lambda function itself, and only underconcat(nested_list, external_array)What you’re describing is exactly the bit I’m wondering how to solve. A lambda function defines a new naming scope. Which is quite different from any other PhysicalExpr. If we need to assume that any PhysicalExpr can introduce a new concept like lambdas, then what kind of api do we need to be able to deal with that? I’m also thinking something along the lines of a tree traversal where the schema can change in the downward phase, and indices can be remapped in the upward phase. You would need to have the schema available for this though, which is not the case when the projection logic is currently triggered.
you only need to pass throught the same naming scope
you only need to pass throught the same naming scope
That's a bit too cryptic for me.
A lambda function defines a new naming scope. [..] I’m also thinking something along the lines of a tree traversal where the schema can change in the downward phase, and indices can be remapped in the upward phase. You would need to have the schema available for this though, which is not the case when the projection logic is currently triggered.
what I'm saying is that for you should stay in the same scope as the your expression (here the case expression), if some child expression create a new scope, you should not get into that scope. and therefor you don't need the schema available
all you need is for expression that create a new scope to not return their children
so for
array_transform(<expr>, value -> <expr-2>), array_transform expression will only return and not the lambda, since the lambda is a new scopeWouldn’t that make it impossible to implement column references from within lambdas?
no, and I already have it in my repo.
because lets say I have this:
select external_column, array_transform(list, item -> item + external_column) from tblso array_transform knows that it has 1 column that it should get from the parent scope
then this is what the implementation will save:
ArrayTransform: child: column list external references mapping: // (name used in the lambda expression, reference expression) "external_column", column(external column) lambda: Add: bound lambda variable: position: 0 bound lambda variable: position 1and when you create the schema in the lambda expression you map from external used references to the internal column position
and when case expression rewrite the columns, it rewrite the external reference mapping
Look at this example implementation of array transform with support for external reference capture that will work with the mapping
struct ArrayTransformExpression { child: Arc<dyn PhysicalExpr>, lambda_function: Arc<dyn PhysicalExpr>, external_references: Vec<( // Renamed schema field name in lambda schema String, // the expression to bound Arc<dyn PhysicalExpr> )> } impl ArrayTransformExpression { fn create_lambda_function_schema(&self, input_schema: &Schema, lambda_args: Vec<FieldRef>) -> SchemaRef { let mut lambda_schema_fields = lambda_args.clone(); let start_index = lambda_schema_fields.len() for (field_name, external_ref) in self.external_references { let field = external_ref.return_field(input_schema).unwrap(); let renamed_field = field.with_name(field_name); lambda_schema_fields.push(renamed_field); } Arc::new(Schema::new(lambda_schema_fields)) } fn create_lambda_function_input_batch(&self, input_batch: &RecordBatch, lambda_args: Vec<(FieldRef, ArrayRef)>) -> SchemaRef { let (mut lambda_schema_fields, mut lambda_arrays) = lambda_args.into_iter().unzip(); let start_index = lambda_schema_fields.len() for (field_name, external_ref) in self.external_references { let field = external_ref.return_field(input_schema).unwrap(); let renamed_field = field.with_name(field_name); lambda_schema_fields.push(renamed_field); let external_reference_value = external_ref.evaluate(input_batch); lambda_arrays.push(external_reference_value); } let lambda_function_schema = Arc::new(Schema::new(lambda_schema_fields)); RecordBatch::try_new(lambda_function_schema, lambda_arrays) } fn children_in_scope(&self) -> Vec<&Arc<dyn PhysicalExpr>> { std::iter::once(&self.child).chain(self.external_references.iter().map(|(_, external_ref)| external_ref)).collect() } fn with_new_children_in_scope( self: Arc<Self>, children_in_scope: Vec<Arc<dyn PhysicalExpr>>, ) -> Result<Arc<dyn PhysicalExpr>> { // child and external references assert_eq!(children_in_scope.len(), 1 + self.external_references.len()) Ok(Arc::new(ArrayTransformExpression { child: children_in_scope[0], lambda_function: self.lambda_function.clone(), // this update the mapping external_references: self.external_references .iter() .map(|(name, _old_ref)| name) .zip(children_in_scope.iter().skip(1)) .collect() })) } }
so in your case, if you work through children_in_scope and keep the same code you already do in CASE WHEN it will work
8 remaining items
the underlying problem is still the same, custom column are not supported while before they were
I know, I was just wondering what examples of custom columns are. The example you gave was lambda variables, but those now exist in the core library. Are there any other examples of custom columns that would make this useful?
Thus I think the "right" way to fix this bug (and avoid potential future bugs) is:
- Add an API to PhysicalExpr to report schema column used
@alamb I agree with you that we should try to make the aspect of 'which columns do you need' and 'project yourself with this column mapping' part of the
PhysicalExprtrait so that each implementation can handle this however it sees fit. That allows us to remove the downcasts fromCaseand wherever else similar logic exists.The part that I'm still not confident about is how we would control recursion. Would it be preferable for the 'collect columns' trait method to return the used columns for the physical expression itself and all its children (in other words, the method implementation itself handles the recursion) or should the recursion be done externally using a visitor?
The tradeoff I see is that internal recursion potentially requires a lot of data copying. If the method returns a
Vecfor instance you have to merge Vecs at each recursion level, filter out duplicates, etc.With external recursion though we can avoid that and collect into a set, but I think we may need to use something like
Result<(Option<usize>, TreeNodeRecursion)>as return type so that implementations can return an optional column index and also have scope control.A possible hybrid is internal recursion, but rather than returning values via the return value, we do so by passing in a consumer callback function. That avoids having to have
TreeNodeRecursionin the method signature, and also avoids the need for recursive result merging, but then we probably have to have a dyn function parameter.Is there any prior art in the library for something like this?
I don't quite follow the options you are proposing @pepijnve
You could implement internal recursion by passing in the result Vec (e.g.
output: &mut Vec<...>) but I am not sure what is what you meanI don't think an API that reports "what columns do you need" would handle "scoped" access as was originally described above -- as every time the recursion entered a new scope you would have to rebind the column definitions somehow
However I am not sure I am answering the questions you had in mind
Sorry, I didn't provide enough context. I scribbled up my current understanding so I have something to point to.
I'm using high-order function calls here as an example, but I think this can be generalised to other things as well. If I understood it correctly, the problem we want to handle is the fact that the lambda function in this diagram will use a local schema that is different from the outer schema. The outer one is
a@0, b@1, c@2while the lambda one isel@0, c@1.My assumption is that somewhere at some point in the code there has to be logic that implements the mapping of the outer
c@2to the innerc@1. In my sketch, my assumption is that that's the 'apply higher order function' node and that is has a mapping table with an entry2 -> 1.A recursive 'collect column indices' function should be able to handle this provided that each node can process the answer from its children. The high-order function application expression node needs to know which of its children are plain expressions, which are functions, what the indices for the lambda variables are, etc. so I would expect it to be able to handle the mapping of the
[0, 1]answer from the function expression to[2]for the outer scope.As I wrote this, I think I've answered my own question. This is probably easiest to implement with internal recursion (i.e. the trait method provides the answer for the target expression itself and all its children). The alternative would require an extra method to handle the remapping. The other questions I had were more coding style questions regarding how you express output parameters in Rust: literally a
mutoutput parameter, a callback, something else?The more important question though is if the model sketched above is a sufficient solution? Am I overlooking something?
I think a classic way to implement this type of resolution is using a stack of named scope / bindings.
So that would be something like
(Query Bindings (column references) -> Local Scope Bindings (any newly defined references) ... -> nested Scope Bindings for a nested lambda, for exampe (any newly defined references)So then when we went to resolve a mapping
"a"to a column index it would look back in the stack of scope bindings.When the recursion passes a lambda function, it would push a new scope on ot the stack and pop off the stack when done
Column index binding happens much earlier than the case projection though, doesn't it? I'll check the outer column reference PR to see what's being proposed there.
Column index binding happens much earlier than the case projection though, doesn't it? I'll check the outer column reference PR to see what's being proposed there.
I guess I was thinking that physical planning would do whatever column binding was needed 🤔
then any reference to a non base column would be via some other type of physical exper (not Column)
Then the implementation of lambda would pass in the appropriate input somehow
Given @gstvg merged something for lambdas recently maybe they already sorted this out. I am not sure
At #21323, during physical planning, lambda scope schemas are created by appending the lambda parameters to the outer schema, so that the outer schema fields still exists in the same position, regardless of being referenced or not
(how it handle name shadowing I believe it's not relevant for this discussion). Given this query:create table t as select [[1, 2]] as b, 2 as b; select array_transform(a, arr -> array_transform(arr, v -> v + b)) from t;
The schemas would be:
[0 => a, 1 => b] [0 => a, 1 => b, 2 => arr] [0 => a, 1 => b, 2 => arr, 3 => v]The planned query with indices, which would be what is visible via tree node traversals:
select array_transform(a@0, arr@2 -> array_transform(arr@2, v@3 -> v@3 + b@1)) from t;
So every scoped schema is a superset of the root schema, and every column reference, regardless of being
within a lambda scope, has an index that it's valid relative to the root schema (lambda variable references uses a different expression type)Then, as an internal optimization, not visible to the external world, similar to the case optimization,
lambda derives a projected body to work with projected batches that doesn't include unrefereced columns or lambda variables:[0 => a, 1 => b] [0 => b, 1 => arr] [0 => b, 1 => v]select array_transform(a@0, arr@1 -> array_transform(arr@1, v@1 -> v@1 + b@0)) from t;
Reacted by Andrew Lamb@gstvg I think the approach where the indices are kept stable in the public physical plan and the rewriting is an internal optimisation is very sensible. Very easy to reason about.
This approach also aligns perfectly with the collect/project idea I sketched above. Would we still want to introduce methods on
PhsyicalExprto support this? We would essentially be replacing the current downcasting code with polymorphic method calls which could support custom column expressions provide we accept the 'extend/append only' model.The change you made in the lambda column capture branch where
IndexSetwas replaced withBTreeSetis an interesting one from an API design pov and is why I was asking questions earlier about how to collect the column indices in a recursive method call. I'm a bit reluctant to make the output collection type a part of the signature since that's up to the caller to decide. One caller may want to count occurrences of each index, one may not care about order while another does, etc. That kind of nudges in the direction of a dyn function callback rather than a mutable reference to some collection. Is there an idiomatic pattern in Rust for something this?So every scoped schema is a superset of the root schema, and every column reference, regardless of being
within a lambda scope, has an index that it's valid relative to the root schema (lambda variable references uses a different expression type)This is very clever. Nice.
This approach also aligns perfectly with the collect/project idea I sketched above. Would we still want to introduce methods on PhsyicalExpr to support this? We would essentially be replacing the current downcasting code with polymorphic method calls which could support custom column expressions provide we accept the 'extend/append only' model.
In my opinion, the general idea of switching to methods on PhysicalExpr rather than downcasting is a good idea in general as it makes the code easier to reason about and the contracts of what PhysicalExprs do / are used for more explicit in the trait (rather than having to look through the codebase at the uses)
hat kind of nudges in the direction of a dyn function callback rather than a mutable reference to some collection. Is there an idiomatic pattern in Rust for something this?
What about a visitor pattern (instead of hard coded APIs)? Maybe the existing treenode api can help
What about a visitor pattern (instead of hard coded APIs)? Maybe the existing treenode api can help
That's an option, but that would then remain an external traversal that requires knowledge of all possible PhysicalExpr implementations it needs to handle. If we want to this to be independent of specific implementations we need extra methods on the trait. I'm still a bit on the fence myself if this is a sensible thing to do at this point. Today there are only
ColumnandLambdaVariablethat we would need to support. With just two implementations that are both standard types providing a new extension point may be a bit premature. I'm having a hard time coming with other examples where this might be useful.- added a commit that references this issue
on Sep 3, 2026
Describe the bug
While upgrading to DataFusion 52 (I know, I'm getting to 53 slowly) tests started failing for CASE WHEN for us because we have cases where we use custom
ColumnPhysicalExprrather than DataFusion one.this should not affect correctness.
This was done as an optimization in this PR by @pepijnve
CASEevaluation #18329because it find the columns referenced by traversing the expression tree and finding the
ColumnPhysical Expression. however, when having customColumnimpl it will not add that used column to the projectionTo Reproduce
This is the same test as:
datafusion/datafusion/physical-expr/src/expressions/case.rs
Lines 1694 to 1732 in 1a0af76
but using Custom impl of
ColumnExpected behavior
to work with any custom column impl, and hopefully still having that optimization
Additional context
also might affect lambda support: