fix: support struct-typed scalar subquery results

.github/workflows/pr_build_linux.yml

@@ -507,6 +507,7 @@ jobs:
               org.apache.comet.CometUuidExpressionSuite
               org.apache.comet.serde.CometScalarFunctionSuite
               org.apache.comet.serde.CometLiteralSuite
+              org.apache.comet.serde.CometScalarSubquerySuite
               org.apache.comet.CometFallbackInvarianceSuite
       fail-fast: false
     name: ${{ matrix.profile.name }} [${{ matrix.suite.name }}]

.github/workflows/pr_build_macos.yml

@@ -252,6 +252,7 @@ jobs:
               org.apache.comet.CometUuidExpressionSuite
               org.apache.comet.serde.CometScalarFunctionSuite
               org.apache.comet.serde.CometLiteralSuite
+              org.apache.comet.serde.CometScalarSubquerySuite
               org.apache.comet.CometFallbackInvarianceSuite
 
       fail-fast: false

docs/source/user-guide/latest/expressions.md

@@ -698,6 +698,8 @@ Comet also accelerates a number of Catalyst expressions that have no Spark SQL f
 
 This list is illustrative, not exhaustive: the per-function tables are not the complete set of expressions Comet can accelerate.
 
+Scalar subqueries can return structs, including those created when Spark merges multiple scalar subqueries. Struct results are transferred from Spark through Arrow IPC during native physical planning and retained as owned, immutable literals for that plan's execution. Supported fields include booleans, numeric types, default-collation strings, binary, dates, timestamps, nulls, and nested structs. Decimal fields require a non-negative scale no greater than their precision. Structs must be non-empty and have distinct field names at each level; arrays, maps, intervals, and other unsupported field types still cause fallback to Spark. Existing non-struct scalar-subquery paths are unchanged.
+
 ## See also
 
 - [Comet Compatibility Guide](compatibility/index.md) - known incompatibilities and edge cases for supported expressions.

native/core/src/execution/expressions/subquery.rs

@@ -16,11 +16,13 @@
 // under the License.
 
 use crate::{
+    errors::CometError,
     execution::utils::bytes_to_i128,
     jvm_bridge::{BinaryWrapper, JVMClasses, StringWrapper},
 };
-use arrow::array::RecordBatch;
+use arrow::array::{Array, ArrayRef, RecordBatch, StructArray};
 use arrow::datatypes::{DataType, Schema, TimeUnit};
+use arrow::ipc::reader::StreamReader;
 use datafusion::common::{internal_err, ScalarValue};
 use datafusion::logical_expr::ColumnarValue;
 use datafusion::physical_expr::PhysicalExpr;
@@ -30,11 +32,12 @@ use jni::{
 };
 use std::{
     fmt::{Display, Formatter},
-    hash::Hash,
+    io::Cursor,
     sync::Arc,
 };
 
-#[derive(Debug, Hash, PartialEq, Eq)]
+/// Runtime lookup for non-struct scalar results. The planner resolves structs to owned literals.
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub struct Subquery {
     /// The ID of the execution context that owns this subquery. We use this ID to retrieve the
     /// subquery result.
@@ -53,6 +56,88 @@ impl Subquery {
             data_type,
         }
     }
