feat: support Spark expression json_array_length

native/spark-expr/src/comet_scalar_funcs.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::hash_funcs::*;
+use crate::json_funcs::JsonArrayLength;
 use crate::map_funcs::spark_map_sort;
 use crate::math_funcs::abs::abs;
 use crate::math_funcs::checked_arithmetic::{checked_add, checked_div, checked_mul, checked_sub};
@@ -216,6 +217,7 @@ fn all_scalar_functions() -> Vec<Arc<ScalarUDF>> {
         Arc::new(ScalarUDF::new_from_impl(SparkMakeDate::default())),
         Arc::new(ScalarUDF::new_from_impl(SparkSecondsToTimestamp::default())),
         Arc::new(ScalarUDF::new_from_impl(SparkSizeFunc::default())),
+        Arc::new(ScalarUDF::new_from_impl(JsonArrayLength::default())),
     ]
 }
 

native/spark-expr/src/json_funcs/json_array_length.rs

@@ -0,0 +1,137 @@
+// 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::{Array, ArrayRef, Int32Builder};
+use arrow::datatypes::DataType;
+use datafusion::common::cast::as_string_array;
+use datafusion::common::{exec_err, Result, ScalarValue};
+use datafusion::logical_expr::{
+    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+
+use std::any::Any;
+use std::sync::Arc;
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct JsonArrayLength {
+    signature: Signature,
+}
+
+impl Default for JsonArrayLength {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl JsonArrayLength {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::variadic(vec![DataType::Utf8], Volatility::Immutable),
+        }
+    }
+}
+
+impl ScalarUDFImpl for JsonArrayLength {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "json_array_length"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Int32)
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        spark_json_array_length(&args.args)
+    }
+}
+
+fn spark_json_array_length(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+    if args.len() != 1 {
+        return exec_err!("json_array_length function takes exactly one argument");
+    }
+    match &args[0] {
+        ColumnarValue::Array(array) => {
+            let result = spark_json_array_length_array(array)?;
+            Ok(ColumnarValue::Array(result))
+        }
+        ColumnarValue::Scalar(scalar) => {
+            let result = spark_json_array_length_scalar(scalar)?;
+            Ok(ColumnarValue::Scalar(result))
+        }
+    }
+}
+
+fn spark_json_array_length_array(array: &ArrayRef) -> Result<ArrayRef> {
+    match array.data_type() {
+        DataType::Utf8 => {
+            let array = as_string_array(array)?;
+            let mut builder = Int32Builder::with_capacity(array.len());
+
+            for row_idx in 0..array.len() {
+                if array.is_null(row_idx) {
+                    builder.append_null();
+                } else {
+                    let json_str = array.value(row_idx);
+                    if let Some(json_array_length) = get_json_array_length(json_str) {
+                        builder.append_value(json_array_length);
+                    } else {
+                        builder.append_null()
+                    }
+                }
+            }
+            Ok(Arc::new(builder.finish()))
+        }
+        other => {
+            exec_err!("Unsupported data type {other:?} for function `json_array_length`")
+        }
+    }
+}
+
+fn spark_json_array_length_scalar(scalar: &ScalarValue) -> Result<ScalarValue> {
+    match scalar {
+        ScalarValue::Utf8(value) => {
+            let length = value
+                .clone()
+                .and_then(|json_str| get_json_array_length(&json_str));
+            Ok(ScalarValue::Int32(length))
+        }
+        other => {
+            exec_err!("Unsupported data type {other:?} for function `json_array_length`")
+        }
+    }
+}
+
+fn get_json_array_length(json_str: &str) -> Option<i32> {
+    match serde_json::from_str::<serde_json::Value>(json_str) {
+        Ok(json_value) => {
+            if json_value.is_array() {
+                Some(json_value.as_array().unwrap().len() as i32)
+            } else {
+                None
+            }
+        }
+        Err(_) => None,
+    }
+}

native/spark-expr/src/json_funcs/mod.rs

