jayzhan211 commented on code in PR #26589:
URL: https://github.com/apache/datafusion/pull/26589#discussion_r4498937541
##########
datafusion/physical-expr/src/expressions/case.rs:
##########
@@ -175,10 +215,25 @@ impl CaseBody {
.collect::<IndexMap<usize, usize>>();
// Construct the projected body by rewriting each expression from the
original body
- // using the column index mapping.
+ // using the column index mapping. Precomputed expressions are
replaced with
+ // references to the columns appended after the projected ones.
let project = |expr: &Arc<dyn PhysicalExpr>| -> Result<Arc<dyn
PhysicalExpr>> {
Arc::clone(expr)
.transform_down(|e| {
+ if !precomputed.is_empty()
+ && is_precomputable(&e)
Review Comment:
This makes planning exponential in nested CASE depth when levels read
different fields. Replacing a field access rebuilds each enclosing `CaseExpr`
through `try_new`, which runs `project()` again on the rewritten subtree. With
`CASE WHEN s['f1'] = 1 THEN 0 ELSE CASE WHEN s['f2'] = 2 ...` at depth 20 in a
debug build, planning goes from 10 ms to 47 s under the default config, and
about 4x every two levels in release.
Could you derive the projection on first evaluation instead? This fixed it
when I tried it, and also fixes an existing blow-up on main for direct columns
nested in reverse order. Please add a test too; the one below passes with this
change.
```rs
#[derive(Debug, Default)]
struct LazyProjectedBody(
OnceLock<std::result::Result<ProjectedCaseBody, Arc<DataFusionError>>>,
);
impl LazyProjectedBody {
fn get(&self, body: &CaseBody) -> Result<&ProjectedCaseBody> {
self.0
.get_or_init(|| body.project().map_err(Arc::new))
.as_ref()
.map_err(|e| DataFusionError::Shared(Arc::clone(e)))
}
}
```
`EvalMethod` then holds `LazyProjectedBody::default()`, and `evaluate` calls
`p.get(&self.body)?`. It needs no-op `Hash` and `Eq` impls, like
`LiteralLookupTable`.
```rs
#[test]
fn nested_case_with_distinct_struct_fields() -> Result<()> {
let fields = (0..=40)
.map(|i| Field::new(format!("f{i}"), DataType::Int32, true))
.collect::<Vec<_>>();
let columns = (0..=40)
.map(|i| Arc::new(Int32Array::from(vec![i, i + 1])) as ArrayRef)
.collect::<Vec<_>>();
let s = StructArray::new(fields.into(), columns, None);
let schema = Arc::new(Schema::new(vec![Field::new(
"s",
s.data_type().clone(),
true,
)]));
let batch = RecordBatch::try_new(Arc::clone(&schema),
vec![Arc::new(s)])?;
let s = col("s", &schema)?;
let mut expr = get_field_expr(Arc::clone(&s), "f0", &schema)?;
for i in 1..=40 {
let when = binary(
get_field_expr(Arc::clone(&s), &format!("f{i}"), &schema)?,
Operator::Eq,
lit(-1i32),
&schema,
)?;
expr = Arc::new(CaseExpr::try_new(None, vec![(when, lit(1i32))],
Some(expr))?);
}
let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?;
assert_eq!(as_int32_array(&result)?, &Int32Array::from(vec![0, 1]));
Ok(())
}
```
##########
datafusion/physical-expr/src/expressions/case.rs:
##########
@@ -224,10 +279,46 @@ impl CaseBody {
Ok(ProjectedCaseBody {
projection,
+ precomputed: precomputed.into_iter().collect(),
body: projected_body,
supports_projection,
})
}
+
+ /// Iterates over all the top-level expressions of this case body.
+ fn exprs(&self) -> impl Iterator<Item = &Arc<dyn PhysicalExpr>> {
+ self.expr
+ .iter()
+ .chain(self.when_then_expr.iter().flat_map(|(w, t)| [w, t]))
+ .chain(self.else_expr.iter())
+ }
+}
+
+/// Returns true if `expr` is a cheap, non-volatile expression over columns
and literals,
+/// such as a struct field access, that can be evaluated for the whole input
batch ahead of
+/// the `CASE` branches.
+///
+/// Expressions with [`ExpressionPlacement::MoveTowardsLeafNodes`] placement
may be pushed
+/// below filters by the optimizer, so evaluating them for rows a branch does
not select is
+/// expected to be safe.
+fn is_precomputable(expr: &Arc<dyn PhysicalExpr>) -> bool {
Review Comment:
`get_field` on a map isn't the cheap case this assumes: it scans entries for
every row and can fail. `SELECT CASE WHEN id < 0 THEN m[7] WHEN id > 100 THEN
m[8] ELSE 0 END FROM mt` returns 0 on main and errors here. With 10 entries and
8 branches, a batch where no branch matches goes from 11 µs to 518 µs, though
uniform branches do get faster (1005 to 738 µs).
Could you leave map lookups out? With the lazy projection, the batch schema
is available to `project()`, so `is_precomputable` can also return false when
`crosses_map(expr, schema)` is true. Struct, `Dictionary<Struct>` and `Null`
bases stay precomputed, and only a path that reaches a `Map` is excluded. I
tried this: it fixes the error and the slowdowns, at the cost of the
uniform-branch win.
```rs
fn crosses_map(expr: &Arc<dyn PhysicalExpr>, schema: &Schema) -> bool {
expr.exists(|e| {
let Some(func) = e.downcast_ref::<ScalarFunctionExpr>() else {
return Ok(false);
};
let [base, path @ ..] = func.args() else {
return Ok(false);
};
let mut data_type = base.data_type(schema)?;
for key in path {
let name = key
.downcast_ref::<Literal>()
.and_then(|l|
l.value().try_as_str().flatten().map(str::to_string));
data_type = match (&data_type, name) {
(DataType::Map(..), _) => return Ok(true),
(DataType::Struct(fields), Some(name)) => {
match fields.iter().find(|f| *f.name() == name) {
Some(f) => f.data_type().clone(),
None => return Ok(false),
}
}
_ => return Ok(false),
};
}
Ok(false)
})
.unwrap_or(false)
}
```
Please add tests for the query above, plus `CASE WHEN s['a'] = 1 THEN
s['m']['k'] WHEN s['a'] = 2 THEN s['m']['k'] END` with `m` a map field. The
flat `get_field(s, 'm', 'k')` should give an empty `precomputed` and keep `s`
in `projection`, while the nested `get_field(get_field(s, 'm'), 'k')` should
precompute `s['a']` and `s['m']`.
##########
datafusion/physical-expr/src/expressions/case.rs:
##########
@@ -272,16 +372,43 @@ struct ProjectedCaseBody {
}
impl ProjectedCaseBody {
- /// Returns a projection when evaluating the derived body can avoid unused
input columns.
- fn projection_for(&self, batch: &RecordBatch) -> Option<Vec<usize>> {
+ /// Returns the batch to evaluate the derived body against, when doing so
can avoid
+ /// filtering unused input columns.
+ fn project_batch(&self, batch: &RecordBatch) ->
Result<Option<RecordBatch>> {
+ if !self.supports_projection {
+ return Ok(None);
+ }
let projection = self
.projection
.iter()
.copied()
.filter(|index| *index < batch.num_columns())
.collect::<Vec<_>>();
- (self.supports_projection && projection.len() < batch.num_columns())
- .then_some(projection)
+ if self.precomputed.is_empty() {
+ return if projection.len() < batch.num_columns() {
+ Ok(Some(batch.project(&projection)?))
+ } else {
+ Ok(None)
+ };
+ }
+
+ let schema = batch.schema();
+ let mut fields = Vec::with_capacity(projection.len() +
self.precomputed.len());
+ let mut columns = Vec::with_capacity(fields.capacity());
+ for index in projection {
+ fields.push(Arc::clone(&schema.fields()[index]));
+ columns.push(Arc::clone(batch.column(index)));
+ }
+ for expr in &self.precomputed {
+ fields.push(expr.return_field(&schema)?);
+ columns.push(expr.evaluate(batch)?.into_array(batch.num_rows())?);
Review Comment:
This adds a failure main doesn't have. The field comes from the plan while
the array comes from the batch, and `try_new_with_options` rejects any
difference in nested types or nullability. A plan that declares `s` non-null
over a batch with null structs fails with `Column 'a' is declared as
non-nullable`, and on main it works. I could only reach this with directly
built physical plans, which is the engine case from the issue. This change
fixed both probes:
```suggestion
let column = expr.evaluate(batch)?.into_array(batch.num_rows())?;
// The array may not match the plan-time field exactly (nested
nullability or
// metadata); only the expressions in the body read this schema.
let field = expr.return_field(&schema)?;
let nullable = field.is_nullable() || column.null_count() > 0;
fields.push(Arc::new(
field
.as_ref()
.clone()
.with_data_type(column.data_type().clone())
.with_nullable(nullable),
));
columns.push(column);
```
##########
datafusion/physical-expr/src/expressions/case.rs:
##########
@@ -132,40 +136,76 @@ struct CaseBody {
impl CaseBody {
/// Derives a [ProjectedCaseBody] from this [CaseBody].
fn project(&self) -> Result<ProjectedCaseBody> {
- // Determine the set of columns that are used in all the expressions
of the case body.
+ // Lambda bodies are evaluated against batches of a different shape, so
+ // expressions are only precomputed when the body contains no lambdas.
+ let can_precompute = !self.exprs().any(|e| {
+ e.exists(|e| Ok(e.is::<LambdaExpr>() || e.is::<LambdaVariable>()))
+ .expect("Closure cannot fail")
+ });
+
+ // Unknown leaves may read input columns without exposing a Column
child.
+ // This is checked over all the nodes, independently of precomputation.
+ let supports_projection = !self.exprs().any(|e| {
+ e.exists(|e| {
+ Ok(e.children().is_empty()
+ && !e.is::<Column>()
+ && !e.is::<LambdaVariable>()
+ && !e.is::<Literal>()
+ && !e.is::<ScalarFunctionExpr>())
+ })
+ .expect("Closure cannot fail")
+ });
+
+ // Find the candidates for precomputation and the columns used outside
of them.
+ let mut direct_column_indices = HashSet::<usize>::new();
+ let mut precompute_candidates = IndexSet::<Arc<dyn
PhysicalExpr>>::new();
+ if can_precompute {
+ self.exprs().for_each(|expr| {
+ expr.apply(|expr| {
+ if is_precomputable(expr) {
+ precompute_candidates.insert(Arc::clone(expr));
+ return Ok(TreeNodeRecursion::Jump);
+ }
+ if let Some(column) = expr.downcast_ref::<Column>() {
+ direct_column_indices.insert(column.index());
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })
+ .expect("Closure cannot fail");
+ });
+ }
+
+ // Precomputing only pays off when it spares filtering a column that
is not
+ // otherwise used. Candidates referring only to used columns are
evaluated in place.
+ let precomputed = precompute_candidates
+ .into_iter()
+ .filter(|expr| {
+ collect_columns(expr)
+ .iter()
+ .any(|column|
!direct_column_indices.contains(&column.index()))
+ })
+ .collect::<IndexSet<_>>();
+
+ // Determine the set of columns that are used by the expressions of
the case body
+ // outside of the precomputed expressions.
// Use an ordered set so lambda variables continue to be positioned
after columns
let mut used_column_indices = BTreeSet::<usize>::new();
- let mut supports_projection = true;
- let mut collect_column_indices = |expr: &Arc<dyn PhysicalExpr>| {
+ self.exprs().for_each(|expr| {
expr.apply(|expr| {
+ if !precomputed.is_empty() && precomputed.contains(expr) {
Review Comment:
Optional cleanup: this hashes every visited node, and a nested `CaseExpr`
hashes its whole subtree. Checking `is_precomputable` first, as the rewrite
below does, limits hashing to real candidates.
```suggestion
if !precomputed.is_empty()
&& is_precomputable(expr)
&& precomputed.contains(expr)
{
```
--
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]
