peterxcli commented on code in PR #5889:
URL: https://github.com/apache/datafusion-comet/pull/5889#discussion_r3999220629
##########
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:
Can the same native `Subquery` instance be evaluated concurrently here?
Separate Spark tasks have separate instances. Understanding the sharing would
help decide whether lazy initialization is needed.
##########
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:
Related to #5025, which centralizes these recursive type checks. Could we
coordinate this with that helper instead of adding another independent
predicate? It already covers duplicate names, string types, and interval
restrictions. This path would still need its struct-only and decimal-scale
restrictions preserved.
##########
native/core/src/execution/expressions/subquery.rs:
##########
@@ -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<H: Hasher>(&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(
Review Comment:
The JVM writer always emits one row, one column, and one batch. Can these
shape checks fail through a supported production path? If not, could we
simplify them?
##########
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:
By native physical planning, Spark has materialized and registered the
subquery result. Could we resolve it into an immutable scalar there? What
requires deferring initialization until `evaluate`?
##########
native/core/src/execution/expressions/subquery.rs:
##########
@@ -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<H: Hasher>(&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<ScalarValue> {
+ let mut reader = StreamReader::try_new(Cursor::new(bytes), None)?;
+ let Some(batch) = reader.next().transpose()? else {
+ return internal_err!("Scalar subquery IPC result contains no batch");
+ };
+ if batch.num_rows() != 1 || batch.num_columns() != 1 {
+ return internal_err!("Scalar subquery IPC result must contain one row
and one column");
+ }
+ 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)?;
Review Comment:
The Rust test covers restoring metadata. Could we also exercise a struct
with Parquet field IDs through the JVM serializer, showing the schema mismatch
that occurs without this alignment?
--
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]
