This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 5b5142b1 feat: Add GetStructField expression (#731)
5b5142b1 is described below
commit 5b5142b1a248e838f70de351baff4509a4f60434
Author: Adam Binford <[email protected]>
AuthorDate: Fri Aug 2 20:35:03 2024 -0400
feat: Add GetStructField expression (#731)
* Add GetStructField support
* Add custom types to CometBatchScanExec
* Remove test explain
* Rust fmt
* Fix struct type support checks
* Support converting StructArray to native
* fix style
* Attempt to fix scalar subquery issue
* Fix other unit test
* Cleanup
* Default query plan supporting complex type to false
* Migrate struct expressions to spark-expr
* Update shouldApplyRowToColumnar comment
* Add nulls to test
* Rename to allowStruct
* Add DataTypeSupport trait
* Fix parquet datatype test
---
.../org/apache/spark/sql/comet/util/Utils.scala | 3 +-
.../src/execution/datafusion/expressions/mod.rs | 1 -
native/core/src/execution/datafusion/planner.rs | 9 +-
native/proto/src/proto/expr.proto | 6 ++
native/spark-expr/src/lib.rs | 2 +
.../src/structs.rs} | 106 ++++++++++++++++++++-
.../apache/comet/CometSparkSessionExtensions.scala | 44 ++++-----
.../scala/org/apache/comet/DataTypeSupport.scala | 72 ++++++++++++++
.../org/apache/comet/serde/QueryPlanSerde.scala | 20 +++-
.../spark/sql/comet/CometBatchScanExec.scala | 4 +-
.../spark/sql/comet/CometRowToColumnarExec.scala | 25 ++++-
.../org/apache/spark/sql/comet/CometScanExec.scala | 4 +-
.../org/apache/comet/CometExpressionSuite.scala | 21 +++-
.../comet/CometSparkSessionExtensionsSuite.scala | 37 -------
.../apache/comet/parquet/ParquetReadSuite.scala | 38 ++++++++
15 files changed, 313 insertions(+), 79 deletions(-)
diff --git a/common/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
b/common/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
index 2300e109..8d6a6334 100644
--- a/common/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
+++ b/common/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
@@ -28,6 +28,7 @@ import scala.collection.JavaConverters._
import org.apache.arrow.c.CDataDictionaryProvider
import org.apache.arrow.vector.{BigIntVector, BitVector, DateDayVector,
DecimalVector, FieldVector, FixedSizeBinaryVector, Float4Vector, Float8Vector,
IntVector, SmallIntVector, TimeStampMicroTZVector, TimeStampMicroVector,
TinyIntVector, ValueVector, VarBinaryVector, VarCharVector, VectorSchemaRoot}
import org.apache.arrow.vector.complex.MapVector
+import org.apache.arrow.vector.complex.StructVector
import org.apache.arrow.vector.dictionary.DictionaryProvider
import org.apache.arrow.vector.ipc.ArrowStreamWriter
import org.apache.arrow.vector.types._
@@ -258,7 +259,7 @@ object Utils {
case v @ (_: BitVector | _: TinyIntVector | _: SmallIntVector | _:
IntVector |
_: BigIntVector | _: Float4Vector | _: Float8Vector | _:
VarCharVector |
_: DecimalVector | _: DateDayVector | _: TimeStampMicroTZVector | _:
VarBinaryVector |
- _: FixedSizeBinaryVector | _: TimeStampMicroVector) =>
+ _: FixedSizeBinaryVector | _: TimeStampMicroVector | _:
StructVector) =>
v.asInstanceOf[FieldVector]
case _ =>
throw new SparkException(s"Unsupported Arrow Vector for $reason:
${valueVector.getClass}")
diff --git a/native/core/src/execution/datafusion/expressions/mod.rs
b/native/core/src/execution/datafusion/expressions/mod.rs
index ecd04956..4435f6b6 100644
--- a/native/core/src/execution/datafusion/expressions/mod.rs
+++ b/native/core/src/execution/datafusion/expressions/mod.rs
@@ -29,7 +29,6 @@ pub mod bloom_filter_might_contain;
pub mod comet_scalar_funcs;
pub mod correlation;
pub mod covariance;
-pub mod create_named_struct;
pub mod negative;
pub mod stats;
pub mod stddev;
diff --git a/native/core/src/execution/datafusion/planner.rs
b/native/core/src/execution/datafusion/planner.rs
index 8038e0fb..b4d723eb 100644
--- a/native/core/src/execution/datafusion/planner.rs
+++ b/native/core/src/execution/datafusion/planner.rs
@@ -94,7 +94,7 @@ use crate::{
},
};
-use super::expressions::{create_named_struct::CreateNamedStruct, EvalMode};
+use super::expressions::EvalMode;
use
crate::execution::datafusion::expressions::comet_scalar_funcs::create_comet_physical_fun;
use datafusion_comet_proto::{
spark_expression::{
@@ -109,7 +109,8 @@ use datafusion_comet_proto::{
spark_partitioning::{partitioning::PartitioningStruct, Partitioning as
SparkPartitioning},
};
use datafusion_comet_spark_expr::{
- Cast, DateTruncExpr, HourExpr, IfExpr, MinuteExpr, RLike, SecondExpr,
TimestampTruncExpr,
+ Cast, CreateNamedStruct, DateTruncExpr, GetStructField, HourExpr, IfExpr,
MinuteExpr, RLike,
+ SecondExpr, TimestampTruncExpr,
};
// For clippy error on type_complexity.
@@ -619,6 +620,10 @@ impl PhysicalPlanner {
let data_type =
to_arrow_datatype(expr.datatype.as_ref().unwrap());
Ok(Arc::new(CreateNamedStruct::new(values, data_type)))
}
+ ExprStruct::GetStructField(expr) => {
+ let child = self.create_expr(expr.child.as_ref().unwrap(),
input_schema.clone())?;
+ Ok(Arc::new(GetStructField::new(child, expr.ordinal as usize)))
+ }
expr => Err(ExecutionError::GeneralError(format!(
"Not implemented: {:?}",
expr
diff --git a/native/proto/src/proto/expr.proto
b/native/proto/src/proto/expr.proto
index 158356df..485fdcfd 100644
--- a/native/proto/src/proto/expr.proto
+++ b/native/proto/src/proto/expr.proto
@@ -78,6 +78,7 @@ message Expr {
UnboundReference unbound = 51;
BloomFilterMightContain bloom_filter_might_contain = 52;
CreateNamedStruct create_named_struct = 53;
+ GetStructField get_struct_field = 54;
}
}
@@ -492,6 +493,11 @@ message CreateNamedStruct {
DataType datatype = 2;
}
+message GetStructField {
+ Expr child = 1;
+ int32 ordinal = 2;
+}
+
enum SortDirection {
Ascending = 0;
Descending = 1;
diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs
index c9606233..cf7fc872 100644
--- a/native/spark-expr/src/lib.rs
+++ b/native/spark-expr/src/lib.rs
@@ -23,6 +23,7 @@ mod kernels;
mod regexp;
pub mod scalar_funcs;
pub mod spark_hash;
+mod structs;
mod temporal;
pub mod timezone;
pub mod utils;
@@ -32,6 +33,7 @@ pub use cast::{spark_cast, Cast};
pub use error::{SparkError, SparkResult};
pub use if_expr::IfExpr;
pub use regexp::RLike;
+pub use structs::{CreateNamedStruct, GetStructField};
pub use temporal::{DateTruncExpr, HourExpr, MinuteExpr, SecondExpr,
TimestampTruncExpr};
/// Spark supports three evaluation modes when evaluating expressions, which
affect
diff --git
a/native/core/src/execution/datafusion/expressions/create_named_struct.rs
b/native/spark-expr/src/structs.rs
similarity index 68%
rename from
native/core/src/execution/datafusion/expressions/create_named_struct.rs
rename to native/spark-expr/src/structs.rs
index c06aacf1..d6844ea8 100644
--- a/native/core/src/execution/datafusion/expressions/create_named_struct.rs
+++ b/native/spark-expr/src/structs.rs
@@ -19,9 +19,9 @@ use arrow::compute::take;
use arrow::record_batch::RecordBatch;
use arrow_array::types::Int32Type;
use arrow_array::{Array, DictionaryArray, StructArray};
-use arrow_schema::{DataType, Schema};
+use arrow_schema::{DataType, Field, Schema};
use datafusion::logical_expr::ColumnarValue;
-use datafusion_common::{DataFusionError, Result as DataFusionResult};
+use datafusion_common::{DataFusionError, Result as DataFusionResult,
ScalarValue};
use datafusion_physical_expr::PhysicalExpr;
use std::{
any::Any,
@@ -30,7 +30,7 @@ use std::{
sync::Arc,
};
-use crate::execution::datafusion::expressions::utils::down_cast_any_ref;
+use crate::utils::down_cast_any_ref;
#[derive(Debug, Hash)]
pub struct CreateNamedStruct {
@@ -142,6 +142,106 @@ impl PartialEq<dyn Any> for CreateNamedStruct {
}
}
+#[derive(Debug, Hash)]
+pub struct GetStructField {
+ child: Arc<dyn PhysicalExpr>,
+ ordinal: usize,
+}
+
+impl GetStructField {
+ pub fn new(child: Arc<dyn PhysicalExpr>, ordinal: usize) -> Self {
+ Self { child, ordinal }
+ }
+
+ fn child_field(&self, input_schema: &Schema) ->
DataFusionResult<Arc<Field>> {
+ match self.child.data_type(input_schema)? {
+ DataType::Struct(fields) => Ok(fields[self.ordinal].clone()),
+ data_type => Err(DataFusionError::Plan(format!(
+ "Expect struct field, got {:?}",
+ data_type
+ ))),
+ }
+ }
+}
+
+impl PhysicalExpr for GetStructField {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn data_type(&self, input_schema: &Schema) -> DataFusionResult<DataType> {
+ Ok(self.child_field(input_schema)?.data_type().clone())
+ }
+
+ fn nullable(&self, input_schema: &Schema) -> DataFusionResult<bool> {
+ Ok(self.child_field(input_schema)?.is_nullable())
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult<ColumnarValue>
{
+ let child_value = self.child.evaluate(batch)?;
+
+ match child_value {
+ ColumnarValue::Array(array) => {
+ let struct_array = array
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .expect("A struct is expected");
+
+ Ok(ColumnarValue::Array(
+ struct_array.column(self.ordinal).clone(),
+ ))
+ }
+ ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) =>
Ok(ColumnarValue::Array(
+ struct_array.column(self.ordinal).clone(),
+ )),
+ value => Err(DataFusionError::Execution(format!(
+ "Expected a struct array, got {:?}",
+ value
+ ))),
+ }
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+ vec![&self.child]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn PhysicalExpr>>,
+ ) -> datafusion_common::Result<Arc<dyn PhysicalExpr>> {
+ Ok(Arc::new(GetStructField::new(
+ children[0].clone(),
+ self.ordinal,
+ )))
+ }
+
+ fn dyn_hash(&self, state: &mut dyn Hasher) {
+ let mut s = state;
+ self.child.hash(&mut s);
+ self.ordinal.hash(&mut s);
+ self.hash(&mut s);
+ }
+}
+
+impl Display for GetStructField {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(
+ f,
+ "GetStructField [child: {:?}, ordinal: {:?}]",
+ self.child, self.ordinal
+ )
+ }
+}
+
+impl PartialEq<dyn Any> for GetStructField {
+ fn eq(&self, other: &dyn Any) -> bool {
+ down_cast_any_ref(other)
+ .downcast_ref::<Self>()
+ .map(|x| self.child.eq(&x.child) && self.ordinal.eq(&x.ordinal))
+ .unwrap_or(false)
+ }
+}
+
#[cfg(test)]
mod test {
use super::CreateNamedStruct;
diff --git
a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 96ec23fa..87697309 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -42,11 +42,10 @@ import
org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExc
import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec,
ShuffledHashJoinExec, SortMergeJoinExec}
import org.apache.spark.sql.execution.window.WindowExec
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types._
import org.apache.comet.CometConf._
import org.apache.comet.CometExplainInfo.getActualPlan
-import org.apache.comet.CometSparkSessionExtensions.{createMessage,
getCometBroadcastNotEnabledReason, getCometShuffleNotEnabledReason,
isANSIEnabled, isCometBroadCastForceEnabled, isCometEnabled,
isCometExecEnabled, isCometJVMShuffleMode, isCometNativeShuffleMode,
isCometOperatorEnabled, isCometScan, isCometScanEnabled, isCometShuffleEnabled,
isSchemaSupported, isSpark34Plus, isSpark40Plus, shouldApplyRowToColumnar,
withInfo, withInfos}
+import org.apache.comet.CometSparkSessionExtensions.{createMessage,
getCometBroadcastNotEnabledReason, getCometShuffleNotEnabledReason,
isANSIEnabled, isCometBroadCastForceEnabled, isCometEnabled,
isCometExecEnabled, isCometJVMShuffleMode, isCometNativeShuffleMode,
isCometOperatorEnabled, isCometScan, isCometScanEnabled, isCometShuffleEnabled,
isSpark34Plus, isSpark40Plus, shouldApplyRowToColumnar, withInfo, withInfos}
import org.apache.comet.parquet.{CometParquetScan, SupportsComet}
import org.apache.comet.serde.OperatorOuterClass.Operator
import org.apache.comet.serde.QueryPlanSerde
@@ -95,8 +94,10 @@ class CometSparkSessionExtensions
// data source V2
case scanExec: BatchScanExec
if scanExec.scan.isInstanceOf[ParquetScan] &&
-
isSchemaSupported(scanExec.scan.asInstanceOf[ParquetScan].readDataSchema) &&
-
isSchemaSupported(scanExec.scan.asInstanceOf[ParquetScan].readPartitionSchema)
&&
+ CometBatchScanExec.isSchemaSupported(
+ scanExec.scan.asInstanceOf[ParquetScan].readDataSchema) &&
+ CometBatchScanExec.isSchemaSupported(
+ scanExec.scan.asInstanceOf[ParquetScan].readPartitionSchema)
&&
// Comet does not support pushedAggregate
scanExec.scan.asInstanceOf[ParquetScan].pushedAggregate.isEmpty =>
val cometScan =
CometParquetScan(scanExec.scan.asInstanceOf[ParquetScan])
@@ -110,11 +111,11 @@ class CometSparkSessionExtensions
case scanExec: BatchScanExec if
scanExec.scan.isInstanceOf[ParquetScan] =>
val requiredSchema =
scanExec.scan.asInstanceOf[ParquetScan].readDataSchema
val info1 = createMessage(
- !isSchemaSupported(requiredSchema),
+ !CometBatchScanExec.isSchemaSupported(requiredSchema),
s"Schema $requiredSchema is not supported")
val readPartitionSchema =
scanExec.scan.asInstanceOf[ParquetScan].readPartitionSchema
val info2 = createMessage(
- !isSchemaSupported(readPartitionSchema),
+ !CometBatchScanExec.isSchemaSupported(readPartitionSchema),
s"Partition schema $readPartitionSchema is not supported")
// Comet does not support pushedAggregate
val info3 = createMessage(
@@ -129,7 +130,7 @@ class CometSparkSessionExtensions
// Iceberg scan, supported cases
case s: SupportsComet
if s.isCometEnabled &&
- isSchemaSupported(scanExec.scan.readSchema()) =>
+
CometBatchScanExec.isSchemaSupported(scanExec.scan.readSchema()) =>
logInfo(s"Comet extension enabled for
${scanExec.scan.getClass.getSimpleName}")
// When reading from Iceberg, we automatically enable type
promotion
SQLConf.get.setConfString(COMET_SCHEMA_EVOLUTION_ENABLED.key,
"true")
@@ -144,7 +145,7 @@ class CometSparkSessionExtensions
"Comet extension is not enabled for " +
s"${scanExec.scan.getClass.getSimpleName}: not enabled on
data source side")
val info2 = createMessage(
- !isSchemaSupported(scanExec.scan.readSchema()),
+
!CometBatchScanExec.isSchemaSupported(scanExec.scan.readSchema()),
"Comet extension is not enabled for " +
s"${scanExec.scan.getClass.getSimpleName}: Schema not
supported")
withInfos(scanExec, Seq(info1, info2).flatten.toSet)
@@ -166,7 +167,9 @@ class CometSparkSessionExtensions
_,
_,
_,
- _) if isSchemaSupported(requiredSchema) &&
isSchemaSupported(partitionSchema) =>
+ _)
+ if CometScanExec.isSchemaSupported(requiredSchema)
+ && CometScanExec.isSchemaSupported(partitionSchema) =>
logInfo("Comet extension enabled for v1 Scan")
CometScanExec(scanExec, session)
@@ -182,10 +185,10 @@ class CometSparkSessionExtensions
_,
_) =>
val info1 = createMessage(
- !isSchemaSupported(requiredSchema),
+ !CometScanExec.isSchemaSupported(requiredSchema),
s"Schema $requiredSchema is not supported")
val info2 = createMessage(
- !isSchemaSupported(partitionSchema),
+ !CometScanExec.isSchemaSupported(partitionSchema),
s"Partition schema $partitionSchema is not supported")
withInfo(scanExec, Seq(info1, info2).flatten.mkString(","))
scanExec
@@ -1109,28 +1112,17 @@ object CometSparkSessionExtensions extends Logging {
COMET_EXEC_ALL_OPERATOR_ENABLED.get(conf)
}
- private[comet] def isSchemaSupported(schema: StructType): Boolean =
- schema.map(_.dataType).forall(isTypeSupported)
-
- private[comet] def isTypeSupported(dt: DataType): Boolean = dt match {
- case BooleanType | ByteType | ShortType | IntegerType | LongType |
FloatType | DoubleType |
- BinaryType | StringType | _: DecimalType | DateType | TimestampType =>
- true
- case t: DataType if t.typeName == "timestamp_ntz" => true
- case dt =>
- logInfo(s"Comet extension is disabled because data type $dt is not
supported")
- false
- }
-
def isCometScan(op: SparkPlan): Boolean = {
op.isInstanceOf[CometBatchScanExec] || op.isInstanceOf[CometScanExec]
}
private def shouldApplyRowToColumnar(conf: SQLConf, op: SparkPlan): Boolean
= {
// Only consider converting leaf nodes to columnar currently, so that all
the following
- // operators can have a chance to be converted to columnar.
+ // operators can have a chance to be converted to columnar. Leaf operators
that output
+ // columnar batches, such as Spark's vectorized readers, will also be
converted to native
+ // comet batches.
// TODO: consider converting other intermediate operators to columnar.
- op.isInstanceOf[LeafExecNode] && !op.supportsColumnar &&
isSchemaSupported(op.schema) &&
+ op.isInstanceOf[LeafExecNode] &&
CometRowToColumnarExec.isSchemaSupported(op.schema) &&
COMET_ROW_TO_COLUMNAR_ENABLED.get(conf) && {
val simpleClassName = Utils.getSimpleName(op.getClass)
val nodeName = simpleClassName.replaceAll("Exec$", "")
diff --git a/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala
b/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala
new file mode 100644
index 00000000..c49a2c46
--- /dev/null
+++ b/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala
@@ -0,0 +1,72 @@
+/*
+ * 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
+
+import org.apache.spark.sql.types._
+
+trait DataTypeSupport {
+
+ /**
+ * This can be overridden to provide support for specific data types that
aren't universally
+ * supported. This will only be called for non-universally supported types.
+ *
+ * @param dt
+ * the data type to check
+ * @return
+ * true if the datatype is supported
+ */
+ def isAdditionallySupported(dt: DataType): Boolean = false
+
+ private def isGloballySupported(dt: DataType): Boolean = dt match {
+ case BooleanType | ByteType | ShortType | IntegerType | LongType |
FloatType | DoubleType |
+ BinaryType | StringType | _: DecimalType | DateType | TimestampType =>
+ true
+ case t: DataType if t.typeName == "timestamp_ntz" => true
+ case _ => false
+ }
+
+ /**
+ * Checks if this schema is supported by checking if each field in the
struct is supported.
+ *
+ * @param struct
+ * the struct to check the fields of
+ * @return
+ * true if all fields in the struct are supported
+ */
+ def isSchemaSupported(struct: StructType): Boolean = {
+ struct.fields.map(_.dataType).forall(isTypeSupported)
+ }
+
+ def isTypeSupported(dt: DataType): Boolean = {
+ if (isGloballySupported(dt) || isAdditionallySupported(dt)) {
+ // If complex types are supported, we additionally want to recurse into
their children
+ dt match {
+ case StructType(fields) =>
fields.map(_.dataType).forall(isTypeSupported)
+ case ArrayType(elementType, _) => isTypeSupported(elementType)
+ case MapType(keyType, valueType, _) =>
+ isTypeSupported(keyType) && isTypeSupported(valueType)
+ // Not a complex type
+ case _ => true
+ }
+ } else {
+ false
+ }
+ }
+}
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 3fbe6d74..98ca1404 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -59,12 +59,13 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
logWarning(s"Comet native execution is disabled due to: $reason")
}
- def supportedDataType(dt: DataType): Boolean = dt match {
+ def supportedDataType(dt: DataType, allowStruct: Boolean = false): Boolean =
dt match {
case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _:
FloatType |
_: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _:
DecimalType |
_: DateType | _: BooleanType | _: NullType =>
true
case dt if isTimestampNTZType(dt) => true
+ case s: StructType if allowStruct =>
s.fields.map(_.dataType).forall(supportedDataType(_))
case dt =>
emitWarning(s"unsupported Spark data type: $dt")
false
@@ -2199,7 +2200,7 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
ExprOuterClass.Expr.newBuilder().setNormalizeNanAndZero(builder).build()
}
- case s @ execution.ScalarSubquery(_, _) =>
+ case s @ execution.ScalarSubquery(_, _) if
supportedDataType(s.dataType) =>
val dataType = serializeDataType(s.dataType)
if (dataType.isEmpty) {
withInfo(s, s"Scalar subquery returns unsupported datatype
${s.dataType}")
@@ -2316,6 +2317,19 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
None
}
+ case GetStructField(child, ordinal, _) =>
+ exprToProto(child, inputs, binding).map { childExpr =>
+ val getStructFieldBuilder = ExprOuterClass.GetStructField
+ .newBuilder()
+ .setChild(childExpr)
+ .setOrdinal(ordinal)
+
+ ExprOuterClass.Expr
+ .newBuilder()
+ .setGetStructField(getStructFieldBuilder)
+ .build()
+ }
+
case _ =>
withInfo(expr, s"${expr.prettyName} is not supported",
expr.children: _*)
None
@@ -2819,7 +2833,7 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
withInfo(join, "SortMergeJoin is not enabled")
None
- case op if isCometSink(op) && op.output.forall(a =>
supportedDataType(a.dataType)) =>
+ case op if isCometSink(op) && op.output.forall(a =>
supportedDataType(a.dataType, true)) =>
// These operators are source of Comet native execution chain
val scanBuilder = OperatorOuterClass.Scan.newBuilder()
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/CometBatchScanExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/CometBatchScanExec.scala
index 82ebed95..6c21b93e 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBatchScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBatchScanExec.scala
@@ -34,7 +34,7 @@ import org.apache.spark.sql.vectorized._
import com.google.common.base.Objects
-import org.apache.comet.MetricsSupport
+import org.apache.comet.{DataTypeSupport, MetricsSupport}
import org.apache.comet.shims.ShimCometBatchScanExec
case class CometBatchScanExec(wrapped: BatchScanExec, runtimeFilters:
Seq[Expression])
@@ -152,3 +152,5 @@ case class CometBatchScanExec(wrapped: BatchScanExec,
runtimeFilters: Seq[Expres
override def supportsColumnar: Boolean = true
}
+
+object CometBatchScanExec extends DataTypeSupport
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/CometRowToColumnarExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/CometRowToColumnarExec.scala
index 5679e865..b8a5329c 100644
---
a/spark/src/main/scala/org/apache/spark/sql/comet/CometRowToColumnarExec.scala
+++
b/spark/src/main/scala/org/apache/spark/sql/comet/CometRowToColumnarExec.scala
@@ -19,6 +19,8 @@
package org.apache.spark.sql.comet
+import scala.collection.JavaConverters._
+
import org.apache.spark.TaskContext
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.rdd.RDD
@@ -28,8 +30,11 @@ import
org.apache.spark.sql.catalyst.plans.physical.Partitioning
import org.apache.spark.sql.comet.execution.arrow.CometArrowConverters
import org.apache.spark.sql.execution.{RowToColumnarTransition, SparkPlan}
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.types._
import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.comet.DataTypeSupport
+
case class CometRowToColumnarExec(child: SparkPlan)
extends RowToColumnarTransition
with CometPlan {
@@ -60,8 +65,17 @@ case class CometRowToColumnarExec(child: SparkPlan)
val timeZoneId = conf.sessionLocalTimeZone
val schema = child.schema
- child
- .execute()
+ val rdd: RDD[InternalRow] = if (child.supportsColumnar) {
+ child
+ .executeColumnar()
+ .mapPartitionsInternal { iter =>
+ iter.flatMap(_.rowIterator().asScala)
+ }
+ } else {
+ child.execute()
+ }
+
+ rdd
.mapPartitionsInternal { iter =>
val context = TaskContext.get()
CometArrowConverters.toArrowBatchIterator(
@@ -82,3 +96,10 @@ case class CometRowToColumnarExec(child: SparkPlan)
copy(child = newChild)
}
+
+object CometRowToColumnarExec extends DataTypeSupport {
+ override def isAdditionallySupported(dt: DataType): Boolean = dt match {
+ case _: StructType => true
+ case _ => false
+ }
+}
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala
index adaefb23..ce040f9f 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala
@@ -43,7 +43,7 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
import org.apache.spark.util.SerializableConfiguration
import org.apache.spark.util.collection._
-import org.apache.comet.{CometConf, MetricsSupport}
+import org.apache.comet.{CometConf, DataTypeSupport, MetricsSupport}
import org.apache.comet.parquet.{CometParquetFileFormat,
CometParquetPartitionReaderFactory}
/**
@@ -439,7 +439,7 @@ case class CometScanExec(
}
}
-object CometScanExec {
+object CometScanExec extends DataTypeSupport {
def apply(scanExec: FileSourceScanExec, session: SparkSession):
CometScanExec = {
// TreeNode.mapProductIterator is protected method.
def mapProductIterator[B: ClassTag](product: Product, f: Any => B):
Array[B] = {
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
index f169538a..27c046d7 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -28,7 +28,7 @@ import scala.util.Random
import org.apache.hadoop.fs.Path
import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
-import org.apache.spark.sql.functions.expr
+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.{Decimal, DecimalType}
@@ -1903,4 +1903,23 @@ class CometExpressionSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
}
+
+ test("get_struct_field") {
+ withSQLConf(
+ CometConf.COMET_ROW_TO_COLUMNAR_ENABLED.key -> "true",
+ CometConf.COMET_ROW_TO_COLUMNAR_SUPPORTED_OPERATOR_LIST.key ->
"FileSourceScan") {
+ withTempPath { dir =>
+ var df = spark
+ .range(5)
+ // Add both a null struct and null inner value
+ .select(when(col("id") > 1, struct(when(col("id") > 2,
col("id")).alias("id")))
+ .alias("nested"))
+
+ df.write.parquet(dir.toString())
+
+ df = spark.read.parquet(dir.toString())
+ checkSparkAnswerAndOperator(df.select("nested.id"))
+ }
+ }
+ }
}
diff --git
a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala
b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala
index fdcd307b..516fb09f 100644
---
a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala
+++
b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala
@@ -22,45 +22,8 @@ package org.apache.comet
import org.apache.spark.SparkConf
import org.apache.spark.sql._
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types._
class CometSparkSessionExtensionsSuite extends CometTestBase {
- test("unsupported Spark types") {
- Seq(
- NullType -> false,
- BooleanType -> true,
- ByteType -> true,
- ShortType -> true,
- IntegerType -> true,
- LongType -> true,
- FloatType -> true,
- DoubleType -> true,
- BinaryType -> true,
- StringType -> true,
- ArrayType(TimestampType) -> false,
- StructType(
- Seq(
- StructField("f1", DecimalType.SYSTEM_DEFAULT),
- StructField("f2", StringType))) -> false,
- MapType(keyType = LongType, valueType = DateType) -> false,
- StructType(Seq(StructField("f1", ByteType), StructField("f2",
StringType))) -> false,
- MapType(keyType = IntegerType, valueType = BinaryType) -> false).foreach
{
- case (dt, expected) =>
- assert(CometSparkSessionExtensions.isTypeSupported(dt) == expected)
- }
- }
-
- test("unsupported Spark schema") {
- Seq(
- Seq(StructField("f1", IntegerType), StructField("f2", BooleanType)) ->
true,
- Seq(StructField("f1", IntegerType), StructField("f2",
ArrayType(IntegerType))) -> false,
- Seq(
- StructField("f1", MapType(keyType = LongType, valueType = StringType)),
- StructField("f2", ArrayType(DoubleType))) -> false).foreach { case
(schema, expected) =>
- assert(CometSparkSessionExtensions.isSchemaSupported(StructType(schema))
== expected)
- }
- }
-
test("isCometEnabled") {
val conf = new SQLConf
diff --git
a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala
b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala
index 9e6cbc0a..3c5ed16c 100644
--- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala
@@ -82,6 +82,44 @@ abstract class ParquetReadSuite extends CometTestBase {
}
}
+ test("unsupported Spark types") {
+ Seq(
+ NullType -> false,
+ BooleanType -> true,
+ ByteType -> true,
+ ShortType -> true,
+ IntegerType -> true,
+ LongType -> true,
+ FloatType -> true,
+ DoubleType -> true,
+ BinaryType -> true,
+ StringType -> true,
+ ArrayType(TimestampType) -> false,
+ StructType(
+ Seq(
+ StructField("f1", DecimalType.SYSTEM_DEFAULT),
+ StructField("f2", StringType))) -> false,
+ MapType(keyType = LongType, valueType = DateType) -> false,
+ StructType(Seq(StructField("f1", ByteType), StructField("f2",
StringType))) -> false,
+ MapType(keyType = IntegerType, valueType = BinaryType) -> false).foreach
{
+ case (dt, expected) =>
+ assert(CometScanExec.isTypeSupported(dt) == expected)
+ assert(CometBatchScanExec.isTypeSupported(dt) == expected)
+ }
+ }
+
+ test("unsupported Spark schema") {
+ Seq(
+ Seq(StructField("f1", IntegerType), StructField("f2", BooleanType)) ->
true,
+ Seq(StructField("f1", IntegerType), StructField("f2",
ArrayType(IntegerType))) -> false,
+ Seq(
+ StructField("f1", MapType(keyType = LongType, valueType = StringType)),
+ StructField("f2", ArrayType(DoubleType))) -> false).foreach { case
(schema, expected) =>
+ assert(CometScanExec.isSchemaSupported(StructType(schema)) == expected)
+ assert(CometBatchScanExec.isSchemaSupported(StructType(schema)) ==
expected)
+ }
+ }
+
test("simple count") {
withParquetTable((0 until 10).map(i => (i, i.toString)), "tbl") {
assert(sql("SELECT * FROM tbl WHERE _1 % 2 == 0").count() == 5)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
