[SPARK-31008][SQL] Support json_array_length function

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -538,6 +538,7 @@ object FunctionRegistry {
     expression[StructsToJson]("to_json"),
     expression[JsonToStructs]("from_json"),
     expression[SchemaOfJson]("schema_of_json"),
+    expression[LengthOfJsonArray]("json_array_length"),
 
     // cast
     expression[Cast]("cast"),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala

@@ -796,3 +796,75 @@ case class SchemaOfJson(
 
   override def prettyName: String = "schema_of_json"
 }
+
+/**
+ * A function that returns the number of elements in outer JSON array.
+ */
+@ExpressionDescription(
+  usage = "_FUNC_(jsonArray) - Returns the number of elements in outer JSON array.",
+  arguments = """
+    Arguments:
+      * jsonArray - A JSON array. An exception is thrown if any other valid JSON strings are passed.
+          `NULL` is returned in case of an invalid JSON.
+  """,
+  examples = """
+    Examples:
+      > SELECT _FUNC_('[1,2,3,4]');
+        4
+      > SELECT _FUNC_('[1,2,3,{"f1":1,"f2":[5,6]},4]');
+        5
+      > SELECT _FUNC_('[1,2');
+        NULL
+  """,
+  since = "3.1.0"
+)
+case class LengthOfJsonArray(child: Expression) extends UnaryExpression
+  with CodegenFallback with ExpectsInputTypes {
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(StringType)
+  override def dataType: DataType = IntegerType
+  override def nullable: Boolean = true
+  override def prettyName: String = "json_array_length"
+
+  override def eval(input: InternalRow): Any = {
+    val json = child.eval(input).asInstanceOf[UTF8String]
+    // return null for null input
+    if (json == null) {
+      return null
+    }
+
+    try {
+      Utils.tryWithResource(CreateJacksonParser.utf8String(SharedFactory.jsonFactory, json)) {
+        parser => {
+          // return null if null array is encountered.
+          if (parser.nextToken() == null) {
+            return null
+          }
+          // Parse the array to compute its length.
+          parseCounter(parser, input)
+        }
+      }
+    } catch {
+      case _: JsonProcessingException | _: IOException => null
+    }
+  }
+
+  private def parseCounter(parser: JsonParser, input: InternalRow): Int = {
+    var length = 0;
+    // Only JSON array are supported for this function.
+    if (parser.currentToken != JsonToken.START_ARRAY) {
+      throw new IllegalArgumentException(s"$prettyName can only be called on JSON array.")
+    }
+    // Keep traversing until the end of JSON array
+    while(parser.nextToken() != JsonToken.END_ARRAY) {
+      // Null indicates end of input.
+      if (parser.currentToken == null) {
+        throw new IllegalArgumentException("Please provide a valid JSON array.")
+      }
+      length += 1
+      // skip all the child of inner object or array
+      parser.skipChildren()
+    }
+    length
+  }
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala

@@ -790,4 +790,27 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper with
         checkDecimalInfer(_, """struct<d:decimal(7,3)>""")
     }
   }
+
+  test("Length of JSON array") {
+    Seq(
+      ("""""", null),
+      ("""[1,2,3]""", 3),
+      ("""[]""", 0),
+      ("""[[1],[2,3],[]]""", 3),
+      ("""[{"a":123},{"b":"hello"}]""", 2),
+      ("""[1,2,3,[33,44],{"key":[2,3,4]}]""", 5),
+      ("""[1,2,3,4,5""", null),
+      ("""Random String""", null)
+    ).foreach{
+      case(literal, expectedValue) =>
+        checkEvaluation(LengthOfJsonArray(Literal(literal)), expectedValue)
+    }
+
+    val not_a_json_array = """{"key":"not a json array"}"""
+
+    checkExceptionInExpression[IllegalArgumentException](
+      LengthOfJsonArray(Literal(not_a_json_array)),
+      expectedErrMsg = "json_array_length can only be called on JSON array"
+    )
+  }
 }

sql/core/src/test/resources/sql-tests/inputs/json-functions.sql

