sunchao commented on code in PR #6458:
URL: https://github.com/apache/datafusion-comet/pull/6458#discussion_r4154017963
##########
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:
[P2] Merge struct fields positionally before using DataFusion coercion. With
default case-insensitive resolution, nullable `i INT` and `d DOUBLE`, Spark
accepts `IF(b, named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)),
named_struct('X', 0, 'x', d))` and returns `STRUCT<x:INT,X:DOUBLE>`. Because
the branches differ in nullability, `type_union_coercion` reaches DataFusion's
name-based struct merge and pairs the first integer field with the second
double field. The new casts consequently change `x` to DOUBLE even on an
all-THEN batch that worked before. With native `to_json` enabled, the result
changes from `{"x":7,"X":5.5}` to `{"x":7.0,"X":5.5}`. Use positional recursive
struct reconciliation while retaining THEN names and combining nullability.
Evidence: An exact-head Rust reproduction constructed the branches with
`CreateNamedStruct` over b=[true,true], i=[7,NULL], d=[9.5,NULL]. The base
planner path, `IfExpr::new`, returned Int32 values [7,NULL] and JSON
{"x":7,"X":5.5}. `create_if_expr` returned Float64 values [7.0,NULL] and JSON
{"x":7.0,"X":5.5}. Spark 3.5.9 independently returned the integer schema and
original JSON. The case-distinct names pass Scala's duplicate-name check.
Native JSON observability requires
`spark.comet.expression.StructsToJson.allowIncompatible=true`. Reproduction
source is retained at `/tmp/review6458-repro.rs`, with output in
`/tmp/review6458-repro.log` and Spark results in
`/tmp/review6458-spark-reference.log`.
--
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]
