Weijun-H commented on code in PR #26589:
URL: https://github.com/apache/datafusion/pull/26589#discussion_r4523104285
##########
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| {
Review Comment:
Echoing the existing comment on planning cost: replacing a child here
rebuilds each enclosing nested `CaseExpr` through `with_new_children` ->
`try_new` -> `project()`, which rewrites the already-rewritten subtree again. A
reachable shape is `C0 = s['f0']`, `Ci = CASE WHEN s['fi'] = -1 THEN 1 ELSE
C(i-1) END` with distinct fields: the outer level renumbers the fields to
columns 0..d, the inner suffix then has shifted indices and is renumbered and
rebuilt again at every level, so the work grows faster than linearly with
depth. (Derived from the source; I did not time it.)
##########
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())?);
+ }
+ let options =
RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+ Ok(Some(RecordBatch::try_new_with_options(
+ Arc::new(Schema::new(fields)),
Review Comment:
`Schema::new(fields)` drops the input schema metadata, while the
non-precompute path (`batch.project`) keeps it. Could this preserve it?
```suggestion
Arc::new(Schema::new_with_metadata(fields,
schema.metadata().clone())),
```
##########
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 {
+ if expr.placement() != ExpressionPlacement::MoveTowardsLeafNodes ||
is_volatile(expr)
Review Comment:
Echoing the existing comment on maps: `get_field` reports
`MoveTowardsLeafNodes` from argument placement alone, so a map lookup with a
literal key is also precomputed. For `CASE WHEN c1 THEN m['a'] WHEN c2 THEN
m['b'] ELSE 0 END` where no row matches either condition, main evaluates
neither THEN, while this now scans the map twice for the whole batch before the
first branch. The result is the same; the scans are new.
--
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]