@@ -58,5 +58,19 @@ select schema_of_json('{"c1":01, "c2":0.1}', map('allowNumericLeadingZeros', 'tr
 select schema_of_json(null);
 CREATE TEMPORARY VIEW jsonTable(jsonField, a) AS SELECT * FROM VALUES ('{"a": 1, "b": 2}', 'a');
 SELECT schema_of_json(jsonField) FROM jsonTable;
+
+-- json_array_length
+select json_array_length(null);
+select json_array_length(2);
+select json_array_length();
+select json_array_length('');
+select json_array_length('[]');
+select json_array_length('[1,2,3]');
+select json_array_length('[[1,2],[5,6,7]]');
+select json_array_length('[{"a":123},{"b":"hello"}]');
+select json_array_length('[1,2,3,[33,44],{"key":[2,3,4]}]');
+select json_array_length('{"key":"not a json array"}');
+select json_array_length('[1,2,3,4,5');
+
 -- Clean up
 DROP VIEW IF EXISTS jsonTable;

sql/core/src/test/resources/sql-tests/results/json-functions.sql.out

@@ -1,5 +1,5 @@
 -- Automatically generated by SQLQueryTestSuite
--- Number of queries: 42
+-- Number of queries: 53
 
 
 -- !query
@@ -346,6 +346,97 @@ org.apache.spark.sql.AnalysisException
 cannot resolve 'schema_of_json(jsontable.`jsonField`)' due to data type mismatch: The input json should be a foldable string expression and not null; however, got jsontable.`jsonField`.; line 1 pos 7
 
 
+-- !query
+select json_array_length(null)
+-- !query schema
+struct<json_array_length(CAST(NULL AS STRING)):int>
+-- !query output
+NULL
+
+
+-- !query
+select json_array_length(2)
+-- !query schema
+struct<>
+-- !query output
+org.apache.spark.sql.AnalysisException
+cannot resolve 'json_array_length(2)' due to data type mismatch: argument 1 requires string type, however, '2' is of int type.; line 1 pos 7
+
+
+-- !query
+select json_array_length()
+-- !query schema
+struct<>
+-- !query output
+org.apache.spark.sql.AnalysisException
+Invalid number of arguments for function json_array_length. Expected: 1; Found: 0; line 1 pos 7
+
+
+-- !query
+select json_array_length('')
+-- !query schema
+struct<json_array_length():int>
+-- !query output
+NULL
+
+
+-- !query
+select json_array_length('[]')
+-- !query schema
+struct<json_array_length([]):int>
+-- !query output
+0
+
+
+-- !query
+select json_array_length('[1,2,3]')
+-- !query schema
+struct<json_array_length([1,2,3]):int>
+-- !query output
+3
+
+
+-- !query
+select json_array_length('[[1,2],[5,6,7]]')
+-- !query schema
+struct<json_array_length([[1,2],[5,6,7]]):int>
+-- !query output
+2
+
+
+-- !query
+select json_array_length('[{"a":123},{"b":"hello"}]')
+-- !query schema
+struct<json_array_length([{"a":123},{"b":"hello"}]):int>
+-- !query output
+2
+
+
+-- !query
+select json_array_length('[1,2,3,[33,44],{"key":[2,3,4]}]')
+-- !query schema
+struct<json_array_length([1,2,3,[33,44],{"key":[2,3,4]}]):int>
+-- !query output
+5
+
+
+-- !query
+select json_array_length('{"key":"not a json array"}')
+-- !query schema
+struct<>
+-- !query output
+java.lang.IllegalArgumentException
+json_array_length can only be called on JSON array.
+
+
+-- !query
+select json_array_length('[1,2,3,4,5')
+-- !query schema
+struct<json_array_length([1,2,3,4,5):int>
+-- !query output
+NULL
+
+
 -- !query
 DROP VIEW IF EXISTS jsonTable
 -- !query schema

sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala

@@ -710,4 +710,11 @@ class JsonFunctionsSuite extends QueryTest with SharedSparkSession {
       Seq(Row("string")))
   }
 
+  test("json_array_length") {
+    val df = Seq(1).toDF("json")
+    val errMsg = intercept[AnalysisException] {
+      df.selectExpr("json_array_length(json)")
+    }.getMessage
+    assert(errMsg.contains("due to data type mismatch"))
+  }
 }