Feat: Add support for array_size

native/core/src/execution/planner.rs

@@ -100,6 +100,7 @@ use datafusion_expr::{
     WindowFunctionDefinition,
 };
 use datafusion_functions_nested::array_has::ArrayHas;
+use datafusion_functions_nested::length::ArrayLength;
 use datafusion_physical_expr::expressions::{Literal, StatsType};
 use datafusion_physical_expr::window::WindowExpr;
 use datafusion_physical_expr::LexOrdering;
@@ -735,6 +736,24 @@ impl PhysicalPlanner {
                 ));
                 Ok(array_has_expr)
             }
+            ExprStruct::ArraySize(expr) => {
+                let src_array_expr =
+                    self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                let array_size_expr = Arc::new(ScalarFunctionExpr::new(
+                    "array_size",
+                    Arc::new(ScalarUDF::new_from_impl(ArrayLength::new())),
+                    vec![src_array_expr],
+                    DataType::UInt64,
+                ));
+                // Note: Return type of array_size if UInt32, casting to Int32 for Spark compatibility
+                let mut cast_opts = SparkCastOptions::new_without_timezone(EvalMode::Legacy, false);
+                cast_opts.allow_cast_unsigned_ints = true;
+                Ok(Arc::new(Cast::new(
+                    array_size_expr,
+                    DataType::Int32,
+                    cast_opts,
+                )))
+            }
             expr => Err(ExecutionError::GeneralError(format!(
                 "Not implemented: {:?}",
                 expr

native/proto/src/proto/expr.proto

@@ -85,6 +85,7 @@ message Expr {
     BinaryExpr array_append = 58;
     ArrayInsert array_insert = 59;
     BinaryExpr array_contains = 60;
+    UnaryExpr array_size = 61;
   }
 }
 

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -2278,6 +2278,14 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
             expr.children(1),
             inputs,
             (builder, binaryExpr) => builder.setArrayAppend(binaryExpr))
+        case _ if expr.prettyName == "array_size" || expr.prettyName == "size" =>
+          if (!expr.children.head.dataType.isInstanceOf[ArrayType]) {
+            return None
+          }
+          createUnaryExpr(
+            expr.children.head,
+            inputs,
+            (builder, unaryExpr) => builder.setArraySize(unaryExpr))
         case _ =>
           withInfo(expr, s"${expr.prettyName} is not supported", expr.children: _*)
           None

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

@@ -2529,4 +2529,15 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         spark.sql("SELECT array_contains((CASE WHEN _2 =_3 THEN array(_4) END), _4) FROM t1"));
     }
   }
+
+  test("array_size") {
+    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_size(array(_2, _3, _4)) FROM t1"))
+      checkSparkAnswerAndOperator(
+        spark.sql("SELECT array_size((CASE WHEN _2 =_3 THEN array(_4) END)) FROM t1"));
+    }
+  }
 }

spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala

@@ -164,6 +164,7 @@ abstract class CometTestBase
       df: => DataFrame,
       includeClasses: Seq[Class[_]],
       excludedClasses: Class[_]*): Unit = {
+    println(df.queryExecution.executedPlan)
     checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan), excludedClasses: _*)
     checkPlanContains(stripAQEPlan(df.queryExecution.executedPlan), includeClasses: _*)
     checkSparkAnswer(df)