LinSimon-901101 commented on code in PR #5889:
URL: https://github.com/apache/datafusion-comet/pull/5889#discussion_r4003213045
##########
spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala:
##########
@@ -21,28 +21,48 @@ 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}
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()
Review Comment:
Thanks for pointing this out. I've adopted the shared supportedDataType API
and implementation proposed in #5025, while retaining the scalar-subquery
bridge's struct-only and decimal-scale restrictions.
Code:
[QueryPlanSerde.scala:579–614](https://github.com/apache/datafusion-comet/blob/fb658a6918149d5b328c2cf5dea10c002b9d75ab/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala#L579-L614),
with bridge-specific restrictions in
[CometScalarSubquery.scala:35–55](https://github.com/apache/datafusion-comet/blob/fb658a6918149d5b328c2cf5dea10c002b9d75ab/spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala#L35-L55).
##########
native/core/src/execution/expressions/subquery.rs:
##########
@@ -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());
Review Comment:
Each Spark task creates its own native execution context, and struct results
are now resolved into owned, immutable Literals during native physical
planning. I've removed OnceLock and evaluation-time initialization, so the
implementation no longer depends on assumptions about concurrent evaluation.
Code: task-local iterator creation and subquery registration in
[CometExecRDD.scala:125–139](https://github.com/apache/datafusion-comet/blob/fb658a6918149d5b328c2cf5dea10c002b9d75ab/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala#L125-L139),
and immutable literal construction in
[planner.rs:786–798](https://github.com/apache/datafusion-comet/blob/fb658a6918149d5b328c2cf5dea10c002b9d75ab/native/core/src/execution/planner.rs#L786-L798).
##########
native/core/src/execution/expressions/subquery.rs:
##########
@@ -75,7 +157,10 @@ impl PhysicalExpr for Subquery {
}
fn evaluate(&self, _: &RecordBatch) ->
datafusion::common::Result<ColumnarValue> {
- JVMClasses::with_env(|env| unsafe {
+ if let Some(value) = self.struct_value.get() {
+ return Ok(ColumnarValue::Scalar(value.clone()));
+ }
Review Comment:
Since Spark completes the materialization and registration of results before
native physical planning, there's no reason to delay until evaluate. It's now
resolved to owned immutable Literals during planning, and OnceLock has been
removed.
Code: planning-time resolution in
[planner.rs:786–798](https://github.com/apache/datafusion-comet/blob/fb658a6918149d5b328c2cf5dea10c002b9d75ab/native/core/src/execution/planner.rs#L786-L798),
using [Subquery::resolve_struct in
subquery.rs:60–86](https://github.com/apache/datafusion-comet/blob/fb658a6918149d5b328c2cf5dea10c002b9d75ab/native/core/src/execution/expressions/subquery.rs#L60-L86).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
