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 712658e2 feat: Implement CAST between struct types (#1074)
712658e2 is described below
commit 712658e2966e33b0f56fcc5449a9d8aab33a5d58
Author: Andy Grove <[email protected]>
AuthorDate: Mon Nov 11 14:25:26 2024 -0700
feat: Implement CAST between struct types (#1074)
* implement basic native code for casting struct to struct
* add another test
* rustdoc
* add scala side
* code cleanup
* clippy
* clippy
* add scala test
* improve test
* remove assert
* clippy
---
native/spark-expr/src/cast.rs | 112 ++++++++++++++++++++-
.../org/apache/comet/expressions/CometCast.scala | 10 ++
.../scala/org/apache/comet/CometCastSuite.scala | 14 +++
3 files changed, 134 insertions(+), 2 deletions(-)
diff --git a/native/spark-expr/src/cast.rs b/native/spark-expr/src/cast.rs
index 0a2f7fef..13263a59 100644
--- a/native/spark-expr/src/cast.rs
+++ b/native/spark-expr/src/cast.rs
@@ -34,7 +34,7 @@ use arrow::{
};
use arrow_array::builder::StringBuilder;
use arrow_array::{DictionaryArray, StringArray, StructArray};
-use arrow_schema::{DataType, Schema};
+use arrow_schema::{DataType, Field, Schema};
use datafusion_common::{
cast::as_generic_string_array, internal_err, Result as DataFusionResult,
ScalarValue,
};
@@ -714,6 +714,14 @@ fn cast_array(
(DataType::Struct(_), DataType::Utf8) => {
Ok(casts_struct_to_string(array.as_struct(), &timezone)?)
}
+ (DataType::Struct(_), DataType::Struct(_)) => Ok(cast_struct_to_struct(
+ array.as_struct(),
+ from_type,
+ to_type,
+ eval_mode,
+ timezone,
+ allow_incompat,
+ )?),
_ if is_datafusion_spark_compatible(from_type, to_type,
allow_incompat) => {
// use DataFusion cast only when we know that it is compatible
with Spark
Ok(cast_with_options(&array, to_type, &CAST_OPTIONS)?)
@@ -811,6 +819,35 @@ fn is_datafusion_spark_compatible(
}
}
+/// Cast between struct types based on logic in
+/// `org.apache.spark.sql.catalyst.expressions.Cast#castStruct`.
+fn cast_struct_to_struct(
+ array: &StructArray,
+ from_type: &DataType,
+ to_type: &DataType,
+ eval_mode: EvalMode,
+ timezone: String,
+ allow_incompat: bool,
+) -> DataFusionResult<ArrayRef> {
+ match (from_type, to_type) {
+ (DataType::Struct(_), DataType::Struct(to_fields)) => {
+ let mut cast_fields: Vec<(Arc<Field>, ArrayRef)> =
Vec::with_capacity(to_fields.len());
+ for i in 0..to_fields.len() {
+ let cast_field = cast_array(
+ Arc::clone(array.column(i)),
+ to_fields[i].data_type(),
+ eval_mode,
+ timezone.clone(),
+ allow_incompat,
+ )?;
+ cast_fields.push((Arc::clone(&to_fields[i]), cast_field));
+ }
+ Ok(Arc::new(StructArray::from(cast_fields)))
+ }
+ _ => unreachable!(),
+ }
+}
+
fn casts_struct_to_string(array: &StructArray, timezone: &str) ->
DataFusionResult<ArrayRef> {
// cast each field to a string
let string_arrays: Vec<ArrayRef> = array
@@ -1929,7 +1966,7 @@ fn trim_end(s: &str) -> &str {
mod tests {
use arrow::datatypes::TimestampMicrosecondType;
use arrow_array::StringArray;
- use arrow_schema::{Field, TimeUnit};
+ use arrow_schema::{Field, Fields, TimeUnit};
use std::str::FromStr;
use super::*;
@@ -2336,4 +2373,75 @@ mod tests {
assert_eq!(r#"{4, d}"#, string_array.value(3));
assert_eq!(r#"{5, e}"#, string_array.value(4));
}
+
+ #[test]
+ fn test_cast_struct_to_struct() {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![
+ Some(1),
+ Some(2),
+ None,
+ Some(4),
+ Some(5),
+ ]));
+ let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d",
"e"]));
+ let c: ArrayRef = Arc::new(StructArray::from(vec![
+ (Arc::new(Field::new("a", DataType::Int32, true)), a),
+ (Arc::new(Field::new("b", DataType::Utf8, true)), b),
+ ]));
+ // change type of "a" from Int32 to Utf8
+ let fields = Fields::from(vec![
+ Field::new("a", DataType::Utf8, true),
+ Field::new("b", DataType::Utf8, true),
+ ]);
+ let cast_array = spark_cast(
+ ColumnarValue::Array(c),
+ &DataType::Struct(fields),
+ EvalMode::Legacy,
+ "UTC",
+ false,
+ )
+ .unwrap();
+ if let ColumnarValue::Array(cast_array) = cast_array {
+ assert_eq!(5, cast_array.len());
+ let a = cast_array.as_struct().column(0).as_string::<i32>();
+ assert_eq!("1", a.value(0));
+ } else {
+ unreachable!()
+ }
+ }
+
+ #[test]
+ fn test_cast_struct_to_struct_drop_column() {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![
+ Some(1),
+ Some(2),
+ None,
+ Some(4),
+ Some(5),
+ ]));
+ let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d",
"e"]));
+ let c: ArrayRef = Arc::new(StructArray::from(vec![
+ (Arc::new(Field::new("a", DataType::Int32, true)), a),
+ (Arc::new(Field::new("b", DataType::Utf8, true)), b),
+ ]));
+ // change type of "a" from Int32 to Utf8 and drop "b"
+ let fields = Fields::from(vec![Field::new("a", DataType::Utf8, true)]);
+ let cast_array = spark_cast(
+ ColumnarValue::Array(c),
+ &DataType::Struct(fields),
+ EvalMode::Legacy,
+ "UTC",
+ false,
+ )
+ .unwrap();
+ if let ColumnarValue::Array(cast_array) = cast_array {
+ assert_eq!(5, cast_array.len());
+ let struct_array = cast_array.as_struct();
+ assert_eq!(1, struct_array.columns().len());
+ let a = struct_array.column(0).as_string::<i32>();
+ assert_eq!("1", a.value(0));
+ } else {
+ unreachable!()
+ }
+ }
}
diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
index eb9800b8..11d6d049 100644
--- a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
+++ b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
@@ -95,6 +95,16 @@ object CometCast {
canCastFromFloat(toType)
case (DataTypes.DoubleType, _) =>
canCastFromDouble(toType)
+ case (from_struct: StructType, to_struct: StructType) =>
+ from_struct.fields.zip(to_struct.fields).foreach { case (a, b) =>
+ isSupported(a.dataType, b.dataType, timeZoneId, evalMode) match {
+ case Compatible(_) =>
+ // all good
+ case other =>
+ return other
+ }
+ }
+ Compatible()
case _ => Unsupported
}
}
diff --git a/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
b/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
index 817545c5..db9a870d 100644
--- a/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
@@ -881,6 +881,20 @@ class CometCastSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("cast StructType to StructType") {
+ Seq(true, false).foreach { dictionaryEnabled =>
+ withTempDir { dir =>
+ val path = new Path(dir.toURI.toString, "test.parquet")
+ makeParquetFileAllTypes(path, dictionaryEnabled = dictionaryEnabled,
10000)
+ withParquetTable(path.toString, "tbl") {
+ checkSparkAnswerAndOperator(
+ "SELECT CAST(CASE WHEN _1 THEN struct(_1, _2, _3, _4) ELSE null
END as " +
+ "struct<_1:string, _2:string, _3:string, _4:string>) FROM tbl")
+ }
+ }
+ }
+ }
+
private def generateFloats(): DataFrame = {
withNulls(gen.generateFloats(dataSize)).toDF("a")
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
