From 8b17442c08d78540e19d94ed9fb7dc97dec513e1 Mon Sep 17 00:00:00 2001
From: Pramod <pramod@ahana.io>
Date: Tue, 16 Aug 2022 12:04:42 -0700
Subject: [PATCH] Add json_array_length Presto Json function (#2294)

Summary:
Adds Json function json_array_length to find the length of Json arrays, if the Json object doesn't represent an array, return NULL.

Pull Request resolved: https://github.com/facebookincubator/velox/pull/2294

Reviewed By: Yuhta

Differential Revision: D38721693

Pulled By: kagamiori

fbshipit-source-id: 39df39eea3f91e9903ffddcb0b6b0e3f69754b37
---
 velox/docs/functions/json.rst                 |  7 ++++++
 velox/functions/prestosql/JsonFunctions.h     | 17 ++++++++++++++
 .../JsonFunctionsRegistration.cpp             |  2 ++
 .../prestosql/tests/JsonFunctionsTest.cpp     | 22 +++++++++++++++++++
 4 files changed, 48 insertions(+)

diff --git a/velox/docs/functions/json.rst b/velox/docs/functions/json.rst
index 70ab868d65a2..91f84ffe1f3c 100644
--- a/velox/docs/functions/json.rst
+++ b/velox/docs/functions/json.rst
@@ -109,6 +109,13 @@ JSON Functions
         SELECT is_json_scalar('1'); *-- true*
         SELECT is_json_scalar('[1, 2, 3]'); *-- false*
 
+.. function:: json_array_length(json) -> bigint
+
+    Returns the array length of ``json`` (a string containing a JSON
+    array). Returns NULL if ``json`` is not an array::
+
+        SELECT json_array_length('[1, 2, 3]');
+
 ============
 JSON Vectors
 ============
diff --git a/velox/functions/prestosql/JsonFunctions.h b/velox/functions/prestosql/JsonFunctions.h
index be8f92558c0e..d6e3da1ac29d 100644
--- a/velox/functions/prestosql/JsonFunctions.h
+++ b/velox/functions/prestosql/JsonFunctions.h
@@ -57,4 +57,21 @@ struct JsonExtractScalarFunction {
   }
 };
 
+template <typename T>
+struct JsonArrayLengthFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE bool call(
+      int64_t& result,
+      const arg_type<Varchar>& json) {
+    auto parsedJson = folly::parseJson(json);
+    if (!parsedJson.isArray()) {
+      return false;
+    }
+
+    result = parsedJson.size();
+    return true;
+  }
+};
+
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp
index 9b0ba50a115b..fa92c22efb39 100644
--- a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp
+++ b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp
@@ -22,6 +22,8 @@ void registerJsonFunctions() {
   registerFunction<IsJsonScalarFunction, bool, Varchar>({"is_json_scalar"});
   registerFunction<JsonExtractScalarFunction, Varchar, Varchar, Varchar>(
       {"json_extract_scalar"});
+  registerFunction<JsonArrayLengthFunction, int64_t, Varchar>(
+      {"json_array_length"});
 }
 
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/tests/JsonFunctionsTest.cpp b/velox/functions/prestosql/tests/JsonFunctionsTest.cpp
index 86cf10424a89..60890c664a3c 100644
--- a/velox/functions/prestosql/tests/JsonFunctionsTest.cpp
+++ b/velox/functions/prestosql/tests/JsonFunctionsTest.cpp
@@ -26,6 +26,10 @@ class JsonFunctionsTest : public functions::test::FunctionBaseTest {
   std::optional<bool> is_json_scalar(std::optional<std::string> json) {
     return evaluateOnce<bool>("is_json_scalar(c0)", json);
   }
+
+  std::optional<int64_t> json_array_length(std::optional<std::string> json) {
+    return evaluateOnce<int64_t>("json_array_length(c0)", json);
+  }
 };
 
 TEST_F(JsonFunctionsTest, isJsonScalar) {
@@ -45,6 +49,24 @@ TEST_F(JsonFunctionsTest, isJsonScalar) {
   EXPECT_EQ(is_json_scalar(R"({"k1":""})"), false);
 }
 
+TEST_F(JsonFunctionsTest, jsonArrayLength) {
+  EXPECT_EQ(json_array_length(R"([])"), 0);
+  EXPECT_EQ(json_array_length(R"([1])"), 1);
+  EXPECT_EQ(json_array_length(R"([1, 2, 3])"), 3);
+  EXPECT_EQ(
+      json_array_length(
+          R"([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20])"),
+      20);
+
+  EXPECT_EQ(json_array_length(R"(1)"), std::nullopt);
+  EXPECT_EQ(json_array_length(R"("hello")"), std::nullopt);
+  EXPECT_EQ(json_array_length(R"("")"), std::nullopt);
+  EXPECT_EQ(json_array_length(R"(true)"), std::nullopt);
+  EXPECT_EQ(json_array_length(R"({"k1":"v1"})"), std::nullopt);
+  EXPECT_EQ(json_array_length(R"({"k1":[0,1,2]})"), std::nullopt);
+  EXPECT_EQ(json_array_length(R"({"k1":[0,1,2], "k2":"v1"})"), std::nullopt);
+}
+
 } // namespace
 
 } // namespace facebook::velox::functions::prestosql