test: cover positional struct reconciliation in native CASE WHEN and COALESCE

native/spark-expr/src/conditional_funcs/case_when.rs

@@ -1659,6 +1659,180 @@ mod tests {
         }
     }
 
+    /// The same reconciliation as above through the arrays, maps and structs that hold the structs,
+    /// with either branch first, and with the `CASE WHEN x IS NOT NULL THEN x ELSE y END` that
+    /// COALESCE is planned as. The result takes the first branch's names and widens neither INT
+    /// nor DOUBLE, as it would if fields were paired by name.
+    #[test]
+    fn conditional_branches_reconcile_nested_case_variant_fields_positionally() {
+        use arrow::array::{Float64Array, ListArray, MapArray};
+        use arrow::datatypes::Fields;
+
+        /// The structure a struct is nested in.
+        #[derive(Clone, Copy, Debug)]
+        enum Layer {
+            List,
+            MapValue,
+            StructField,
+        }
+
+        /// Nests `inner` in `layers`, innermost first, with one element per row at every level.
+        /// Only the outermost level has `nulls`.
+        fn nest(inner: ArrayRef, layers: &[Layer], nulls: NullBuffer) -> ArrayRef {
+            let len = inner.len();
+            let mut array = inner;
+            for (i, layer) in layers.iter().enumerate() {
+                let nulls = (i + 1 == layers.len()).then(|| nulls.clone());
+                let offsets = OffsetBuffer::from_lengths(vec![1; len]);
+                array = match layer {
+                    Layer::List => Arc::new(ListArray::new(
+                        Arc::new(Field::new("item", array.data_type().clone(), true)),
+                        offsets,
+                        array,
+                        nulls,
+                    )),
+                    Layer::MapValue => {
+                        let fields: Fields = vec![
+                            Field::new("key", DataType::Utf8, false),
+                            Field::new("value", array.data_type().clone(), true),
+                        ]
+                        .into();
+                        let keys: StringArray = (0..len).map(|_| Some("k")).collect();
+                        let entries = StructArray::new(fields, vec![Arc::new(keys), array], None);
+                        Arc::new(MapArray::new(
+                            Arc::new(Field::new("entries", entries.data_type().clone(), false)),
+                            offsets,
+                            entries,
+                            nulls,
+                            false,
+                        ))
+                    }
+                    Layer::StructField => Arc::new(StructArray::new(
+                        vec![Field::new("s", array.data_type().clone(), true)].into(),
+                        vec![array],
+                        nulls,
+                    )),
+                };
+            }
+            array
+        }
+
+        // The first field of each struct is an INT and the second a DOUBLE, whatever their names
+        let then_fields: Fields = vec![
+            Field::new("x", DataType::Int32, true),
+            Field::new("X", DataType::Float64, false),
+        ]
+        .into();
+        let else_fields: Fields = vec![
+            Field::new("X", DataType::Int32, false),
+            Field::new("x", DataType::Float64, true),
+        ]
+        .into();
+        let result_fields = |names: [&str; 2]| -> Fields {
+            vec![
+                Field::new(names[0], DataType::Int32, true),
+                Field::new(names[1], DataType::Float64, true),
+            ]
+            .into()
+        };
+        let structs = |fields: &Fields, ints: Vec<Option<i32>>, doubles: Vec<Option<f64>>| {
+            Arc::new(StructArray::new(
+                fields.clone(),
+                vec![
+                    Arc::new(Int32Array::from(ints)) as ArrayRef,
+                    Arc::new(Float64Array::from(doubles)),
+                ],
+                None,
+            )) as ArrayRef
+        };
+        // Each branch is NULL in the row the other chooses, so whichever is first, the first row
+        // takes the THEN value and the second row the ELSE value.
+        let then_nulls = NullBuffer::from(vec![true, false]);
+        let else_nulls = NullBuffer::from(vec![false, true]);
+
+        for layers in [
+            vec![],
+            vec![Layer::List],
+            vec![Layer::MapValue],
+            vec![Layer::StructField, Layer::List],
+            vec![Layer::List, Layer::MapValue],
+            vec![Layer::MapValue, Layer::StructField],
+        ] {
+            let then_inner = structs(
+                &then_fields,
+                vec![Some(7), None],
+                vec![Some(5.5), Some(5.5)],
+            );
+            let else_inner = structs(&else_fields, vec![Some(0), Some(0)], vec![Some(9.5), None]);
+            let then_array = if layers.is_empty() {
+                // There is no outer array to carry the NULL, so the struct carries it itself
+                let columns = then_inner.as_struct().columns().to_vec();
+                Arc::new(StructArray::new(
+                    then_fields.clone(),
+                    columns,
+                    Some(then_nulls.clone()),
+                )) as ArrayRef
+            } else {
+                nest(then_inner, &layers, then_nulls.clone())
+            };
+            let else_array = if layers.is_empty() {
+                let columns = else_inner.as_struct().columns().to_vec();
+                Arc::new(StructArray::new(
+                    else_fields.clone(),
+                    columns,
+                    Some(else_nulls.clone()),
+                )) as ArrayRef
+            } else {
+                nest(else_inner, &layers, else_nulls.clone())
+            };
+            let schema = Schema::new(vec![
+                Field::new("t", then_array.data_type().clone(), true),
+                Field::new("e", else_array.data_type().clone(), true),
+            ]);
+            let batch = RecordBatch::try_new(
+                Arc::new(schema.clone()),
+                vec![Arc::clone(&then_array), Arc::clone(&else_array)],
+            )
+            .unwrap();
+            let c = |name: &str| col(name, &schema).unwrap();
+
+            // Which branch is first, and so which struct field names the result has
+            for (first, second, names) in [("t", "e", ["x", "X"]), ("e", "t", ["X", "x"])] {
+                let expected_inner = structs(
+                    &result_fields(names),
+                    vec![Some(7), Some(0)],
+                    vec![Some(5.5), None],
+                );
+                let expected = nest(expected_inner, &layers, NullBuffer::new_valid(2));
+                let exprs = [
+                    create_if_expr(
+                        Arc::new(IsNotNullExpr::new(c(first))),
+                        c(first),
+                        c(second),
+                        &schema,
+                    )
+                    .unwrap(),
+                    create_case_when(
+                        vec![(Arc::new(IsNotNullExpr::new(c(first))), c(first))],
+                        Some(c(second)),
+                        &schema,
+                    )
+                    .unwrap(),
+                ];
+                for expr in exprs {
+                    let context = format!("{layers:?} {expr}");
+                    assert_eq!(
+                        expr.data_type(&schema).unwrap(),
+                        *expected.data_type(),
+                        "{context}"
+                    );
+                    let result = evaluate_conditional(&expr, &batch);
+                    assert_eq!(result.as_ref(), expected.as_ref(), "{context}");
+                }
+            }
+        }
+    }
+
     /// IF branches that share a Spark timestamp type but carry different Arrow timezone labels are
     /// reconciled by relabelling one of them, as CASE branches are. The result has the THEN
     /// branch's label, which `IfExpr` reported before the ELSE branch was cast, unless it has none.

