diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 6814d75769b..66295582e91 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.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 }}] diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 66a1b62ba3d..1e6922333c1 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.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 diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 208a5f3f124..83d9f6746e6 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 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. diff --git a/native/core/src/execution/expressions/subquery.rs b/native/core/src/execution/expressions/subquery.rs index fc7b8104d2a..8517c4c0cd6 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, + 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 { + 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 { + 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 { + 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:?}") + } + } } 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 { + 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.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); + } +} diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 5bd4bfebb77..3310dfb5438 100644 --- a/native/core/src/execution/planner.rs +++ b/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( 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/native/spark-expr/src/struct_funcs/get_struct_field.rs b/native/spark-expr/src/struct_funcs/get_struct_field.rs index b815b4ed9a0..07bd6261f65 100644 --- a/native/spark-expr/src/struct_funcs/get_struct_field.rs +++ b/native/spark-expr/src/struct_funcs/get_struct_field.rs @@ -99,9 +99,13 @@ impl PhysicalExpr for GetStructField { self.ordinal, )?)) } - ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => Ok(ColumnarValue::Array( - child_with_parent_nulls(&struct_array, self.ordinal)?, - )), + ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => { + let child = child_with_parent_nulls(&struct_array, self.ordinal)?; + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + child.as_ref(), + 0, + )?)) + } value => Err(DataFusionError::Execution(format!( "Expected a struct array, got {value:?}" ))), @@ -139,7 +143,106 @@ mod tests { use arrow::array::{ArrayRef, Int64Array}; use arrow::buffer::NullBuffer; use arrow::datatypes::Fields; - use datafusion::physical_expr::expressions::Column; + use datafusion::physical_expr::expressions::{Column, Literal}; + + fn assert_scalar_for_batch_sizes(expr: &GetStructField, expected: ScalarValue) { + for num_rows in [4, 1, 0] { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let input = Arc::new(Int64Array::from_iter_values(0..num_rows as i64)); + let batch = RecordBatch::try_new(schema, vec![input]).unwrap(); + let result = expr.evaluate(&batch).unwrap(); + + // A projection materializes scalars to its input batch length. Returning the + // one-element struct child as an array fails this check for multi-row/empty batches. + let output = result.clone().into_array_of_size(num_rows).unwrap(); + assert_eq!(output.len(), num_rows); + assert_eq!( + output.as_ref(), + expected.to_array_of_size(num_rows).unwrap().as_ref() + ); + match result { + ColumnarValue::Scalar(value) => assert_eq!(value, expected), + other => panic!("expected a scalar struct field, got {other:?}"), + } + } + } + + #[test] + fn scalar_field_is_broadcast_to_batch_length() { + let fields = Fields::from(vec![Field::new("value", DataType::Int64, false)]); + let child = Arc::new(Int64Array::from(vec![42_i64])) as ArrayRef; + let scalar = ScalarValue::Struct(Arc::new(StructArray::new(fields, vec![child], None))); + let expr = GetStructField::new(Arc::new(Literal::new(scalar)), 0); + + assert_scalar_for_batch_sizes(&expr, ScalarValue::Int64(Some(42))); + } + + #[test] + fn scalar_field_of_null_struct_is_null() { + let fields = Fields::from(vec![Field::new("value", DataType::Int64, false)]); + // The null parent hides a populated, non-nullable child buffer. + let child = Arc::new(Int64Array::from(vec![42_i64])) as ArrayRef; + let scalar = ScalarValue::Struct(Arc::new(StructArray::new( + fields, + vec![child], + Some(NullBuffer::from(vec![false])), + ))); + let expr = GetStructField::new(Arc::new(Literal::new(scalar)), 0); + + assert_scalar_for_batch_sizes(&expr, ScalarValue::Int64(None)); + } + + #[test] + fn scalar_null_field_is_null() { + let fields = Fields::from(vec![Field::new("value", DataType::Int64, true)]); + let child = Arc::new(Int64Array::from(vec![None::])) as ArrayRef; + let scalar = ScalarValue::Struct(Arc::new(StructArray::new(fields, vec![child], None))); + let expr = GetStructField::new(Arc::new(Literal::new(scalar)), 0); + + assert_scalar_for_batch_sizes(&expr, ScalarValue::Int64(None)); + } + + #[test] + fn nested_scalar_field_retains_scalar_semantics() { + for outer_valid in [true, false] { + for inner_valid in [true, false] { + for value in [Some(42_i64), None] { + let inner_fields = + Fields::from(vec![Field::new("value", DataType::Int64, true)]); + let inner = Arc::new(StructArray::new( + inner_fields.clone(), + vec![Arc::new(Int64Array::from(vec![value]))], + Some(NullBuffer::from(vec![inner_valid])), + )); + let outer_fields = Fields::from(vec![Field::new( + "nested", + DataType::Struct(inner_fields.clone()), + true, + )]); + let scalar = ScalarValue::Struct(Arc::new(StructArray::new( + outer_fields, + vec![Arc::clone(&inner) as ArrayRef], + Some(NullBuffer::from(vec![outer_valid])), + ))); + let nested = GetStructField::new(Arc::new(Literal::new(scalar)), 0); + let expected_inner = if outer_valid { + inner + } else { + Arc::new(StructArray::new_null(inner_fields, 1)) + }; + assert_scalar_for_batch_sizes(&nested, ScalarValue::Struct(expected_inner)); + + let leaf = GetStructField::new(Arc::new(nested), 0); + let expected_value = if outer_valid && inner_valid { + value + } else { + None + }; + assert_scalar_for_batch_sizes(&leaf, ScalarValue::Int64(expected_value)); + } + } + } + } // A field of a NULL struct must be NULL (Spark semantics) even when the child buffer holds a // non-null value at that row -- Arrow stores child validity independently of the parent 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..2926ae1b54d 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,44 @@ 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)) { + // The shared type gate handles general capabilities. Only the Arrow IPC scalar-subquery + // bridge's additional struct-only and decimal-scale restrictions belong here. + private def supportedStructShape(dt: DataType): Boolean = dt match { + case s: StructType => s.fields.forall(f => supportedStructShape(f.dataType)) + case _: ArrayType | _: MapType => false + case d: DecimalType => d.scale >= 0 && d.scale <= d.precision + case _ => true + } + + override def getSupportLevel(expr: ScalarSubquery): SupportLevel = { + val supported = expr.dataType match { + case s: StructType => + supportedDataType( + s, + allowComplex = true, + allowIntervals = false, + allowCalendarInterval = false, + allowTimeType = false, + allowAnyStringType = false, + allowDuplicateStructFieldNames = false) && supportedStructShape(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/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index be4bc9c3412..11e4b2df8a1 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -557,21 +557,60 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { builder.build() } - def supportedDataType(dt: DataType, allowComplex: Boolean = false): Boolean = dt match { - case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType | - _: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType | - _: DecimalType | _: DateType | _: BooleanType | _: NullType | CalendarIntervalType => - true - case dt if isTimeType(dt) => - true - case s: StructType if allowComplex => - s.fields.nonEmpty && s.fields.map(_.dataType).forall(supportedDataType(_, allowComplex)) - case a: ArrayType if allowComplex => - supportedDataType(a.elementType, allowComplex) - case m: MapType if allowComplex => - supportedDataType(m.keyType, allowComplex) && supportedDataType(m.valueType, allowComplex) - case _ => - false + /** + * Returns whether `dt` is supported at a caller's data-type boundary. All options apply + * recursively. Defaults preserve existing expression-serde behavior: complex and ANSI interval + * types are rejected, while calendar intervals, time types and all string variants are + * accepted. Callers must additionally check restrictions specific to their value-transfer path. + * + * @param allowComplex + * allow non-empty structs, arrays, and maps + * @param allowIntervals + * allow year-month and day-time intervals + * @param allowCalendarInterval + * allow calendar intervals + * @param allowTimeType + * allow Spark `TimeType` + * @param allowAnyStringType + * allow non-default string variants such as collated strings + * @param allowDuplicateStructFieldNames + * allow duplicate field names in structs + */ + def supportedDataType( + dt: DataType, + allowComplex: Boolean = false, + allowIntervals: Boolean = false, + allowCalendarInterval: Boolean = true, + allowTimeType: Boolean = true, + allowAnyStringType: Boolean = true, + allowDuplicateStructFieldNames: Boolean = true): Boolean = { + def supported(dt: DataType): Boolean = dt match { + case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType | + _: DoubleType | _: BinaryType | _: TimestampType | _: TimestampNTZType | + _: DecimalType | _: DateType | _: BooleanType | _: NullType => + true + case CalendarIntervalType if allowCalendarInterval => + true + case st: StringType if allowAnyStringType || st == StringType => + true + case _: YearMonthIntervalType | _: DayTimeIntervalType if allowIntervals => + true + case dt if allowTimeType && isTimeType(dt) => + true + case s: StructType if allowComplex => + s.fields.nonEmpty && + (allowDuplicateStructFieldNames || + s.fields.map(_.name).distinct.length == s.fields.length) && + s.fields.forall(f => supported(f.dataType)) + case a: ArrayType if allowComplex => + supported(a.elementType) + case m: MapType if allowComplex => + supported(m.keyType) && supported(m.valueType) + case _ => + false + } + + supported(dt) } /** 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/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 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..76745fbc2aa 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,230 @@ class CometExecSuite extends CometTestBase { } } + test("scalar subqueries merged into a struct") { + val numRows = 1024 + withTempPath { dir => + // Keep the consuming input in one file: local[5] writes the old five-row input as five + // one-row files, hiding a scalar field incorrectly returned as a one-element array. + (0 until numRows) + .map(i => (i, i + 10)) + .toDF("_1", "_2") + .coalesce(1) + .write + .parquet(dir.getCanonicalPath) + Seq(false, true).foreach { aqeEnabled => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString, + SQLConf.SUBQUERY_REUSE_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> numRows.toString, + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + withParquetTable(dir.getCanonicalPath, "tbl") { + assert(spark.table("tbl").inputFiles.length == 1) + // 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 nativeProjections = stripAQEPlan(cometPlan) + .collect { case p: CometProjectExec => + val fields = p.projectList.flatMap(_.collect { + case g @ GetStructField(s: ScalarSubquery, _, _) + if s.dataType.isInstanceOf[StructType] => + g + }) + (p, fields) + } + .filter(_._2.nonEmpty) + assert( + nativeProjections.nonEmpty, + s"Expected CometProjectExec to consume a struct scalar subquery:\n$cometPlan") + nativeProjections.foreach { case (projection, fields) => + assert(fields.map(_.ordinal).toSet == Set(0, 1, 2)) + assert(fields.forall(_.child.dataType.asInstanceOf[StructType].length == 3)) + // Projection preserves its input batch length. Read sizes inside each task while + // the iterator owns the batches, and bring only integers back to the driver. + val batchSizes = projection + .executeColumnar() + .mapPartitions(batches => batches.map(_.numRows())) + .collect() + assert(batchSizes.sum == numRows) + assert( + batchSizes.exists(_ > 1), + s"Expected a multi-row consuming batch, got ${batchSizes.mkString(", ")}") + } + } + } + } + } + } + + test("struct scalar subquery with nested Parquet field IDs") { + import org.apache.spark.sql.types.{IntegerType, LongType, MetadataBuilder, StringType, StructField} + + import org.apache.comet.vector.CometVector + + def fieldId(id: Long) = new MetadataBuilder().putLong("parquet.field.id", id).build() + val nestedType = StructType( + Seq( + StructField("number", IntegerType, nullable = true, fieldId(5)), + StructField("optional", LongType, nullable = true, fieldId(6)))) + val payloadType = StructType( + Seq( + StructField("label", StringType, nullable = true, fieldId(3)), + StructField("nested", nestedType, nullable = true, fieldId(4)))) + val schema = StructType( + Seq( + StructField("id", IntegerType, nullable = true, fieldId(1)), + StructField("payload", payloadType, nullable = true, fieldId(2)))) + val rows = + Seq(Row(1, Row("kept", Row(17, null))), Row(2, null), Row(3, Row("other", Row(-91, 42L)))) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key -> "true", + SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "false") { + withTempPath { dir => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .createDataFrame(spark.sparkContext.parallelize(rows, 1), schema) + .write + .parquet(dir.getCanonicalPath) + } + withTempView("struct_metadata_subquery") { + spark.read + .schema(schema) + .parquet(dir.getCanonicalPath) + .createOrReplaceTempView("struct_metadata_subquery") + Seq(1, 2).foreach { id => + val df = sql(s""" + |SELECT id, + | (SELECT payload FROM struct_metadata_subquery WHERE id = $id) AS s + |FROM struct_metadata_subquery + |""".stripMargin) + val structSubqueries = stripAQEPlan(df.queryExecution.executedPlan).collect { + case p: CometProjectExec => + p.projectList.flatMap(_.collect { + case s: ScalarSubquery if s.dataType.isInstanceOf[StructType] => s + }) + }.flatten + assert(structSubqueries.nonEmpty) + structSubqueries.foreach { subquery => + val resultType = subquery.dataType.asInstanceOf[StructType] + assert(resultType == payloadType) + assert(resultType("label").metadata.getLong("parquet.field.id") == 3L) + assert(resultType("nested").metadata.getLong("parquet.field.id") == 4L) + val nested = resultType("nested").dataType.asInstanceOf[StructType] + assert(nested("number").metadata.getLong("parquet.field.id") == 5L) + assert(nested("optional").metadata.getLong("parquet.field.id") == 6L) + } + // Exercise native whole-struct output for both present and NULL results. The companion + // JVM IPC test pins the metadata difference between the wire and planned types. + val (_, cometPlan) = + checkSparkAnswerAndOperator(df, Seq(classOf[CometProjectExec])) + val nativeProjections = stripAQEPlan(cometPlan).collect { + case p: CometProjectExec if p.projectList.exists(_.exists { + case s: ScalarSubquery => s.dataType == payloadType + case _ => false + }) => + p + } + assert(nativeProjections.nonEmpty) + nativeProjections.foreach { projection => + assert(projection.output.map(_.name) == Seq("id", "s")) + // Inspect the schema imported from native output while the iterator owns each + // batch. Return only ordinary Scala values, never live Arrow fields or vectors. + val batches = projection + .executeColumnar() + .mapPartitions { iter => + iter.map { batch => + val value = batch.column(1).asInstanceOf[CometVector] + val fields = value.getValueVector.getField.getChildren + val nested = fields.get(1).getChildren + val metadata = Seq(fields.get(0), fields.get(1), nested.get(0), nested.get(1)) + .map(field => + field.getName -> Option(field.getMetadata.get("PARQUET:field_id"))) + val nulls = (0 until batch.numRows()).count(value.isNullAt) + (batch.numRows(), metadata, nulls) + } + } + .collect() + assert(batches.map(_._1).sum == rows.size) + batches.foreach { case (size, metadata, nulls) => + assert( + metadata == Seq( + "label" -> Some("3"), + "nested" -> Some("4"), + "number" -> Some("5"), + "optional" -> Some("6"))) + assert(nulls == (if (id == 2) size else 0)) + } + } + } + } + } + } + } + + 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 diff --git a/spark/src/test/scala/org/apache/comet/serde/CometScalarSubquerySuite.scala b/spark/src/test/scala/org/apache/comet/serde/CometScalarSubquerySuite.scala new file mode 100644 index 00000000000..46d385750f4 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/serde/CometScalarSubquerySuite.scala @@ -0,0 +1,204 @@ +/* + * 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. + */ + +package org.apache.comet.serde + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.expressions.{Alias, Literal, NamedExpression} +import org.apache.spark.sql.execution.{ProjectExec, ScalarSubquery, SubqueryExec} +import org.apache.spark.sql.types._ + +import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus} +import org.apache.comet.serde.QueryPlanSerde.supportedDataType +import org.apache.comet.shims.CometTypeShim + +/** Direct type-gate tests avoid optimizer folding and exercise types Parquet cannot store. */ +class CometScalarSubquerySuite extends CometTestBase with CometTypeShim { + + private def struct(dt: DataType): StructType = StructType(Seq(StructField("value", dt))) + + private lazy val emptyInput = spark.range(0).queryExecution.sparkPlan + + private def subquery(dt: DataType): ScalarSubquery = { + // Inspect the declared result type without executing or optimizing the typed NULL away. + val plan = ProjectExec(Seq(Alias(Literal.create(null, dt), "result")()), emptyInput) + ScalarSubquery(SubqueryExec("type-check", plan), NamedExpression.newExprId) + } + + private def supported(dt: DataType): Boolean = + CometScalarSubquery.getSupportLevel(subquery(dt)) == Compatible() + + private def versionSpecificTypes: Seq[DataType] = { + val strings = + if (isSpark40Plus) Seq(DataType.fromDDL("STRING COLLATE UTF8_LCASE")) else Seq.empty + val times = if (isSpark41Plus) Seq(DataType.fromDDL("TIME")) else Seq.empty + strings ++ times ++ variantType.toSeq + } + + private val scalarTypes: Seq[DataType] = Seq( + BooleanType, + ByteType, + ShortType, + IntegerType, + LongType, + FloatType, + DoubleType, + StringType, + BinaryType, + DateType, + TimestampType, + TimestampNTZType, + NullType, + DecimalType(1, 0), + DecimalType(10, 2), + DecimalType(38, 38)) + + // Freeze the pre-refactor shared predicate: existing callers must retain their accepted types. + private def legacySupported(dt: DataType, allowComplex: Boolean): Boolean = dt match { + case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType | + _: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType | + _: DecimalType | _: DateType | _: BooleanType | _: NullType | CalendarIntervalType => + true + case dt if isTimeType(dt) => true + case s: StructType if allowComplex => + s.nonEmpty && s.fields.forall(f => legacySupported(f.dataType, allowComplex)) + case a: ArrayType if allowComplex => legacySupported(a.elementType, allowComplex) + case m: MapType if allowComplex => + legacySupported(m.keyType, allowComplex) && legacySupported(m.valueType, allowComplex) + case _ => false + } + + private def nestedTypes(dt: DataType): Seq[DataType] = Seq( + dt, + struct(dt), + struct(struct(dt)), + ArrayType(dt), + MapType(dt, IntegerType), + MapType(IntegerType, dt), + struct(ArrayType(MapType(StringType, dt)))) + + test("shared type gate preserves existing defaults and non-struct scalar subqueries") { + val types = scalarTypes ++ versionSpecificTypes ++ Seq( + CalendarIntervalType, + YearMonthIntervalType(), + DayTimeIntervalType(), + CharType(5), + VarcharType(5), + StructType(Nil), + StructType(Seq(StructField("same", IntegerType), StructField("same", LongType)))) + types.flatMap(nestedTypes).foreach { dt => + Seq(false, true).foreach { allowComplex => + withClue(s"$dt, allowComplex=$allowComplex: ") { + assert(supportedDataType(dt, allowComplex) == legacySupported(dt, allowComplex)) + } + } + if (!dt.isInstanceOf[StructType]) { + assert(supported(dt) == legacySupported(dt, allowComplex = false), dt) + } + } + } + + test("shared capability flags apply recursively to structs arrays and map keys and values") { + val duplicate = + StructType(Seq(StructField("same", IntegerType), StructField("same", LongType))) + val cases: Seq[(DataType, DataType => Boolean, DataType => Boolean)] = Seq( + ( + CalendarIntervalType, + (t: DataType) => supportedDataType(t, allowComplex = true), + (t: DataType) => + supportedDataType(t, allowComplex = true, allowCalendarInterval = false)), + ( + YearMonthIntervalType(), + (t: DataType) => supportedDataType(t, allowComplex = true, allowIntervals = true), + (t: DataType) => supportedDataType(t, allowComplex = true)), + ( + DayTimeIntervalType(), + (t: DataType) => supportedDataType(t, allowComplex = true, allowIntervals = true), + (t: DataType) => supportedDataType(t, allowComplex = true)), + ( + duplicate, + (t: DataType) => supportedDataType(t, allowComplex = true), + (t: DataType) => + supportedDataType(t, allowComplex = true, allowDuplicateStructFieldNames = false))) ++ + versionSpecificTypes + .filter(isTimeType) + .map(dt => + ( + dt, + (t: DataType) => supportedDataType(t, allowComplex = true), + (t: DataType) => supportedDataType(t, allowComplex = true, allowTimeType = false))) ++ + versionSpecificTypes.collect { case dt: StringType => + ( + dt, + (t: DataType) => supportedDataType(t, allowComplex = true), + (t: DataType) => supportedDataType(t, allowComplex = true, allowAnyStringType = false)) + } + cases.foreach { case (dt, accepts, rejects) => + nestedTypes(dt).foreach { nested => + withClue(s"$nested: ") { + assert(accepts(nested)) + assert(!rejects(nested)) + } + } + } + } + + test("struct scalar subqueries retain supported fields and case-distinct names") { + val distinct = + StructType(Seq(StructField("a", IntegerType), StructField("A", LongType))) + (scalarTypes.map(struct) ++ scalarTypes.map(dt => struct(struct(dt))) ++ + Seq(distinct, struct(distinct))).foreach { dt => + assert(supported(dt), dt) + assert(CometScalarSubquery.convert(subquery(dt), Seq.empty, binding = false).isDefined, dt) + } + } + + test("struct scalar subqueries reject unsupported fields and shapes recursively") { + val rejected = versionSpecificTypes ++ Seq( + ArrayType(IntegerType), + MapType(StringType, IntegerType), + StructType(Nil), + StructType(Seq(StructField("same", IntegerType), StructField("same", LongType))), + CalendarIntervalType, + YearMonthIntervalType(), + DayTimeIntervalType(), + CharType(5), + VarcharType(5)) + rejected.foreach { dt => + Seq(struct(dt), struct(struct(dt))).foreach { nested => + assert(!supported(nested), nested) + } + } + rejected.collect { case s: StructType => s }.foreach(s => assert(!supported(s), s)) + } + + test("struct decimal scale restrictions do not change the legacy non-struct gate") { + withSQLConf("spark.sql.legacy.allowNegativeScaleOfDecimal" -> "true") { + val negative = DecimalType(10, -2) + assert(supportedDataType(struct(negative), allowComplex = true)) + assert(supported(negative)) + assert(!supported(struct(negative))) + assert(!supported(struct(struct(negative)))) + Seq(DecimalType(1, 0), DecimalType(38, 0), DecimalType(38, 38)).foreach { dt => + assert(supported(struct(dt)), dt) + assert(supported(struct(struct(dt))), dt) + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index d95a4c718e5..fc3dbe84c83 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.comet.execution.arrow +import java.io.ByteArrayInputStream import java.nio.ByteOrder import scala.jdk.CollectionConverters._ @@ -28,14 +29,17 @@ import org.scalatest.matchers.should.Matchers import org.apache.arrow.memory.{AllocationListener, RootAllocator} import org.apache.arrow.vector.{BaseFixedWidthVector, BaseValueVector, BigIntVector, BitVector, DecimalVector, IntervalMonthDayNanoVector, IntVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.StructVector +import org.apache.arrow.vector.ipc.ArrowStreamReader import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, SpecializedGetters} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.vectorized.{ConstantColumnVector, Dictionary, OffHeapColumnVector, OnHeapColumnVector} -import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType, YearMonthIntervalType} +import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, Metadata, MetadataBuilder, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType, YearMonthIntervalType} import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch, ColumnVector} -import org.apache.spark.unsafe.types.CalendarInterval +import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String} +import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.vector.{CometPlainVector, CometVector, NativeUtil} /** @@ -58,6 +62,83 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { new ColumnarBatch(vectors.toArray, numRows) } + test("scalar subquery IPC omits nested Parquet field IDs retained by the planned type") { + def fieldId(id: Long) = new MetadataBuilder().putLong("parquet.field.id", id).build() + val nestedType = StructType( + Seq( + StructField("label", StringType, nullable = true, fieldId(41)), + StructField("value", LongType, nullable = false, fieldId(42)))) + val valueType = StructType( + Seq( + StructField("number", IntegerType, nullable = false, fieldId(10)), + StructField("untagged", StringType, nullable = true), + StructField("nested", nestedType, nullable = true, fieldId(30)))) + val row = new GenericInternalRow( + Array[Any](17, null, new GenericInternalRow(Array[Any](UTF8String.fromString("尾"), -91L)))) + + val planned = QueryPlanSerde.serializeDataType(valueType).get.getTypeInfo.getStruct + val plannedMetadata = planned.getFieldMetadataList.asScala.map(_.getMetadataMap.asScala.toMap) + plannedMetadata.toSeq shouldBe Seq( + Map("PARQUET:field_id" -> "10"), + Map.empty[String, String], + Map("PARQUET:field_id" -> "30")) + val plannedNested = planned.getFieldDatatypes(2).getTypeInfo.getStruct + val plannedNestedMetadata = + plannedNested.getFieldMetadataList.asScala.map(_.getMetadataMap.asScala.toMap) + plannedNestedMetadata.toSeq shouldBe Seq( + Map("PARQUET:field_id" -> "41"), + Map("PARQUET:field_id" -> "42")) + + // Read the real JVM serializer's IPC bytes. A synthetic Rust-only schema would not establish + // whether the JVM actually drops the metadata that native struct reconstruction restores. + val bytes = CometArrowConverters.serializeScalarSubquery(row, valueType) + val allocator = new RootAllocator(Long.MaxValue) + val reader = new ArrowStreamReader(new ByteArrayInputStream(bytes), allocator) + try { + reader.loadNextBatch() shouldBe true + val root = reader.getVectorSchemaRoot + root.getRowCount shouldBe 1 + root.getFieldVectors.size() shouldBe 1 + val valueField = root.getSchema.getFields.get(0) + valueField.getName shouldBe "value" + valueField.isNullable shouldBe true + val fields = valueField.getChildren.asScala + fields.map(_.getName).toSeq shouldBe Seq("number", "untagged", "nested") + fields.map(_.isNullable).toSeq shouldBe Seq(false, true, true) + val wireMetadata = fields.map(_.getMetadata.asScala.toMap) + wireMetadata.toSeq shouldBe Seq.fill(3)(Map.empty[String, String]) + wireMetadata should not equal plannedMetadata + val nestedFields = fields(2).getChildren.asScala + nestedFields.map(_.getName).toSeq shouldBe Seq("label", "value") + nestedFields.map(_.isNullable).toSeq shouldBe Seq(true, false) + val wireNestedMetadata = nestedFields.map(_.getMetadata.asScala.toMap) + wireNestedMetadata.toSeq shouldBe Seq.fill(2)(Map.empty[String, String]) + wireNestedMetadata should not equal plannedNestedMetadata + + // The mismatch is metadata, not a different result, field order, type, or nullability. + val untaggedNestedType = StructType(nestedType.map(_.copy(metadata = Metadata.empty))) + val untaggedType = StructType( + Seq( + StructField("number", IntegerType, nullable = false), + StructField("untagged", StringType, nullable = true), + StructField("nested", untaggedNestedType, nullable = true))) + Utils.fromArrowField(valueField) shouldBe untaggedType + Utils.fromArrowField(valueField) should not equal valueType + val value = root.getVector(0).asInstanceOf[StructVector] + value.isNull(0) shouldBe false + value.getChildByOrdinal(0).asInstanceOf[IntVector].get(0) shouldBe 17 + value.getChildByOrdinal(1).isNull(0) shouldBe true + val nested = value.getChildByOrdinal(2).asInstanceOf[StructVector] + nested.isNull(0) shouldBe false + nested.getChildByOrdinal(0).getObject(0).toString shouldBe "尾" + nested.getChildByOrdinal(1).asInstanceOf[BigIntVector].get(0) shouldBe -91L + reader.loadNextBatch() shouldBe false + } finally { + reader.close() + allocator.close() + } + } + test("CalendarIntervalType round-trips through Arrow writer and Comet vector") { val allocator = new RootAllocator(Integer.MAX_VALUE) val field = Utils.toArrowField("interval", CalendarIntervalType, nullable = true, "UTC")