LLDay commented on code in PR #20009:
URL: https://github.com/apache/datafusion/pull/20009#discussion_r2732251301
##########
datafusion/physical-plan/src/execution_plan.rs:
##########
@@ -1152,6 +1194,57 @@ pub fn check_default_invariants<P: ExecutionPlan +
?Sized>(
Ok(())
}
+/// Verifies that the [`ExecutionPlan::physical_expressions`] and
+/// [`ExecutionPlan::with_physical_expressions`] implementations are
consistent.
+///
+/// It verifies:
+/// 1. `with_physical_expressions` with the same expressions returns a plan
with the same expressions.
+/// 2. `with_physical_expressions` with a different number of expressions
returns an error.
+///
+/// Returns rewritten plan if it supported expression replacement.
+pub fn check_physical_expressions(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+ let Some(exprs) = plan.physical_expressions() else {
+ if let Some(new_plan) = plan.with_physical_expressions(vec![])? {
+ assert_or_internal_err!(plan.physical_expressions().is_none());
+
+ let new_expr = Arc::new(expressions::Column::new("", 0));
+ let result = plan.with_physical_expressions(vec![new_expr]);
+ assert_or_internal_err!(result.is_err());
+
+ return Ok(new_plan);
+ };
+
+ return Ok(plan);
+ };
+
+ let mut exprs = exprs.collect::<Vec<_>>();
+
+ // Should return plan with same expressions
+ let new_plan = plan.with_physical_expressions(exprs.clone())?.unwrap();
+ let new_plan_exprs = new_plan.physical_expressions().unwrap();
+
+ assert_or_internal_err!(exprs.iter().cloned().eq(new_plan_exprs));
+
+ // Should return an error if number of expressions differs from
physical_expressions()
+ let new_expr = Arc::new(expressions::Column::new("", 0));
+ exprs.push(new_expr);
+
+ let result = new_plan.with_physical_expressions(exprs.clone());
+ assert_or_internal_err!(result.is_err());
+
+ if exprs.len() > 2 {
+ exprs.pop().unwrap();
+ exprs.pop().unwrap();
+
+ let result = new_plan.with_physical_expressions(exprs);
+ assert_or_internal_err!(result.is_err());
Review Comment:
Added error description.
--
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]