+
+    /// Resolve Spark's already materialized struct result during native physical planning.
+    /// Registration precedes the first executePlan call, which creates the physical plan.
+    /// The planner stores the owned result in a Literal, so evaluation needs no JVM lookup or
+    /// mutable initialization state. Separate physical expressions may resolve the same ID.
+    pub fn resolve_struct(
+        exec_context_id: i64,
+        id: i64,
+        data_type: &DataType,
+    ) -> datafusion::common::Result<ScalarValue> {
+        if !matches!(data_type, DataType::Struct(_)) {
+            return internal_err!("Expected a struct scalar subquery, got {data_type:?}");
+        }
+        JVMClasses::with_env(|env| unsafe {
+            let is_null = jni_static_call!(env,
+                comet_exec.is_null(exec_context_id, id) -> jboolean
+            )?;
+            if is_null {
+                return ScalarValue::try_from(data_type);
+            }
+            let bytes = jni_static_call!(env,
+                comet_exec.get_struct(exec_context_id, id) -> BinaryWrapper
+            )?;
+            let bytes = JByteArray::from_raw(env, bytes.get().as_raw());
+            let bytes = env.convert_byte_array(bytes).map_err(CometError::from)?;
+            decode_struct_result(&bytes, data_type)
+        })
+    }
+}
+
+/// serializeScalarSubquery emits one batch with one row and one struct column. Check the bounds
+/// needed for scalar extraction, but do not scan for additional batches from this internal
+/// producer. Arrow IPC and planned-type validation still apply to the returned value.
+fn decode_struct_result(
+    bytes: &[u8],
+    data_type: &DataType,
+) -> datafusion::common::Result<ScalarValue> {
+    let mut reader = StreamReader::try_new(Cursor::new(bytes), None)?;
+    let Some(batch) = reader.next().transpose()? else {
+        return internal_err!("Scalar subquery IPC result contains no batch");
+    };
+    if batch.num_rows() != 1 || batch.num_columns() != 1 {
+        return internal_err!("Scalar subquery IPC result must contain one row and one column");
+    }
+    let value = align_struct_metadata(batch.column(0), data_type)?;
+    ScalarValue::try_from_array(&value, 0)
+}
+
+// Utils.toArrowSchema preserves field order, names, types and nullability but not Parquet field
+// ID metadata. Restore only that metadata from the planned type, without permitting type casts.
+fn align_struct_metadata(
+    value: &ArrayRef,
+    expected: &DataType,
+) -> datafusion::common::Result<ArrayRef> {
+    match (value.data_type(), expected) {
+        (DataType::Struct(actual), DataType::Struct(fields))
+            if actual.len() == fields.len()
+                && actual
+                    .iter()
+                    .zip(fields.iter())
+                    .all(|(a, b)| a.name() == b.name() && a.is_nullable() == b.is_nullable()) =>
+        {
+            let Some(value) = value.as_any().downcast_ref::<StructArray>() else {
+                return internal_err!("Scalar subquery IPC result is not a struct array");
+            };
+            let children = value
+                .columns()
+                .iter()
+                .zip(fields.iter())
+                .map(|(child, field)| align_struct_metadata(child, field.data_type()))
+                .collect::<datafusion::common::Result<Vec<_>>>()?;
+            Ok(Arc::new(StructArray::try_new(
+                fields.clone(),
+                children,
+                value.nulls().cloned(),
+            )?))
+        }
+        (actual, expected) if actual == expected => Ok(Arc::clone(value)),
+        (actual, expected) => {
+            internal_err!("Scalar subquery IPC result has type {actual:?}, expected {expected:?}")
+        }
+    }
 }
 
 impl Display for Subquery {
@@ -195,3 +280,175 @@ impl PhysicalExpr for Subquery {
         Ok(self)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::{
+        array::{new_null_array, AsArray, Int32Array, StringArray},
+        datatypes::Field,
+        ipc::writer::StreamWriter,
+    };
+    use datafusion::physical_expr::expressions::Literal;
+
+    fn encode(schema: &Schema, batches: &[RecordBatch]) -> Vec<u8> {
+        let mut bytes = Vec::new();
+        {
+            let mut writer = StreamWriter::try_new(&mut bytes, schema).unwrap();
+            for batch in batches {
+                writer.write(batch).unwrap();
+            }
+            writer.finish().unwrap();
+        }
+        bytes
+    }
+
+    fn batch(value: ArrayRef) -> RecordBatch {
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "value",
+            value.data_type().clone(),
+            true,
+        )]));
+        RecordBatch::try_new(schema, vec![value]).unwrap()
+    }
+
+    fn struct_value() -> ArrayRef {
+        Arc::new(StructArray::new(
+            vec![
+                Field::new("number", DataType::Int32, false),
+                Field::new("text", DataType::Utf8, true),
+            ]
+            .into(),
+            vec![
+                Arc::new(Int32Array::from(vec![42])),
+                Arc::new(StringArray::from(vec!["Comet 彗星"])),
+            ],
+            None,
+        ))
+    }
+
+    #[test]
+    fn struct_ipc_round_trip() {
+        let value = struct_value();
+        let batch = batch(Arc::clone(&value));
+        let bytes = encode(batch.schema().as_ref(), std::slice::from_ref(&batch));
+        assert_eq!(
+            decode_struct_result(&bytes, value.data_type()).unwrap(),
+            ScalarValue::try_from_array(&value, 0).unwrap()
+        );
+    }
+
+    #[test]
+    fn struct_ipc_distinguishes_null_struct_from_null_fields() {
+        let fields = vec![Field::new("number", DataType::Int32, true)].into();
+        let all_null_fields: ArrayRef = Arc::new(StructArray::new(
+            fields,
+            vec![new_null_array(&DataType::Int32, 1)],
+            None,
+        ));
+        let null_struct = new_null_array(all_null_fields.data_type(), 1);
+        for (value, expected_null) in [(all_null_fields, false), (null_struct, true)] {
+            let batch = batch(Arc::clone(&value));
+            let bytes = encode(batch.schema().as_ref(), std::slice::from_ref(&batch));
+            let scalar = decode_struct_result(&bytes, value.data_type()).unwrap();
+            assert_eq!(scalar.is_null(), expected_null);
+            assert_eq!(scalar, ScalarValue::try_from_array(&value, 0).unwrap());
+        }
+    }
+
+    #[test]
+    fn struct_ipc_restores_nested_field_metadata() {
+        let inner = struct_value();
+        let outer: ArrayRef = Arc::new(StructArray::new(
+            vec![Field::new("nested", inner.data_type().clone(), true)].into(),
+            vec![inner],
+            None,
+        ));
+        let with_id = |field: Field, id: &str| {
+            field.with_metadata([("PARQUET:field_id".to_owned(), id.to_owned())].into())
+        };
+        let expected = DataType::Struct(
+            vec![with_id(
+                Field::new(
+                    "nested",
+                    DataType::Struct(
+                        vec![
+                            with_id(Field::new("number", DataType::Int32, false), "2"),
+                            Field::new("text", DataType::Utf8, true),
+                        ]
+                        .into(),
+                    ),
+                    true,
+                ),
+                "1",
+            )]
+            .into(),
+        );
+        let batch = batch(Arc::clone(&outer));
+        let bytes = encode(batch.schema().as_ref(), std::slice::from_ref(&batch));
+        let scalar = decode_struct_result(&bytes, &expected).unwrap();
+        assert_eq!(scalar.data_type(), expected);
+        let ScalarValue::Struct(result) = scalar else {
+            panic!("Expected struct scalar");
+        };
+        let nested = result
+            .column(0)
+            .as_any()
+            .downcast_ref::<StructArray>()
+            .unwrap();
+        assert_eq!(
+            nested.column(0),
+            outer.as_struct().column(0).as_struct().column(0)
+        );
+    }
+
+    #[test]
+    fn struct_ipc_rejects_invalid_shape_and_type() {
+        let batch = batch(struct_value());
+        let schema = batch.schema();
+        let data_type = batch.column(0).data_type();
+        assert!(decode_struct_result(b"invalid IPC", data_type).is_err());
+        assert!(decode_struct_result(&encode(&schema, &[]), data_type).is_err());
+        assert!(decode_struct_result(&encode(&schema, &[batch.slice(0, 0)]), data_type).is_err());
+        let bytes = encode(&schema, std::slice::from_ref(&batch));
+        let wrong_type = DataType::Struct(
+            vec![
+                Field::new("number", DataType::Int64, false),
+                Field::new("text", DataType::Utf8, true),
+            ]
+            .into(),
+        );
+        assert!(decode_struct_result(&bytes, &wrong_type).is_err());
+    }
+
+    #[test]
+    fn resolved_struct_literal_owns_its_value() {
+        let expected = ScalarValue::try_from_array(&struct_value(), 0).unwrap();
+        let literal = {
+            let batch = batch(struct_value());
+            let bytes = encode(batch.schema().as_ref(), std::slice::from_ref(&batch));
+            Literal::new(decode_struct_result(&bytes, &expected.data_type()).unwrap())
+        };
+        // The IPC bytes and input batch have been dropped, and no JVM registry is initialized.
+        let input = RecordBatch::new_empty(Arc::new(Schema::empty()));
+        for _ in 0..64 {
+            let ColumnarValue::Scalar(result) = literal.evaluate(&input).unwrap() else {
+                panic!("Expected scalar result");
+            };
+            assert_eq!(result, expected);
+        }
+    }
+
+    #[test]
+    fn resolved_null_struct_literal_preserves_its_type() {
+        let data_type = struct_value().data_type().clone();
+        // This is the typed-NULL branch used when Spark reports no subquery result, before IPC.
+        let literal = Literal::new(ScalarValue::try_from(&data_type).unwrap());
+        let input = RecordBatch::new_empty(Arc::new(Schema::empty()));
+        let ColumnarValue::Scalar(result) = literal.evaluate(&input).unwrap() else {
+            panic!("Expected scalar result");
+        };
+        assert!(result.is_null());
+        assert_eq!(result.data_type(), data_type);
+    }
+}

