andygrove commented on code in PR #6458:
URL: https://github.com/apache/datafusion-comet/pull/6458#discussion_r4154800179
##########
native/spark-expr/src/conditional_funcs/case_when.rs:
##########
@@ -63,34 +64,66 @@ pub fn create_case_when(
else {
return Ok(Arc::new(CaseWhenExpr::try_new(when_then, else_expr)?));
};
- // The branches share a Spark type, so any difference is in the Arrow
representation. For a
- // timestamp that is the timezone label, and the cast only relabels it,
but Comet's cast still
- // needs a timezone. Every `TimestampType` value in a native plan is
labelled UTC.
- let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false);
- // A branch that already has the common type is not wrapped in a cast,
which would do nothing
- // but hide what the branch is from the evaluation.
- let coerce = |expr: Arc<dyn PhysicalExpr>, data_type: &DataType| ->
Arc<dyn PhysicalExpr> {
- if data_type == &coerce_type {
- expr
- } else {
- Arc::new(Cast::new(
- expr,
- coerce_type.clone(),
- cast_options.clone(),
- None,
- None,
- ))
- }
- };
let when_then = when_then
.into_iter()
.zip(&then_types)
- .map(|((when, then), then_type)| (when, coerce(then, then_type)))
+ .map(|((when, then), then_type)| (when, coerce_branch(then, then_type,
&coerce_type)))
.collect();
- let else_expr = else_expr.map(|e| coerce(e, &else_type));
+ let else_expr = else_expr.map(|e| coerce_branch(e, &else_type,
&coerce_type));
Ok(Arc::new(CaseWhenExpr::try_new(when_then, else_expr)?))
}
+/// Creates an `IF`, casting a branch whose type differs from the other's to
their common type.
+///
+/// Spark adds no cast when the branches differ only in whether a nested field
can be NULL, or in
+/// the case of a struct field name. But [`IfExpr`] reports the THEN branch's
type, and returns the
+/// ELSE branch's array unchanged when no row of a batch takes the THEN branch.
+pub fn create_if_expr(
+ if_expr: Arc<dyn PhysicalExpr>,
+ true_expr: Arc<dyn PhysicalExpr>,
+ false_expr: Arc<dyn PhysicalExpr>,
+ input_schema: &Schema,
+) -> Result<Arc<dyn PhysicalExpr>> {
+ let true_type = true_expr.data_type(input_schema)?;
+ let false_type = false_expr.data_type(input_schema)?;
+ // The coercion that `get_coerce_type_for_case_expression` folds over the
branches of a CASE
+ // WHEN, starting from the ELSE branch. Here the THEN branch goes first,
so the common type
+ // takes its struct field names, as Spark's `If.dataType` does.
+ let Some(common_type) = type_union_coercion(&true_type, &false_type) else {
Review Comment:
Fixed in ee1b213a2 and updated the branch-1.1 backport (#6491) in 6a83c1b94.
Struct fields now reconcile recursively by position, retaining THEN names and
combining nullability through structs, arrays and maps. Scalar coercion is
unchanged.
Confirmed the exact JSON difference and a nullable-ELSE panic on both
previous heads. Added native schema/value assertions and SQL coverage for the
final JSON result, both branch orders, all-THEN/all-ELSE/mixed batches, nulls
and nested containers. The expanded 21-query fixture passes on Spark 3.5 and
4.1 for both patches; the Rust conditional suites (19 source / 6 backport) and
workspace Clippy also pass. New-head CI is pending.
--
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]
