feat: Array element extraction (#899)

apache · Sep 5, 2024 · 0f80df2 · 0f80df2
1 parent 60cad71
commit 0f80df2
Show file tree

Hide file tree

Showing 8 changed files with 486 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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};