spark/src/test/resources/sql-tests/expressions/conditional/case_when_field_names.sql

@@ -65,3 +65,69 @@ query expect_native(casewhen)
 SELECT to_json(CASE WHEN p THEN named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))
   ELSE named_struct('X', 0, 'x', d) END)
 FROM test_case_field_names
+
+-- Case-distinct names repeat across positions in the opposite order. Spark aligns fields by
+-- position, so x stays an INT. Pairing fields by name would widen it to DOUBLE, and to_json would
+-- emit 7.0. The rows of an INSERT ... VALUES are written as separate one-row files, and a batch
+-- never spans files, so the first two INSERTs have every batch take THEN or ELSE alone. COALESCE(1)
+-- writes the third as a single file, so its rows reach CASE WHEN as one batch that mixes the two.
+-- https://github.com/apache/datafusion-comet/issues/6532
+statement
+CREATE TABLE test_case_positional(q boolean, i int, d double) USING parquet
+
+statement
+INSERT INTO test_case_positional VALUES (true, 7, 9.5), (true, NULL, NULL)
+
+statement
+INSERT INTO test_case_positional VALUES (false, 7, 9.5), (false, NULL, NULL)
+
+statement
+INSERT INTO test_case_positional SELECT /*+ COALESCE(1) */ * FROM VALUES (true, 7, 9.5), (false, 8, 10.5), (NULL, NULL, NULL), (false, 3, NULL) AS v(q, i, d)
+
+-- Both branch orders
+query expect_native(casewhen)
+SELECT
+  to_json(CASE WHEN q THEN named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))
+    ELSE named_struct('X', 0, 'x', d) END),
+  to_json(CASE WHEN q THEN named_struct('X', 0, 'x', d)
+    ELSE named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)) END)
+FROM test_case_positional
+
+-- Several branches, none of which has the others' names
+query expect_native(casewhen)
+SELECT
+  to_json(CASE WHEN q THEN named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))
+    WHEN i > 7 THEN named_struct('X', 1, 'x', d)
+    ELSE named_struct('x', 2, 'X', d) END),
+  to_json(CASE WHEN q THEN named_struct('X', 1, 'x', d)
+    WHEN i > 7 THEN named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))
+    ELSE named_struct('X', 2, 'x', d) END)
+FROM test_case_positional
+
+-- Native to_json supports structs, so extract the struct after reconciling the array or map.
+query expect_native(casewhen)
+SELECT
+  to_json(CASE WHEN q THEN array_repeat(named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)), 1)
+    ELSE array_repeat(named_struct('X', 0, 'x', d), 1) END[0]),
+  to_json(CASE WHEN q THEN array_repeat(named_struct('X', 0, 'x', d), 1)
+    ELSE array_repeat(named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)), 1) END[0])
+FROM test_case_positional
+
+query expect_native(casewhen)
+SELECT
+  to_json(CASE WHEN q THEN map('k', named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)))
+    ELSE map('k', named_struct('X', 0, 'x', d)) END['k']),
+  to_json(CASE WHEN q THEN map('k', named_struct('X', 0, 'x', d))
+    ELSE map('k', named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))) END['k'])
+FROM test_case_positional
+
+-- Nested structs, and an array inside a map value
+query expect_native(casewhen)
+SELECT
+  to_json(CASE WHEN q THEN named_struct('s', named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)))
+    ELSE named_struct('s', named_struct('X', 0, 'x', d)) END),
+  to_json(CASE WHEN q THEN map('k', array(named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))))
+    ELSE map('k', array(named_struct('X', 0, 'x', d))) END['k'][0]),
+  to_json(CASE WHEN q THEN map('k', array(named_struct('X', 0, 'x', d)))
+    ELSE map('k', array(named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)))) END['k'][0])
+FROM test_case_positional

