From 8b9f6831a536fd331e8490986dfe399b3e870ca4 Mon Sep 17 00:00:00 2001 From: LinSimon-901101 Date: Sun, 13 Sep 2026 03:23:02 +0800 Subject: [PATCH 1/2] fix: support struct-typed scalar subquery results --- docs/source/user-guide/latest/expressions.md | 2 + .../src/execution/expressions/subquery.rs | 285 +++++++++++++++++- native/jni-bridge/src/comet_exec.rs | 8 + .../spark/sql/comet/CometScalarSubquery.java | 10 + .../comet/serde/CometScalarSubquery.scala | 30 +- .../arrow/CometArrowConverters.scala | 57 +++- .../misc/scalar_subquery_struct.sql | 143 +++++++++ .../apache/comet/exec/CometExecSuite.scala | 91 +++++- 8 files changed, 613 insertions(+), 13 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/misc/scalar_subquery_struct.sql diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 208a5f3f124..712561d2dc3 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/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 to native execution through Arrow IPC and cached for the native expression's execution context. 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. diff --git a/native/core/src/execution/expressions/subquery.rs b/native/core/src/execution/expressions/subquery.rs index fc7b8104d2a..dba58e5b9e5 100644 --- a/native/core/src/execution/expressions/subquery.rs +++ b/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, - sync::Arc, + hash::{Hash, Hasher}, + io::Cursor, + sync::{Arc, OnceLock}, }; -#[derive(Debug, Hash, PartialEq, Eq)] +#[derive(Debug)] pub struct Subquery { /// The ID of the execution context that owns this subquery. We use this ID to retrieve the /// subquery result. @@ -43,6 +46,10 @@ pub struct Subquery { pub id: i64, /// The data type of the subquery result. pub data_type: DataType, + // Spark materializes a scalar subquery before native execution. Cache the owned struct + // result for this execution context so IPC serialization/decoding is not paid per batch. + // Do not include this execution state in expression equality or hashing. + struct_value: OnceLock, } impl Subquery { @@ -51,6 +58,81 @@ impl Subquery { exec_context_id, id, data_type, + struct_value: OnceLock::new(), + } + } +} + +impl PartialEq for Subquery { + fn eq(&self, other: &Self) -> bool { + self.exec_context_id == other.exec_context_id + && self.id == other.id + && self.data_type == other.data_type + } +} + +impl Eq for Subquery {} + +impl Hash for Subquery { + fn hash(&self, state: &mut H) { + self.exec_context_id.hash(state); + self.id.hash(state); + self.data_type.hash(state); + } +} + +/// The JVM bridge emits one row with one struct column. Validate the wire shape and type before +/// creating the scalar; Arrow IPC validation also keeps malformed strings out of native arrays. +fn decode_struct_result( + bytes: &[u8], + data_type: &DataType, +) -> datafusion::common::Result { + 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"); + } + if reader.next().transpose()?.is_some() { + return internal_err!("Scalar subquery IPC result contains more than one batch"); + } + 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 { + 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::() 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::>>()?; + 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:?}") } } } @@ -75,7 +157,10 @@ impl PhysicalExpr for Subquery { } fn evaluate(&self, _: &RecordBatch) -> datafusion::common::Result { - JVMClasses::with_env(|env| unsafe { + if let Some(value) = self.struct_value.get() { + return Ok(ColumnarValue::Scalar(value.clone())); + } + let result = JVMClasses::with_env(|env| unsafe { let is_null = jni_static_call!(env, comet_exec.is_null(self.exec_context_id, self.id) -> jboolean )?; @@ -87,6 +172,17 @@ impl PhysicalExpr for Subquery { } match &self.data_type { + DataType::Struct(_) => { + let bytes = jni_static_call!(env, + comet_exec.get_struct(self.exec_context_id, self.id) -> BinaryWrapper + )?; + let bytes = JByteArray::from_raw(env, bytes.get().as_raw()); + let bytes = env.convert_byte_array(bytes).map_err(CometError::from)?; + Ok(ColumnarValue::Scalar(decode_struct_result( + &bytes, + &self.data_type, + )?)) + } DataType::Boolean => { let r = jni_static_call!(env, comet_exec.get_bool(self.exec_context_id, self.id) -> jboolean @@ -181,7 +277,15 @@ impl PhysicalExpr for Subquery { } _ => internal_err!("Unsupported scalar subquery data type {:?}", self.data_type), } - }) + })?; + if matches!(self.data_type, DataType::Struct(_)) { + if let ColumnarValue::Scalar(value) = &result { + // Concurrent first evaluations may both initialize the same immutable result. + // Failed evaluations are never cached. + let _ = self.struct_value.set(value.clone()); + } + } + Ok(result) } fn children(&self) -> Vec<&Arc> { @@ -195,3 +299,172 @@ 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 std::collections::hash_map::DefaultHasher; + + fn encode(schema: &Schema, batches: &[RecordBatch]) -> Vec { + 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::() + .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.clone(), batch.clone()]), 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 struct_cache_does_not_change_expression_identity() { + let value = ScalarValue::try_from_array(&struct_value(), 0).unwrap(); + let cached = Subquery::new(1, 2, value.data_type()); + let same = Subquery::new(1, 2, value.data_type()); + let other_context = Subquery::new(3, 2, value.data_type()); + let hash = |expr: &Subquery| { + let mut hasher = DefaultHasher::new(); + expr.hash(&mut hasher); + hasher.finish() + }; + let before = hash(&cached); + cached.struct_value.set(value.clone()).unwrap(); + assert_eq!(cached, same); + assert_eq!(hash(&cached), before); + assert_ne!(cached, other_context); + // A cached result is owned by the expression and needs no live JVM registry entry. + let input = RecordBatch::new_empty(Arc::new(Schema::empty())); + let ColumnarValue::Scalar(result) = cached.evaluate(&input).unwrap() else { + panic!("Expected scalar result"); + }; + assert_eq!(result, value); + } +} diff --git a/native/jni-bridge/src/comet_exec.rs b/native/jni-bridge/src/comet_exec.rs index a0b39d0eaca..ecea48fc61c 100644 --- a/native/jni-bridge/src/comet_exec.rs +++ b/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"), diff --git a/spark/src/main/java/org/apache/spark/sql/comet/CometScalarSubquery.java b/spark/src/main/java/org/apache/spark/sql/comet/CometScalarSubquery.java index 29984ebb5ac..642a810ce6f 100644 --- a/spark/src/main/java/org/apache/spark/sql/comet/CometScalarSubquery.java +++ b/spark/src/main/java/org/apache/spark/sql/comet/CometScalarSubquery.java @@ -21,8 +21,11 @@ import java.util.HashMap; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.comet.execution.arrow.CometArrowConverters$; import org.apache.spark.sql.execution.ScalarSubquery; import org.apache.spark.sql.types.Decimal; +import org.apache.spark.sql.types.StructType; import org.apache.spark.unsafe.types.UTF8String; import org.apache.comet.CometRuntimeException; @@ -119,4 +122,11 @@ public static String getString(long planId, long id) { public static byte[] getBinary(long planId, long id) { return (byte[]) getSubquery(planId, id); } + + /** Get a struct subquery result as a one-row Arrow IPC stream. Called from native code. */ + public static byte[] getStruct(long planId, long id) { + InternalRow result = (InternalRow) getSubquery(planId, id); + StructType dataType = (StructType) subqueryMap.get(planId).get(id).dataType(); + return CometArrowConverters$.MODULE$.serializeScalarSubquery(result, dataType); + } } diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala index 1a82d789645..c808990e652 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala @@ -21,6 +21,7 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.execution.ScalarSubquery +import org.apache.spark.sql.types._ import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.serde.QueryPlanSerde.{serializeDataType, supportedDataType} @@ -28,21 +29,40 @@ import org.apache.comet.serde.QueryPlanSerde.{serializeDataType, supportedDataTy object CometScalarSubquery extends CometExpressionSerde[ScalarSubquery] { override def getUnsupportedReasons(): Seq[String] = Seq( - "Not all data types are supported for scalar subquery results") + "Not all data types are supported for scalar subquery results", + "Struct fields must have supported types and distinct names at every nesting level") - override def getSupportLevel(expr: ScalarSubquery): SupportLevel = - if (supportedDataType(expr.dataType)) { + // This is the value-transfer gate, not just a test that the type can be serialized to protobuf. + // Keep the scalar path unchanged; the Arrow IPC bridge only extends it to these struct shapes. + private def supportedStructField(dt: DataType): Boolean = dt match { + case s: StructType => + s.nonEmpty && s.fieldNames.distinct.length == s.length && + s.fields.forall(f => supportedStructField(f.dataType)) + case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | + StringType | BinaryType | DateType | TimestampType | TimestampNTZType | NullType => + true + case d: DecimalType => d.scale >= 0 && d.scale <= d.precision + case _ => false + } + + override def getSupportLevel(expr: ScalarSubquery): SupportLevel = { + val supported = expr.dataType match { + case s: StructType => supportedStructField(s) + case dt => supportedDataType(dt) + } + if (supported) { Compatible() } else { Unsupported(Some(s"Unsupported data type: ${expr.dataType}")) } + } override def convert( expr: ScalarSubquery, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - // getSupportLevel has already screened the data type with `supportedDataType`. That is a - // different predicate from `serializeDataType`, which can still decline, so keep this check. + // getSupportLevel has already checked value-transfer support. Type serialization can still + // decline, so keep this separate check. val dataType = serializeDataType(expr.dataType) if (dataType.isEmpty) { withFallbackReason( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala index e68eee6b79b..2f74af0c8a8 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala @@ -19,17 +19,25 @@ package org.apache.spark.sql.comet.execution.arrow +import java.io.ByteArrayOutputStream +import java.nio.channels.Channels + +import scala.util.Using import scala.util.control.NonFatal import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.ipc.ArrowStreamWriter import org.apache.arrow.vector.types.pojo.Schema import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.types.{StringType, StructField, StructType} import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.unsafe.types.UTF8String +import org.apache.comet.CometArrowAllocator import org.apache.comet.vector.NativeUtil /** @@ -46,6 +54,53 @@ import org.apache.comet.vector.NativeUtil */ object CometArrowConverters extends Logging { + /** + * Serialize a scalar subquery result as one struct column containing one row. Keeping the + * struct as a column preserves the distinction between a null struct and a struct whose fields + * are all null. Native execution expects TimestampType fields to carry the UTC zone. + */ + def serializeScalarSubquery(row: InternalRow, dataType: StructType): Array[Byte] = { + val schema = StructType(Seq(StructField("value", dataType, nullable = true))) + val output = new ByteArrayOutputStream() + Using.resource( + VectorSchemaRoot.create(Utils.toArrowSchema(schema, "UTC"), CometArrowAllocator)) { root => + val rowWriter = ArrowWriter.create(root, 1) + rowWriter.write(InternalRow(normalizeScalarSubqueryRow(row, dataType))) + rowWriter.finish() + Using.resource(new ArrowStreamWriter(root, null, Channels.newChannel(output))) { writer => + writer.start() + writer.writeBatch() + writer.end() + } + output.toByteArray + } + } + + /** + * Match CometScalarSubquery.getString's conversion through a JVM String, including replacement + * of malformed UTF-8. ArrowWriter copies raw UTF8String bytes, which Arrow IPC cannot represent + * as a valid string. The scalar-subquery support gate admits only structs and scalar leaves. + */ + private def normalizeScalarSubqueryRow(row: InternalRow, dataType: StructType): InternalRow = { + if (row == null) { + return null + } + val values = new Array[Any](dataType.length) + dataType.fields.zipWithIndex.foreach { case (field, ordinal) => + values(ordinal) = if (row.isNullAt(ordinal)) { + null + } else { + field.dataType match { + case _: StringType => UTF8String.fromString(row.getUTF8String(ordinal).toString) + case struct: StructType => + normalizeScalarSubqueryRow(row.getStruct(ordinal, struct.length), struct) + case dt => row.get(ordinal, dt) + } + } + } + new GenericInternalRow(values) + } + /** * Convert an iterator of Spark `InternalRow`s into an iterator of Arrow `ColumnarBatch`es. * diff --git a/spark/src/test/resources/sql-tests/expressions/misc/scalar_subquery_struct.sql b/spark/src/test/resources/sql-tests/expressions/misc/scalar_subquery_struct.sql new file mode 100644 index 00000000000..3710a5691fa --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/misc/scalar_subquery_struct.sql @@ -0,0 +1,143 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Config: spark.sql.session.timeZone=UTC + +statement +CREATE TABLE test_struct_subq( + id int, + payload struct>) USING parquet + +statement +INSERT INTO test_struct_subq VALUES + (1, named_struct( + 'flag', true, 'tiny', -128, 'small', -32768, 'number', -2147483648, + 'large', -9223372036854775808, 'single', 1.25, 'dbl', -2.5, + 'amount', 1234567890123456789012345678.1234567890, 'compact', -12345.67, + 'text', '中文-é', 'bytes', X'00FF41', 'day', DATE '1969-12-31', + 'instant', TIMESTAMP '1969-12-31 23:59:59.123456', + 'local_time', TIMESTAMP_NTZ '2024-02-29 12:34:56.654321', + 'nested', named_struct('last', 'tail', 'first', 9876543210))), + (2, named_struct( + 'flag', NULL, 'tiny', NULL, 'small', NULL, 'number', NULL, 'large', NULL, + 'single', NULL, 'dbl', NULL, 'amount', NULL, 'compact', NULL, 'text', NULL, + 'bytes', NULL, 'day', NULL, 'instant', NULL, 'local_time', NULL, 'nested', NULL)), + (3, NULL), + (4, named_struct( + 'flag', false, 'tiny', 127, 'small', 32767, 'number', 2147483647, + 'large', 9223372036854775807, 'single', 0.0, 'dbl', 3.5, + 'amount', -9999999999999999999999999999.9999999999, 'compact', 0.00, + 'text', '', 'bytes', X'', 'day', DATE '2000-02-29', + 'instant', TIMESTAMP '2024-02-29 12:34:56.654321', + 'local_time', TIMESTAMP_NTZ '1969-12-31 23:59:59.123456', + 'nested', named_struct('last', NULL, 'first', NULL))) + +-- Materialize the entire non-null scalar struct for every outer row. +query +SELECT id, (SELECT payload FROM test_struct_subq WHERE id = 1) AS s +FROM test_struct_subq + +-- Distinct field values and deliberately nonalphabetical nested names catch ordinal mixups. +query +SELECT id, s.flag, s.tiny, s.small, s.number, s.large, s.single, s.dbl, + s.amount, s.compact, s.text, s.bytes, s.day, s.instant, s.local_time, + s.nested.last, s.nested.first +FROM (SELECT id, (SELECT payload FROM test_struct_subq WHERE id = 1) AS s + FROM test_struct_subq) + +-- A present struct with all-null fields must not become a null struct. +query +SELECT id, s, s IS NULL, s.number, s.nested IS NULL +FROM (SELECT id, (SELECT payload FROM test_struct_subq WHERE id = 2) AS s + FROM test_struct_subq) + +-- A null struct result stays null when materialized and when its fields are extracted. +query +SELECT id, s, s IS NULL, s.number, s.nested, s.nested.first +FROM (SELECT id, (SELECT payload FROM test_struct_subq WHERE id = 3) AS s + FROM test_struct_subq) + +-- A present nested struct whose children are null has its own validity bit. +query +SELECT id, s, s.nested IS NULL, s.nested.last, s.nested.first +FROM (SELECT id, (SELECT payload FROM test_struct_subq WHERE id = 4) AS s + FROM test_struct_subq) + +-- A scalar subquery with no rows returns a null struct of the declared type. +query +SELECT id, s, s IS NULL, s.number, s.nested.first +FROM (SELECT id, (SELECT payload FROM test_struct_subq WHERE id = 99) AS s + FROM test_struct_subq) + +-- Separate scalar subqueries of the same type must retain separate results. +query +SELECT id, + (SELECT payload FROM test_struct_subq WHERE id = 1), + (SELECT payload FROM test_struct_subq WHERE id = 2), + (SELECT payload FROM test_struct_subq WHERE id = 3), + (SELECT payload FROM test_struct_subq WHERE id = 4) +FROM test_struct_subq + +-- Untyped null fields cannot be stored in Parquet, so construct them in the subquery. +query +SELECT id, (SELECT named_struct('untyped', NULL, 'value', max(id), + 'nested', named_struct('untyped', NULL, 'value', min(id))) + FROM test_struct_subq) AS s +FROM test_struct_subq + +statement +CREATE TABLE test_struct_subq_unsupported(id int, items array, entries map) +USING parquet + +statement +INSERT INTO test_struct_subq_unsupported VALUES + (1, array(10, NULL, 30), map('a', 10, 'b', NULL)), + (2, NULL, NULL) + +-- Supporting structs must not enable array or map scalar results. +query expect_fallback(Unsupported data type) +SELECT id, (SELECT items FROM test_struct_subq_unsupported WHERE id = 1) +FROM test_struct_subq_unsupported + +query expect_fallback(Unsupported data type) +SELECT id, (SELECT entries FROM test_struct_subq_unsupported WHERE id = 1) +FROM test_struct_subq_unsupported + +-- Unsupported fields must also be rejected recursively inside structs. +query expect_fallback(Unsupported data type) +SELECT id, (SELECT named_struct('nested', named_struct('items', items)) + FROM test_struct_subq_unsupported WHERE id = 1) +FROM test_struct_subq_unsupported + +query expect_fallback(Unsupported data type) +SELECT id, (SELECT named_struct('nested', named_struct('entries', entries)) + FROM test_struct_subq_unsupported WHERE id = 1) +FROM test_struct_subq_unsupported + +-- Duplicate field names are unsupported at the top level and in nested structs. +query expect_fallback(Unsupported data type) +SELECT id, (SELECT named_struct('same', id, 'same', id + 1) + FROM test_struct_subq_unsupported WHERE id = 1) +FROM test_struct_subq_unsupported + +query expect_fallback(Unsupported data type) +SELECT id, (SELECT named_struct('nested', named_struct('same', id, 'same', id + 1)) + FROM test_struct_subq_unsupported WHERE id = 1) +FROM test_struct_subq_unsupported diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 5b5d43dbe8a..0f8e02d330b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -31,7 +31,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier} import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogStatistics, CatalogTable} -import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, ExpressionInfo, Hex, Literal} +import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, ExpressionInfo, GetStructField, Hex, Literal, ScalarSubquery => LogicalScalarSubquery} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, BloomFilterAggregate} import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec} @@ -47,6 +47,7 @@ import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf.SESSION_LOCAL_TIMEZONE +import org.apache.spark.sql.types.StructType import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.{CometConf, CometExecIterator, ExtendedExplainInfo} @@ -2296,6 +2297,94 @@ class CometExecSuite extends CometTestBase { } } + test("scalar subqueries merged into a struct") { + Seq(false, true).foreach { aqeEnabled => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString, + SQLConf.SUBQUERY_REUSE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + withParquetTable((0 until 5).map(i => (i, i + 10)), "tbl") { + // MergeScalarSubqueries combines the aggregate results into one struct and reads its + // fields at each original scalar-subquery site. There is no explicit struct in the SQL. + val df = sql(""" + |SELECT _1, + | (SELECT max(_1) AS maximum FROM tbl) AS maximum, + | (SELECT sum(_2) AS total FROM tbl) AS total, + | (SELECT avg(_2) AS mean FROM tbl) AS mean + |FROM tbl + |""".stripMargin) + val mergedSubqueries = df.queryExecution.optimizedPlan.collect { case p => + p.expressions.flatMap(_.collect { + case GetStructField(s: LogicalScalarSubquery, _, _) + if s.dataType.isInstanceOf[StructType] => + s + }) + }.flatten + assert( + mergedSubqueries.nonEmpty, + s"Expected merged struct scalar subqueries:\n${df.queryExecution.optimizedPlan}") + assert(mergedSubqueries.exists(_.dataType.asInstanceOf[StructType].length == 3)) + + val (_, cometPlan) = + checkSparkAnswerAndOperator(df, Seq(classOf[CometProjectExec])) + val nativeStructSubqueries = stripAQEPlan(cometPlan).collect { + case p: CometProjectExec => + p.projectList.flatMap(_.collect { + case GetStructField(s: ScalarSubquery, _, _) + if s.dataType.isInstanceOf[StructType] => + s + }) + }.flatten + assert( + nativeStructSubqueries.nonEmpty, + s"Expected CometProjectExec to consume a struct scalar subquery:\n$cometPlan") + } + } + } + } + + test("merged one-row aggregate subplans retain native projection and union") { + assume(isSpark42Plus, "MergeSubplans merges bare aggregate subplans in Spark 4.2+") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SUBQUERY_REUSE_ENABLED.key -> "true") { + withParquetTable((0 until 100).map(i => (i, i * 2)), "tbl") { + // Regression for #5834: this SQL has no scalar subqueries. MergeSubplans introduces + // them, and a Spark projection at either site also prevents native union execution. + // Distinct aliases keep the merged struct outside the duplicate-name limitation. + val df = sql(""" + |SELECT sum(s) FROM ( + | SELECT max(_1) AS s FROM tbl + | UNION ALL + | SELECT min(_2) AS t FROM tbl) + |""".stripMargin) + val mergedSubqueries = df.queryExecution.optimizedPlan.collect { case p => + p.expressions.flatMap(_.collect { + case s: LogicalScalarSubquery if s.dataType.isInstanceOf[StructType] => s + }) + }.flatten + assert( + mergedSubqueries.exists(_.dataType.asInstanceOf[StructType].length == 2), + s"Expected MergeSubplans to introduce a struct scalar:\n${df.queryExecution.optimizedPlan}") + + val (_, cometPlan) = checkSparkAnswerAndOperator( + df, + Seq( + classOf[CometProjectExec], + classOf[CometUnionExec], + classOf[CometHashAggregateExec])) + val nativeStructSubqueries = stripAQEPlan(cometPlan).collect { case p: CometProjectExec => + p.projectList.flatMap(_.collect { + case s: ScalarSubquery if s.dataType.isInstanceOf[StructType] => s + }) + }.flatten + assert( + nativeStructSubqueries.nonEmpty, + s"Expected CometProjectExec to consume the introduced struct scalar:\n$cometPlan") + } + } + } + // Regression test for https://github.com/apache/datafusion-comet/issues/4787 // A scalar subquery inside a RepartitionByExpression (DISTRIBUTE BY) lives in the shuffle's // partitioning expressions, not the native child subtree, so it must be registered separately From 39b8d515a5670febec55e1ca3a3054c9ad24aeb2 Mon Sep 17 00:00:00 2001 From: LinSimon-901101 Date: Sun, 13 Sep 2026 13:27:49 +0800 Subject: [PATCH 2/2] test: update q9 plan for struct scalar subqueries --- .../approved-plans-v1_4/q9/extended.txt | 116 +++++++++--------- 1 file changed, 58 insertions(+), 58 deletions(-) diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q9/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q9/extended.txt index a6278c73836..3d2d7bb0cef 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q9/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q9/extended.txt @@ -1,61 +1,61 @@ - Project [COMET: Unsupported data type: StructType(StructField(count(1),LongType,false),StructField(avg(ss_ext_discount_amt),DecimalType(11,6),true),StructField(avg(ss_net_paid),DecimalType(11,6),true))] -: :- Subquery -: : +- CometColumnarToRow -: : +- CometProject -: : +- CometHashAggregate -: : +- CometExchange -: : +- CometHashAggregate -: : +- CometProject -: : +- CometFilter -: : +- CometNativeScan parquet spark_catalog.default.store_sales -: :- ReusedSubquery -: :- ReusedSubquery -: :- Subquery -: : +- CometColumnarToRow -: : +- CometProject -: : +- CometHashAggregate -: : +- CometExchange -: : +- CometHashAggregate -: : +- CometProject -: : +- CometFilter -: : +- CometNativeScan parquet spark_catalog.default.store_sales -: :- ReusedSubquery -: :- ReusedSubquery -: :- Subquery -: : +- CometColumnarToRow -: : +- CometProject -: : +- CometHashAggregate -: : +- CometExchange -: : +- CometHashAggregate -: : +- CometProject -: : +- CometFilter -: : +- CometNativeScan parquet spark_catalog.default.store_sales -: :- ReusedSubquery -: :- ReusedSubquery -: :- Subquery -: : +- CometColumnarToRow -: : +- CometProject -: : +- CometHashAggregate -: : +- CometExchange -: : +- CometHashAggregate -: : +- CometProject -: : +- CometFilter -: : +- CometNativeScan parquet spark_catalog.default.store_sales -: :- ReusedSubquery -: :- ReusedSubquery -: :- Subquery -: : +- CometColumnarToRow -: : +- CometProject -: : +- CometHashAggregate -: : +- CometExchange -: : +- CometHashAggregate -: : +- CometProject -: : +- CometFilter -: : +- CometNativeScan parquet spark_catalog.default.store_sales -: :- ReusedSubquery -: +- ReusedSubquery -+- CometColumnarToRow +CometColumnarToRow ++- CometProject + : :- Subquery + : : +- CometColumnarToRow + : : +- CometProject + : : +- CometHashAggregate + : : +- CometExchange + : : +- CometHashAggregate + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales + : :- ReusedSubquery + : :- ReusedSubquery + : :- Subquery + : : +- CometColumnarToRow + : : +- CometProject + : : +- CometHashAggregate + : : +- CometExchange + : : +- CometHashAggregate + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales + : :- ReusedSubquery + : :- ReusedSubquery + : :- Subquery + : : +- CometColumnarToRow + : : +- CometProject + : : +- CometHashAggregate + : : +- CometExchange + : : +- CometHashAggregate + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales + : :- ReusedSubquery + : :- ReusedSubquery + : :- Subquery + : : +- CometColumnarToRow + : : +- CometProject + : : +- CometHashAggregate + : : +- CometExchange + : : +- CometHashAggregate + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales + : :- ReusedSubquery + : :- ReusedSubquery + : :- Subquery + : : +- CometColumnarToRow + : : +- CometProject + : : +- CometHashAggregate + : : +- CometExchange + : : +- CometHashAggregate + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.store_sales + : :- ReusedSubquery + : +- ReusedSubquery +- CometFilter +- CometNativeScan parquet spark_catalog.default.reason -Comet accelerated 37 out of 43 eligible operators (86%). Final plan contains 6 transitions between Spark and Comet. Accelerated expressions: 11 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 38 out of 43 eligible operators (88%). Final plan contains 6 transitions between Spark and Comet. Accelerated expressions: 15 native, 0 codegen dispatch. \ No newline at end of file