Copilot commented on code in PR #27281:
URL: https://github.com/apache/datafusion/pull/27281#discussion_r4521006909
##########
datafusion/physical-expr/src/expressions/case.rs:
##########
@@ -295,20 +295,187 @@ pub struct CaseExpr {
body: CaseBody,
/// Evaluation method to use
eval_method: EvalMethod,
+ required_operand: Option<RequiredOperandReuse>,
}
-// eval_method is functionally derived from body, so excluding it from
-// Hash/Eq avoids redundantly hashing the expression tree twice. For
+/// Caller-supplied permission for reusing an operand at its first required
evaluation.
+/// The authorized operand must evaluate to an array.
+/// This permission does not make the operand non-volatile or safe to hoist
elsewhere.
+/// Decisions must be immutable. The permission object's identity is part of
CASE
+/// equality so generic expression substitution cannot introduce a different
grant.
+pub trait RequiredOperandReusePolicy: Debug + Send + Sync {
+ /// The child evaluated unconditionally, exactly once, before this
expression can
+ /// fail or evaluate any other child. It must receive the original input
batch,
+ /// with the same rows and column positions, without prior slicing or
projection.
+ /// Replacing that child with a column must preserve the parent's
behavior, and
+ /// `with_new_children` must recompute any derived state that depends on
the child.
+ /// Return `None` when any of these conditions is unknown.
+ fn first_required_child(&self, expr: &Arc<dyn PhysicalExpr>) ->
Option<usize>;
+
+ /// Whether a rewrite retained the authorized call, excluding its
rewritten arguments.
+ /// Complete argument and type equality is checked separately between the
two uses.
+ fn same_call(
+ &self,
+ authorized: &Arc<dyn PhysicalExpr>,
+ rewritten: &Arc<dyn PhysicalExpr>,
+ ) -> bool;
+}
+
+#[derive(Debug)]
+struct RequiredOperandReuse {
+ policy: Arc<dyn RequiredOperandReusePolicy>,
+ first_path: Vec<usize>,
+ // Prepared expressions only: never cache values or share a batch between
evaluations.
+ prepared: Mutex<Option<(usize, Arc<CaseExpr>)>>,
+}
+
+impl RequiredOperandReuse {
+ fn try_new(
+ body: &CaseBody,
+ authorized: &Arc<dyn PhysicalExpr>,
+ policy: Arc<dyn RequiredOperandReusePolicy>,
+ ) -> Option<Self> {
+ if body.expr.is_some() {
+ return None;
+ }
+ let operand = body.else_expr.as_ref()?;
+ if !policy.same_call(authorized, operand) {
+ return None;
+ }
+ let mut current = Arc::clone(&body.when_then_expr.first()?.0);
+ let mut first_path = Vec::new();
+ while !current.eq(operand) {
+ let index = policy.first_required_child(¤t)?;
+ let next = Arc::clone(current.children().get(index)?);
+ first_path.push(index);
+ current = next;
+ }
+ let reuse = Self {
+ policy,
+ first_path,
+ prepared: Mutex::new(None),
+ };
+ // Exclude unknown leaves in the remaining body before evaluating the
operand.
+ // The operand itself is replaced, so it need not support column
projection.
+ let bound = reuse
+ .bind_case(body, Arc::new(Column::new("__case_required_operand",
0)))
+ .ok()?;
+ match &bound.eval_method {
+ EvalMethod::ExpressionOrExpression(projected)
+ | EvalMethod::NoExpression(projected)
+ if projected.supports_projection =>
+ {
+ Some(reuse)
+ }
+ _ => None,
+ }
+ }
+
+ fn replace_first(
+ expr: Arc<dyn PhysicalExpr>,
+ path: &[usize],
+ value: &Arc<dyn PhysicalExpr>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ let Some((&index, rest)) = path.split_first() else {
+ return Ok(Arc::clone(value));
+ };
+ let mut children: Vec<_> =
expr.children().into_iter().cloned().collect();
+ children[index] = Self::replace_first(Arc::clone(&children[index]),
rest, value)?;
+ expr.with_new_children(children)
+ }
+
+ fn prepared_case(
+ &self,
+ owner: &CaseExpr,
+ input_columns: usize,
+ ) -> Result<Arc<CaseExpr>> {
+ let mut cache = self
+ .prepared
+ .lock()
+ .map_err(|_| internal_datafusion_err!("CASE preparation lock
poisoned"))?;
+ if let Some((columns, prepared)) = cache.as_ref()
+ && *columns == input_columns
+ {
+ return Ok(Arc::clone(prepared));
+ }
+ let value: Arc<dyn PhysicalExpr> =
+ Arc::new(Column::new("__case_required_operand", input_columns));
+ let prepared = Arc::new(self.bind_case(&owner.body, value)?);
+ *cache = Some((input_columns, Arc::clone(&prepared)));
+ Ok(prepared)
+ }
+
+ fn bind_case(
+ &self,
+ body: &CaseBody,
+ value: Arc<dyn PhysicalExpr>,
+ ) -> Result<CaseExpr> {
+ let mut branches = body.when_then_expr.clone();
+ branches[0].0 =
+ Self::replace_first(Arc::clone(&branches[0].0), &self.first_path,
&value)?;
+ // Derive projection only after replacing the two authorized uses.
+ CaseExpr::try_new(None, branches, Some(value))
+ }
+
+ fn evaluate(&self, owner: &CaseExpr, batch: &RecordBatch) ->
Result<ColumnarValue> {
+ // Bind only the required operand. Let the prepared CASE retain its
existing
+ // uniform-predicate fast paths and avoid evaluating an unselected
branch.
+ let operand = owner
+ .body
+ .else_expr
+ .as_ref()
+ .expect("authorized CASE has ELSE");
Review Comment:
This evaluates the ELSE instance rather than the operand instance reached in
the first WHEN. Structural equality permits distinct volatile instances with
separate execution state, so this can change the predicate or result and
executes the ELSE instance even when every predicate is true. Traverse
`first_path` from the first WHEN instead, and add a regression with equal but
distinct instances.
##########
datafusion/physical-expr/src/expressions/case.rs:
##########
@@ -1372,11 +1568,24 @@ impl PhysicalExpr for CaseExpr {
),
(false, false) => (None, &children[0..children.len()],
None),
};
- Ok(Arc::new(CaseExpr::try_new(
+ let mut rebuilt = CaseExpr::try_new(
expr.cloned(),
when_then_expr.iter().cloned().tuples().collect(),
else_expr.cloned(),
- )?))
+ )?;
+ if let Some(required) = &self.required_operand {
+ // The old ELSE is the authorization baseline for the
rewritten uses.
+ let authorized = self
+ .body
+ .else_expr
+ .as_ref()
+ .expect("authorized CASE has ELSE");
+ rebuilt = rebuilt.with_required_operand_reuse(
+ authorized,
+ Arc::clone(&required.policy),
+ );
Review Comment:
A matching argument rewrite can retain permission while invalidating the
array-valued precondition: unprojecting a column to a literal makes the
existing `Counted` operand return a scalar, but both uses still match.
Evaluation then fails with `required CASE operand must evaluate to an array`
instead of returning a valid result. Support scalar results after rewriting or
revalidate the array-valued guarantee before retaining permission, and cover
literal unprojection in the rewrite tests.
--
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]