spark/src/test/resources/sql-tests/expressions/conditional/coalesce_field_names.sql

@@ -34,3 +34,81 @@ SELECT to_json(coalesce(
   CASE WHEN q THEN named_struct('X', i + 100) END,
   named_struct('X', 7)))
 FROM test_coalesce_field_names
+
+-- Case-distinct names repeat across positions in the opposite order. Spark aligns fields by
+-- position, so x stays an INT. Pairing fields by name would widen it to DOUBLE, and to_json would
+-- emit 7.0. The rows of an INSERT ... VALUES are written as separate one-row files, and a batch
+-- never spans files, so the first two INSERTs have every batch take the first or the last argument
+-- alone. COALESCE(1) writes the third as a single file, so its rows reach COALESCE as one batch
+-- that mixes the two.
+-- https://github.com/apache/datafusion-comet/issues/6532
+statement
+CREATE TABLE test_coalesce_positional(p boolean, i int, d double) USING parquet
+
+statement
+INSERT INTO test_coalesce_positional VALUES (true, 7, 9.5), (true, NULL, NULL)
+
+statement
+INSERT INTO test_coalesce_positional VALUES (false, 7, 9.5), (false, NULL, NULL)
+
+statement
+INSERT INTO test_coalesce_positional SELECT /*+ COALESCE(1) */ * FROM VALUES (true, 7, 9.5), (false, 8, 10.5), (NULL, NULL, NULL), (false, 3, NULL) AS v(p, i, d)
+
+-- Both argument orders
+query expect_native(coalesce)
+SELECT
+  to_json(coalesce(
+    CASE WHEN p THEN named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)) END,
+    named_struct('X', 0, 'x', d))),
+  to_json(coalesce(
+    CASE WHEN p THEN named_struct('X', 0, 'x', d) END,
+    named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))))
+FROM test_coalesce_positional
+
+-- Several arguments, none of which has the others' names
+query expect_native(coalesce)
+SELECT
+  to_json(coalesce(
+    CASE WHEN p THEN named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)) END,
+    CASE WHEN i > 7 THEN named_struct('X', 1, 'x', d) END,
+    named_struct('x', 2, 'X', d))),
+  to_json(coalesce(
+    CASE WHEN p THEN named_struct('X', 1, 'x', d) END,
+    CASE WHEN i > 7 THEN named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)) END,
+    named_struct('X', 2, 'x', d)))
+FROM test_coalesce_positional
+
+-- Native to_json supports structs, so extract the struct after reconciling the array or map.
+query expect_native(coalesce)
+SELECT
+  to_json(coalesce(
+    CASE WHEN p THEN array_repeat(named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)), 1) END,
+    array_repeat(named_struct('X', 0, 'x', d), 1))[0]),
+  to_json(coalesce(
+    CASE WHEN p THEN array_repeat(named_struct('X', 0, 'x', d), 1) END,
+    array_repeat(named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)), 1))[0])
+FROM test_coalesce_positional
+
+query expect_native(coalesce)
+SELECT
+  to_json(coalesce(
+    CASE WHEN p THEN map('k', named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))) END,
+    map('k', named_struct('X', 0, 'x', d)))['k']),
+  to_json(coalesce(
+    CASE WHEN p THEN map('k', named_struct('X', 0, 'x', d)) END,
+    map('k', named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))))['k'])
+FROM test_coalesce_positional
+
+-- Nested structs, and an array inside a map value
+query expect_native(coalesce)
+SELECT
+  to_json(coalesce(
+    CASE WHEN p THEN named_struct('s', named_struct('x', i, 'X', CAST(5.5 AS DOUBLE))) END,
+    named_struct('s', named_struct('X', 0, 'x', d)))),
+  to_json(coalesce(
+    CASE WHEN p THEN map('k', array(named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)))) END,
+    map('k', array(named_struct('X', 0, 'x', d))))['k'][0]),
+  to_json(coalesce(
+    CASE WHEN p THEN map('k', array(named_struct('X', 0, 'x', d))) END,
+    map('k', array(named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)))))['k'][0])
+FROM test_coalesce_positional