Merge branch 'main' of github.com:Groennbeck/datafusion-comet into ar…

…ray-size

native/core/src/execution/planner.rs

@@ -84,9 +84,11 @@ use datafusion_comet_proto::{
     spark_partitioning::{partitioning::PartitioningStruct, Partitioning as SparkPartitioning},
 };
 use datafusion_comet_spark_expr::{
-    ArrayInsert, ArraySize, Cast, CreateNamedStruct, DateTruncExpr, GetArrayStructFields,
-    GetStructField, HourExpr, IfExpr, ListExtract, MinuteExpr, RLike, SecondExpr,
-    TimestampTruncExpr, ToJson,
+    ArrayInsert, ArraySize, Avg, AvgDecimal, BitwiseNotExpr, Cast, CheckOverflow, Contains, Correlation,
+    Covariance, CreateNamedStruct, DateTruncExpr, EndsWith, GetArrayStructFields, GetStructField,
+    HourExpr, IfExpr, Like, ListExtract, MinuteExpr, NormalizeNaNAndZero, RLike, SecondExpr,
+    SparkCastOptions, StartsWith, Stddev, StringSpaceExpr, SubstringExpr, SumDecimal,
+    TimestampTruncExpr, ToJson, UnboundColumn, Variance,
 };
 use datafusion_common::scalar::ScalarStructBuilder;
 use datafusion_common::{
@@ -719,6 +721,20 @@ impl PhysicalPlanner {
                     expr.legacy_negative_index,
                 )))
             }
+            ExprStruct::ArrayContains(expr) => {
+                let src_array_expr =
+                    self.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                let key_expr =
+                    self.create_expr(expr.right.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                let args = vec![Arc::clone(&src_array_expr), key_expr];
+                let array_has_expr = Arc::new(ScalarFunctionExpr::new(
+                    "array_has",
+                    Arc::new(ScalarUDF::new_from_impl(ArrayHas::new())),
+                    args,
+                    DataType::Boolean,
+                ));
+                Ok(array_has_expr)
+            }
             ExprStruct::ArraySize(expr) => {
                 let child = self.create_expr(
                     expr.src_array_expr.as_ref().unwrap(),

native/proto/src/proto/expr.proto

@@ -84,7 +84,8 @@ message Expr {
     GetArrayStructFields get_array_struct_fields = 57;
     BinaryExpr array_append = 58;
     ArrayInsert array_insert = 59;
-    ArraySize array_size = 60;
+    BinaryExpr array_contains = 60;
+    ArraySize array_size = 61;
   }
 }
 

native/spark-expr/src/list.rs

@@ -685,21 +685,14 @@ impl Display for ArrayInsert {
     }
 }
 
-impl PartialEq<dyn Any> for ArrayInsert {
-    fn eq(&self, other: &dyn Any) -> bool {
-        down_cast_any_ref(other)
-            .downcast_ref::<Self>()
-            .map(|x| {
-                self.src_array_expr.eq(&x.src_array_expr)
-                    && self.pos_expr.eq(&x.pos_expr)
-                    && self.item_expr.eq(&x.item_expr)
-                    && self.legacy_negative_index.eq(&x.legacy_negative_index)
-            })
-            .unwrap_or(false)
+impl PartialEq for ArraySize {
+    fn eq(&self, other: &Self) -> bool {
+        self.src_array_expr.eq(&other.src_array_expr)
     }
 }
 
-#[derive(Debug, Hash)]
+
+#[derive(Debug, Eq)]
 pub struct ArraySize {
     src_array_expr: Arc<dyn PhysicalExpr>,
 }
@@ -716,12 +709,9 @@ impl Display for ArraySize {
     }
 }
 
-impl PartialEq<dyn Any> for ArraySize {
-    fn eq(&self, other: &dyn Any) -> bool {
-        down_cast_any_ref(other)
-            .downcast_ref::<Self>()
-            .map(|x| self.src_array_expr.eq(&x.src_array_expr))
-            .unwrap_or(false)
+impl Hash for ArraySize {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.src_array_expr.hash(state);
     }
 }
 
@@ -777,12 +767,6 @@ impl PhysicalExpr for ArraySize {
             _ => internal_err!("ArraySize should have exactly one child"),
         }
     }
-
-    fn dyn_hash(&self, state: &mut dyn Hasher) {
-        let mut s = state;
-        self.src_array_expr.hash(&mut s);
-        self.hash(&mut s);
-    }
 }
 
 #[cfg(test)]

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -2534,4 +2534,16 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         checkSparkAnswerAndOperator(df.select("arrSizeResultNull"))
       })
   }
+
+  test("array_contains") {
+    withTempDir { dir =>
+      val path = new Path(dir.toURI.toString, "test.parquet")
+      makeParquetFileAllTypes(path, dictionaryEnabled = false, n = 10000)
+      spark.read.parquet(path.toString).createOrReplaceTempView("t1");
+      checkSparkAnswerAndOperator(
+        spark.sql("SELECT array_contains(array(_2, _3, _4), _2) FROM t1"))
+      checkSparkAnswerAndOperator(
+        spark.sql("SELECT array_contains((CASE WHEN _2 =_3 THEN array(_4) END), _4) FROM t1"));
+    }
+  }
 }