andygrove commented on code in PR #7322:
URL: https://github.com/apache/datafusion-comet/pull/7322#discussion_r4487695298
##########
native/spark-expr/src/conditional_funcs/null_short_circuit.rs:
##########
@@ -290,27 +289,46 @@ pub(super) fn evaluate_projected(
}
/// The input columns that `expr` reads, and `expr` reading them from a batch
of just those columns.
+///
+/// A [`LambdaVariable`] reads a column of the batch its lambda is evaluated
on exactly as a
+/// [`Column`] does, so it is projected and re-indexed the same way.
pub(super) fn project(expr: &Arc<dyn PhysicalExpr>) -> Result<(Vec<usize>,
Arc<dyn PhysicalExpr>)> {
- let mut projection = collect_columns(expr)
- .iter()
- .map(Column::index)
- .collect::<Vec<_>>();
+ let mut projection = vec![];
+ expr.apply(|e| {
+ if let Some(column) = e.downcast_ref::<Column>() {
+ projection.push(column.index());
+ } else if let Some(variable) = e.downcast_ref::<LambdaVariable>() {
+ projection.push(variable.index());
Review Comment:
Good catch. `project()` collected every `LambdaVariable` as a column of the
outer batch, but a lambda inside the projected expression binds its own
parameter, which is not a column of that batch, so `batch.project` was asked
for an index past the end. I reproduced it with a higher-order function as the
right operand of `AND` and `OR`, which goes through the same function, and the
same applied to `filter`.
`project()` now skips a variable whose lambda is inside the expression it
projects, and leaves it unchanged when it re-indexes. Lambda parameter names
are unique per lambda (`x#id`), so collecting the names a lambda in the
expression binds is enough to tell them apart from the free variables that the
earlier fix needs, where a short-circuited argument inside a lambda body reads
the lambda's own parameter.
Tests: your query, `array_contains(ints, size(filter(ints, x -> x > 1)))`,
and `AND`/`OR` with a lambda on the right are in
`array_filter_native_lambda.sql`, and a Rust planner test covers `AND` and
`OR`. I also compared these shapes against Spark with ANSI on and off. Fixed in
165695918.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