@@ -16,7 +16,9 @@
 // under the License.
 
 mod from_json;
+mod json_array_length;
 mod to_json;
 
 pub use from_json::FromJson;
+pub use json_array_length::JsonArrayLength;
 pub use to_json::ToJson;

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

@@ -249,6 +249,9 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
   private val conversionExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
     classOf[Cast] -> CometCast)
 
+  private val jsonExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
+    classOf[LengthOfJsonArray] -> CometScalarFunction("json_array_length"))
+
   private[comet] val miscExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
     // TODO PromotePrecision
     classOf[Alias] -> CometAlias,
@@ -273,7 +276,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     mathExpressions ++ hashExpressions ++ stringExpressions ++
       conditionalExpressions ++ mapExpressions ++ predicateExpressions ++
       structExpressions ++ bitwiseExpressions ++ miscExpressions ++ arrayExpressions ++
-      temporalExpressions ++ conversionExpressions
+      temporalExpressions ++ conversionExpressions ++ jsonExpressions
 
   /**
    * Mapping of Spark aggregate expression class to Comet expression handler.

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

@@ -20,6 +20,7 @@
 package org.apache.comet.serde
 
 import org.apache.spark.sql.catalyst.expressions.{Attribute, ExpressionImplUtils, Literal, UrlCodec}
+import org.apache.spark.sql.catalyst.expressions.json.JsonExpressionUtils
 import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke
 import org.apache.spark.sql.catalyst.util.CharVarcharCodegenUtils
 
@@ -38,7 +39,9 @@ object CometStaticInvoke extends CometExpressionSerde[StaticInvoke] {
         "read_side_padding"),
       ("isLuhnNumber", classOf[ExpressionImplUtils]) -> CometScalarFunction("luhn_check"),
       ("encode", UrlCodec.getClass) -> CometUrlEncodeStaticInvoke,
-      ("decode", UrlCodec.getClass) -> CometUrlDecodeStaticInvoke)
+      ("decode", UrlCodec.getClass) -> CometUrlDecodeStaticInvoke,
+      ("lengthOfJsonArray", classOf[JsonExpressionUtils]) -> CometScalarFunction(
+        "json_array_length"))
 
   override def convert(
       expr: StaticInvoke,

spark/src/test/resources/sql-tests/expressions/json/json_array_length.sql

@@ -0,0 +1,61 @@
+-- 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.
+
+statement
+CREATE TABLE test_json_array_length(j string) USING parquet
+
+statement
+INSERT INTO test_json_array_length VALUES
+  ('[1,2,3,4]'),
+  ('[]'),
+  ('[1]'),
+  (NULL),
+  ('[1,2,3,{"f1":1,"f2":[5,6]},4]'),
+  ('[[1,2],[3,4],[5,6]]'),
+  ('[{"a":1},{"b":2},{"c":3}]'),
+  ('[1,2'),
+  ('[1,2,3,]'),
+  ('not a json'),
+  ('{"object": "not array"}'),
+  (''),
+  ('   '),
+  ('[true, false, null]'),
+  ('["string1", "string2", "string3"]'),
+  ('[1, "mixed", true, null, {"key":"value"}]'),
+  ('[1,2,3,4,5,6,7,8,9,10]'),
+  ('["line1\nline2", "tab\tseparated", "quote\"here"]'),
+  ('{"outer": [1,2,3], "inner": [[1,2],[3,4]]}'),
+  ('{"arrays": {"first": [1,2], "second": [3,4,5]}}'),
+  ('[{"arr": [1,2,3]}, {"arr": [4,5]}]')
+
+query
+SELECT json_array_length(j) FROM test_json_array_length
+
+query
+SELECT json_array_length('[1,2,3,4]')
+
+query
+SELECT json_array_length('not an array')
+
+query
+SELECT json_array_length('{"key":"value"}')
+
+query
+SELECT json_array_length(NULL)
+
+query
+SELECT json_array_length('[]')