native/core/src/execution/planner.rs

@@ -786,7 +786,15 @@ impl PhysicalPlanner {
             ExprStruct::Subquery(expr) => {
                 let id = expr.id;
                 let data_type = to_arrow_datatype(expr.datatype.as_ref().unwrap());
-                Ok(Arc::new(Subquery::new(self.exec_context_id, id, data_type)))
+                if matches!(data_type, DataType::Struct(_)) {
+                    // Spark has materialized and registered the result before this physical
+                    // plan is created in executePlan. Keep an owned constant in the plan rather
+                    // than initializing mutable state while evaluating input batches.
+                    let value = Subquery::resolve_struct(self.exec_context_id, id, &data_type)?;
+                    Ok(Arc::new(DataFusionLiteral::new(value)))
+                } else {
+                    Ok(Arc::new(Subquery::new(self.exec_context_id, id, data_type)))
+                }
             }
             ExprStruct::BloomFilterMightContain(expr) => {
                 let bloom_filter_expr = self.create_expr(

native/jni-bridge/src/comet_exec.rs

@@ -46,6 +46,8 @@ pub struct CometExec<'a> {
     pub method_get_string_ret: ReturnType,
     pub method_get_binary: JStaticMethodID,
     pub method_get_binary_ret: ReturnType,
+    pub method_get_struct: JStaticMethodID,
+    pub method_get_struct_ret: ReturnType,
     pub method_is_null: JStaticMethodID,
     pub method_is_null_ret: ReturnType,
 }
@@ -117,6 +119,12 @@ impl<'a> CometExec<'a> {
                 jni::jni_sig!("(JJ)[B"),
             )?,
             method_get_binary_ret: ReturnType::Array,
+            method_get_struct: env.get_static_method_id(
+                JNIString::new(Self::JVM_CLASS),
+                jni::jni_str!("getStruct"),
+                jni::jni_sig!("(JJ)[B"),
+            )?,
+            method_get_struct_ret: ReturnType::Array,
             method_is_null: env.get_static_method_id(
                 JNIString::new(Self::JVM_CLASS),
                 jni::jni_str!("isNull"),