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 0f80df2e feat: Array element extraction (#899)
0f80df2e is described below
commit 0f80df2e421969976505626de9e63c742f5a0f4a
Author: Adam Binford <[email protected]>
AuthorDate: Thu Sep 5 07:02:37 2024 -0400
feat: Array element extraction (#899)
---
docs/source/user-guide/expressions.md | 12 +-
docs/spark_expressions_support.md | 5 +-
native/core/src/execution/datafusion/planner.rs | 22 +-
native/proto/src/proto/expr.proto | 9 +
native/spark-expr/src/lib.rs | 2 +
native/spark-expr/src/list.rs | 325 +++++++++++++++++++++
.../org/apache/comet/serde/QueryPlanSerde.scala | 49 ++++
.../org/apache/comet/CometExpressionSuite.scala | 71 +++++
8 files changed, 486 insertions(+), 9 deletions(-)
diff --git a/docs/source/user-guide/expressions.md
b/docs/source/user-guide/expressions.md
index c2b37269..3a179054 100644
--- a/docs/source/user-guide/expressions.md
+++ b/docs/source/user-guide/expressions.md
@@ -184,11 +184,13 @@ The following Spark expressions are currently available.
Any known compatibility
## Complex Types
-| Expression | Notes |
-| ----------------- | ----- |
-| CreateNamedStruct | |
-| GetElementAt | |
-| StructsToJson | |
+| Expression | Notes |
+| ----------------- | ----------- |
+| CreateNamedStruct | |
+| ElementAt | Arrays only |
+| GetArrayItem | |
+| GetStructField | |
+| StructsToJson | |
## Other
diff --git a/docs/spark_expressions_support.md
b/docs/spark_expressions_support.md
index 8fb97586..297b15d2 100644
--- a/docs/spark_expressions_support.md
+++ b/docs/spark_expressions_support.md
@@ -80,7 +80,7 @@
- [x] variance
### array_funcs
- - [ ] array
+ - [x] array
- [ ] array_append
- [ ] array_compact
- [ ] array_contains
@@ -97,8 +97,9 @@
- [ ] array_union
- [ ] arrays_overlap
- [ ] arrays_zip
+ - [x] element_at
- [ ] flatten
- - [ ] get
+ - [x] get
- [ ] sequence
- [ ] shuffle
- [ ] slice
diff --git a/native/core/src/execution/datafusion/planner.rs
b/native/core/src/execution/datafusion/planner.rs
index b9b88282..6fca4678 100644
--- a/native/core/src/execution/datafusion/planner.rs
+++ b/native/core/src/execution/datafusion/planner.rs
@@ -96,8 +96,8 @@ use datafusion_comet_proto::{
spark_partitioning::{partitioning::PartitioningStruct, Partitioning as
SparkPartitioning},
};
use datafusion_comet_spark_expr::{
- Cast, CreateNamedStruct, DateTruncExpr, GetStructField, HourExpr, IfExpr,
MinuteExpr, RLike,
- SecondExpr, TimestampTruncExpr, ToJson,
+ Cast, CreateNamedStruct, DateTruncExpr, GetStructField, HourExpr, IfExpr,
ListExtract,
+ MinuteExpr, RLike, SecondExpr, TimestampTruncExpr, ToJson,
};
use datafusion_common::scalar::ScalarStructBuilder;
use datafusion_common::{
@@ -660,6 +660,24 @@ impl PhysicalPlanner {
let child = self.create_expr(expr.child.as_ref().unwrap(),
input_schema)?;
Ok(Arc::new(ToJson::new(child, &expr.timezone)))
}
+ ExprStruct::ListExtract(expr) => {
+ let child =
+ self.create_expr(expr.child.as_ref().unwrap(),
Arc::clone(&input_schema))?;
+ let ordinal =
+ self.create_expr(expr.ordinal.as_ref().unwrap(),
Arc::clone(&input_schema))?;
+ let default_value = expr
+ .default_value
+ .as_ref()
+ .map(|e| self.create_expr(e, Arc::clone(&input_schema)))
+ .transpose()?;
+ Ok(Arc::new(ListExtract::new(
+ child,
+ ordinal,
+ default_value,
+ expr.one_based,
+ expr.fail_on_error,
+ )))
+ }
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 5fd96ac0..88940f38 100644
--- a/native/proto/src/proto/expr.proto
+++ b/native/proto/src/proto/expr.proto
@@ -80,6 +80,7 @@ message Expr {
CreateNamedStruct create_named_struct = 53;
GetStructField get_struct_field = 54;
ToJson to_json = 55;
+ ListExtract list_extract = 56;
}
}
@@ -508,6 +509,14 @@ message GetStructField {
int32 ordinal = 2;
}
+message ListExtract {
+ Expr child = 1;
+ Expr ordinal = 2;
+ Expr default_value = 3;
+ bool one_based = 4;
+ bool fail_on_error = 5;
+}
+
enum SortDirection {
Ascending = 0;
Descending = 1;
diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs
index 6233a29e..c4b1c99b 100644
--- a/native/spark-expr/src/lib.rs
+++ b/native/spark-expr/src/lib.rs
@@ -24,6 +24,7 @@ mod error;
mod if_expr;
mod kernels;
+mod list;
mod regexp;
pub mod scalar_funcs;
pub mod spark_hash;
@@ -37,6 +38,7 @@ mod xxhash64;
pub use cast::{spark_cast, Cast};
pub use error::{SparkError, SparkResult};
pub use if_expr::IfExpr;
+pub use list::ListExtract;
pub use regexp::RLike;
pub use structs::{CreateNamedStruct, GetStructField};
pub use temporal::{DateTruncExpr, HourExpr, MinuteExpr, SecondExpr,
TimestampTruncExpr};
diff --git a/native/spark-expr/src/list.rs b/native/spark-expr/src/list.rs
new file mode 100644
index 00000000..ec953b66
--- /dev/null
+++ b/native/spark-expr/src/list.rs
@@ -0,0 +1,325 @@
+// 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.
+
+use arrow::{array::MutableArrayData, datatypes::ArrowNativeType,
record_batch::RecordBatch};
+use arrow_array::{Array, GenericListArray, Int32Array, OffsetSizeTrait};
+use arrow_schema::{DataType, FieldRef, Schema};
+use datafusion::logical_expr::ColumnarValue;
+use datafusion_common::{
+ cast::{as_int32_array, as_large_list_array, as_list_array},
+ internal_err, DataFusionError, Result as DataFusionResult, ScalarValue,
+};
+use datafusion_physical_expr::PhysicalExpr;
+use std::{
+ any::Any,
+ fmt::{Display, Formatter},
+ hash::{Hash, Hasher},
+ sync::Arc,
+};
+
+use crate::utils::down_cast_any_ref;
+
+#[derive(Debug, Hash)]
+pub struct ListExtract {
+ child: Arc<dyn PhysicalExpr>,
+ ordinal: Arc<dyn PhysicalExpr>,
+ default_value: Option<Arc<dyn PhysicalExpr>>,
+ one_based: bool,
+ fail_on_error: bool,
+}
+
+impl ListExtract {
+ pub fn new(
+ child: Arc<dyn PhysicalExpr>,
+ ordinal: Arc<dyn PhysicalExpr>,
+ default_value: Option<Arc<dyn PhysicalExpr>>,
+ one_based: bool,
+ fail_on_error: bool,
+ ) -> Self {
+ Self {
+ child,
+ ordinal,
+ default_value,
+ one_based,
+ fail_on_error,
+ }
+ }
+
+ fn child_field(&self, input_schema: &Schema) -> DataFusionResult<FieldRef>
{
+ match self.child.data_type(input_schema)? {
+ DataType::List(field) | DataType::LargeList(field) => Ok(field),
+ data_type => Err(DataFusionError::Internal(format!(
+ "Unexpected data type in ListExtract: {:?}",
+ data_type
+ ))),
+ }
+ }
+}
+
+impl PhysicalExpr for ListExtract {
+ 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> {
+ // Only non-nullable if fail_on_error is enabled and the element is
non-nullable
+ Ok(!self.fail_on_error ||
self.child_field(input_schema)?.is_nullable())
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult<ColumnarValue>
{
+ let child_value =
self.child.evaluate(batch)?.into_array(batch.num_rows())?;
+ let ordinal_value =
self.ordinal.evaluate(batch)?.into_array(batch.num_rows())?;
+
+ let default_value = self
+ .default_value
+ .as_ref()
+ .map(|d| {
+ d.evaluate(batch).map(|value| match value {
+ ColumnarValue::Scalar(scalar)
+ if
!scalar.data_type().equals_datatype(child_value.data_type()) =>
+ {
+ scalar.cast_to(child_value.data_type())
+ }
+ ColumnarValue::Scalar(scalar) => Ok(scalar),
+ v => Err(DataFusionError::Execution(format!(
+ "Expected scalar default value for ListExtract, got
{:?}",
+ v
+ ))),
+ })
+ })
+ .transpose()?
+ .unwrap_or(self.data_type(&batch.schema())?.try_into())?;
+
+ let adjust_index = if self.one_based {
+ one_based_index
+ } else {
+ zero_based_index
+ };
+
+ match child_value.data_type() {
+ DataType::List(_) => {
+ let list_array = as_list_array(&child_value)?;
+ let index_array = as_int32_array(&ordinal_value)?;
+
+ list_extract(
+ list_array,
+ index_array,
+ &default_value,
+ self.fail_on_error,
+ adjust_index,
+ )
+ }
+ DataType::LargeList(_) => {
+ let list_array = as_large_list_array(&child_value)?;
+ let index_array = as_int32_array(&ordinal_value)?;
+
+ list_extract(
+ list_array,
+ index_array,
+ &default_value,
+ self.fail_on_error,
+ adjust_index,
+ )
+ }
+ data_type => Err(DataFusionError::Internal(format!(
+ "Unexpected child type for ListExtract: {:?}",
+ data_type
+ ))),
+ }
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+ vec![&self.child, &self.ordinal]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn PhysicalExpr>>,
+ ) -> datafusion_common::Result<Arc<dyn PhysicalExpr>> {
+ match children.len() {
+ 2 => Ok(Arc::new(ListExtract::new(
+ Arc::clone(&children[0]),
+ Arc::clone(&children[1]),
+ self.default_value.clone(),
+ self.one_based,
+ self.fail_on_error,
+ ))),
+ _ => internal_err!("ListExtract should have exactly two children"),
+ }
+ }
+
+ fn dyn_hash(&self, state: &mut dyn Hasher) {
+ let mut s = state;
+ self.child.hash(&mut s);
+ self.ordinal.hash(&mut s);
+ self.default_value.hash(&mut s);
+ self.one_based.hash(&mut s);
+ self.fail_on_error.hash(&mut s);
+ self.hash(&mut s);
+ }
+}
+
+fn one_based_index(index: i32, len: usize) -> DataFusionResult<Option<usize>> {
+ if index == 0 {
+ return Err(DataFusionError::Execution(
+ "Invalid index of 0 for one-based ListExtract".to_string(),
+ ));
+ }
+
+ let abs_index = index.abs().as_usize();
+ if abs_index <= len {
+ if index > 0 {
+ Ok(Some(abs_index - 1))
+ } else {
+ Ok(Some(len - abs_index))
+ }
+ } else {
+ Ok(None)
+ }
+}
+
+fn zero_based_index(index: i32, len: usize) -> DataFusionResult<Option<usize>>
{
+ if index < 0 {
+ Ok(None)
+ } else {
+ let positive_index = index.as_usize();
+ if positive_index < len {
+ Ok(Some(positive_index))
+ } else {
+ Ok(None)
+ }
+ }
+}
+
+fn list_extract<O: OffsetSizeTrait>(
+ list_array: &GenericListArray<O>,
+ index_array: &Int32Array,
+ default_value: &ScalarValue,
+ fail_on_error: bool,
+ adjust_index: impl Fn(i32, usize) -> DataFusionResult<Option<usize>>,
+) -> DataFusionResult<ColumnarValue> {
+ let values = list_array.values();
+ let offsets = list_array.offsets();
+
+ let data = values.to_data();
+
+ let default_data = default_value.to_array()?.to_data();
+
+ let mut mutable = MutableArrayData::new(vec![&data, &default_data], true,
index_array.len());
+
+ for (row, (offset_window, index)) in
offsets.windows(2).zip(index_array.values()).enumerate() {
+ let start = offset_window[0].as_usize();
+ let len = offset_window[1].as_usize() - start;
+
+ if let Some(i) = adjust_index(*index, len)? {
+ mutable.extend(0, start + i, start + i + 1);
+ } else if list_array.is_null(row) {
+ mutable.extend_nulls(1);
+ } else if fail_on_error {
+ return Err(DataFusionError::Execution(
+ "Index out of bounds for array".to_string(),
+ ));
+ } else {
+ mutable.extend(1, 0, 1);
+ }
+ }
+
+ let data = mutable.freeze();
+ Ok(ColumnarValue::Array(arrow::array::make_array(data)))
+}
+
+impl Display for ListExtract {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(
+ f,
+ "ListExtract [child: {:?}, ordinal: {:?}, default_value: {:?},
one_based: {:?}, fail_on_error: {:?}]",
+ self.child, self.ordinal, self.default_value, self.one_based,
self.fail_on_error
+ )
+ }
+}
+
+impl PartialEq<dyn Any> for ListExtract {
+ 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)
+ && (self.default_value.is_none() ==
x.default_value.is_none())
+ && self
+ .default_value
+ .as_ref()
+ .zip(x.default_value.as_ref())
+ .map(|(s, x)| s.eq(x))
+ .unwrap_or(true)
+ && self.one_based.eq(&x.one_based)
+ && self.fail_on_error.eq(&x.fail_on_error)
+ })
+ .unwrap_or(false)
+ }
+}
+
+#[cfg(test)]
+mod test {
+ use crate::list::{list_extract, zero_based_index};
+
+ use arrow::datatypes::Int32Type;
+ use arrow_array::{Array, Int32Array, ListArray};
+ use datafusion_common::{Result, ScalarValue};
+ use datafusion_expr::ColumnarValue;
+
+ #[test]
+ fn test_list_extract_default_value() -> Result<()> {
+ let list = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+ Some(vec![Some(1)]),
+ None,
+ Some(vec![]),
+ ]);
+ let indices = Int32Array::from(vec![0, 0, 0]);
+
+ let null_default = ScalarValue::Int32(None);
+
+ let ColumnarValue::Array(result) =
+ list_extract(&list, &indices, &null_default, false,
zero_based_index)?
+ else {
+ unreachable!()
+ };
+
+ assert_eq!(
+ &result.to_data(),
+ &Int32Array::from(vec![Some(1), None, None]).to_data()
+ );
+
+ let zero_default = ScalarValue::Int32(Some(0));
+
+ let ColumnarValue::Array(result) =
+ list_extract(&list, &indices, &zero_default, false,
zero_based_index)?
+ else {
+ unreachable!()
+ };
+
+ assert_eq!(
+ &result.to_data(),
+ &Int32Array::from(vec![Some(1), None, Some(0)]).to_data()
+ );
+ Ok(())
+ }
+}
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 bce8295e..d77fac47 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -2460,6 +2460,55 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
None
}
+ case GetArrayItem(child, ordinal, failOnError) =>
+ val childExpr = exprToProto(child, inputs, binding)
+ val ordinalExpr = exprToProto(ordinal, inputs, binding)
+
+ if (childExpr.isDefined && ordinalExpr.isDefined) {
+ val listExtractBuilder = ExprOuterClass.ListExtract
+ .newBuilder()
+ .setChild(childExpr.get)
+ .setOrdinal(ordinalExpr.get)
+ .setOneBased(false)
+ .setFailOnError(failOnError)
+
+ Some(
+ ExprOuterClass.Expr
+ .newBuilder()
+ .setListExtract(listExtractBuilder)
+ .build())
+ } else {
+ withInfo(expr, "unsupported arguments for GetArrayItem", child,
ordinal)
+ None
+ }
+
+ case ElementAt(child, ordinal, defaultValue, failOnError)
+ if child.dataType.isInstanceOf[ArrayType] =>
+ val childExpr = exprToProto(child, inputs, binding)
+ val ordinalExpr = exprToProto(ordinal, inputs, binding)
+ val defaultExpr = defaultValue.flatMap(exprToProto(_, inputs,
binding))
+
+ if (childExpr.isDefined && ordinalExpr.isDefined &&
+ defaultExpr.isDefined == defaultValue.isDefined) {
+ val arrayExtractBuilder = ExprOuterClass.ListExtract
+ .newBuilder()
+ .setChild(childExpr.get)
+ .setOrdinal(ordinalExpr.get)
+ .setOneBased(true)
+ .setFailOnError(failOnError)
+
+ defaultExpr.foreach(arrayExtractBuilder.setDefaultValue(_))
+
+ Some(
+ ExprOuterClass.Expr
+ .newBuilder()
+ .setListExtract(arrayExtractBuilder)
+ .build())
+ } else {
+ withInfo(expr, "unsupported arguments for ElementAt", child,
ordinal)
+ None
+ }
+
case _ =>
withInfo(expr, s"${expr.prettyName} is not supported",
expr.children: _*)
None
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
index f5d0e0e7..3701be5f 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -27,6 +27,7 @@ import scala.util.Random
import org.apache.hadoop.fs.Path
import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.catalyst.optimizer.SimplifyExtractValueOps
import org.apache.spark.sql.comet.CometProjectExec
import org.apache.spark.sql.execution.{ColumnarToRowExec, InputAdapter,
WholeStageCodegenExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
@@ -2094,4 +2095,74 @@ class CometExpressionSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
}
+
+ test("ListExtract") {
+ def assertBothThrow(df: DataFrame): Unit = {
+ checkSparkMaybeThrows(df) match {
+ case (Some(_), Some(_)) => ()
+ case (spark, comet) =>
+ fail(
+ s"Expected Spark and Comet to throw exception, but got\nSpark:
$spark\nComet: $comet")
+ }
+ }
+
+ Seq(true, false).foreach { dictionaryEnabled =>
+ withTempDir { dir =>
+ val path = new Path(dir.toURI.toString, "test.parquet")
+ makeParquetFileAllTypes(path, dictionaryEnabled = dictionaryEnabled,
100)
+
+ Seq(true, false).foreach { ansiEnabled =>
+ withSQLConf(
+ CometConf.COMET_ANSI_MODE_ENABLED.key -> "true",
+ SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString(),
+ // Prevent the optimizer from collapsing an extract value of a
create array
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
SimplifyExtractValueOps.ruleName) {
+ val df = spark.read.parquet(path.toString)
+
+ val stringArray = df.select(array(col("_8"), col("_8"),
lit(null)).alias("arr"))
+ checkSparkAnswerAndOperator(
+ stringArray
+ .select(col("arr").getItem(0), col("arr").getItem(1),
col("arr").getItem(2)))
+
+ checkSparkAnswerAndOperator(
+ stringArray.select(
+ element_at(col("arr"), -3),
+ element_at(col("arr"), -2),
+ element_at(col("arr"), -1),
+ element_at(col("arr"), 1),
+ element_at(col("arr"), 2),
+ element_at(col("arr"), 3)))
+
+ // 0 is an invalid index for element_at
+ assertBothThrow(stringArray.select(element_at(col("arr"), 0)))
+
+ if (ansiEnabled) {
+ assertBothThrow(stringArray.select(col("arr").getItem(-1)))
+ assertBothThrow(stringArray.select(col("arr").getItem(3)))
+ assertBothThrow(stringArray.select(element_at(col("arr"), -4)))
+ assertBothThrow(stringArray.select(element_at(col("arr"), 4)))
+ } else {
+
checkSparkAnswerAndOperator(stringArray.select(col("arr").getItem(-1)))
+
checkSparkAnswerAndOperator(stringArray.select(col("arr").getItem(3)))
+
checkSparkAnswerAndOperator(stringArray.select(element_at(col("arr"), -4)))
+
checkSparkAnswerAndOperator(stringArray.select(element_at(col("arr"), 4)))
+ }
+
+ val intArray =
+ df.select(when(col("_4").isNotNull, array(col("_4"),
col("_4"))).alias("arr"))
+ checkSparkAnswerAndOperator(
+ intArray
+ .select(col("arr").getItem(0), col("arr").getItem(1)))
+
+ checkSparkAnswerAndOperator(
+ intArray.select(
+ element_at(col("arr"), 1),
+ element_at(col("arr"), 2),
+ element_at(col("arr"), -1),
+ element_at(col("arr"), -2)))
+ }
+ }
+ }
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
