From 7d5a5bd2cc1e3ed2bc3a90f67520bc5509361fc9 Mon Sep 17 00:00:00 2001
From: Pramod <pramod@ahana.io>
Date: Mon, 15 Aug 2022 10:01:08 -0500
Subject: [PATCH] Add Json function - json_array_length

---
 velox/docs/functions/json.rst                 |  7 +++
 .../{JsonExtractScalar.h => JsonFunctions.h}  | 16 ++++++
 .../JsonFunctionsRegistration.cpp             |  4 +-
 .../functions/prestosql/tests/CMakeLists.txt  |  1 +
 .../prestosql/tests/JsonFunctionsTest.cpp     | 50 +++++++++++++++++++
 velox/functions/sparksql/Register.cpp         |  2 +-
 6 files changed, 78 insertions(+), 2 deletions(-)
 rename velox/functions/prestosql/{JsonExtractScalar.h => JsonFunctions.h} (83%)
 create mode 100644 velox/functions/prestosql/tests/JsonFunctionsTest.cpp

diff --git a/velox/docs/functions/json.rst b/velox/docs/functions/json.rst
index 83131527191ab..1f57fd26504a9 100644
--- a/velox/docs/functions/json.rst
+++ b/velox/docs/functions/json.rst
@@ -101,6 +101,13 @@ JSON Functions
 
     .. _JSONPath: http://goessner.net/articles/JsonPath/
 
+.. 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/JsonExtractScalar.h b/velox/functions/prestosql/JsonFunctions.h
similarity index 83%
rename from velox/functions/prestosql/JsonExtractScalar.h
rename to velox/functions/prestosql/JsonFunctions.h
index 98c4d5108eae0..9b64e3054c1fc 100644
--- a/velox/functions/prestosql/JsonExtractScalar.h
+++ b/velox/functions/prestosql/JsonFunctions.h
@@ -46,4 +46,20 @@ struct JsonExtractScalarFunction {
   }
 };
 
+template <typename T>
+struct jsonArrayLengthFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      int64_t& result,
+      const arg_type<Varchar>& json) {
+    auto parsedJson = folly::parseJson(json);
+    if(!parsedJson.isArray()){
+      result = NULL;
+    } else {
+      result = parsedJson.size();
+    }
+  }
+};
+
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp
index 5f43904d05cb7..e1803628102ad 100644
--- a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp
+++ b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp
@@ -15,12 +15,14 @@
  */
 
 #include "velox/functions/Registerer.h"
-#include "velox/functions/prestosql/JsonExtractScalar.h"
+#include "velox/functions/prestosql/JsonFunctions.h"
 
 namespace facebook::velox::functions {
 void registerJsonFunctions() {
   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/CMakeLists.txt b/velox/functions/prestosql/tests/CMakeLists.txt
index 440affb3f0d1c..1139d08ca0d5e 100644
--- a/velox/functions/prestosql/tests/CMakeLists.txt
+++ b/velox/functions/prestosql/tests/CMakeLists.txt
@@ -45,6 +45,7 @@ add_executable(
   InPredicateTest.cpp
   JsonCastTest.cpp
   JsonExtractScalarTest.cpp
+  JsonFunctionsTest.cpp
   MapConcatTest.cpp
   MapEntriesTest.cpp
   MapFilterTest.cpp
diff --git a/velox/functions/prestosql/tests/JsonFunctionsTest.cpp b/velox/functions/prestosql/tests/JsonFunctionsTest.cpp
new file mode 100644
index 0000000000000..6f3cd5f020a55
--- /dev/null
+++ b/velox/functions/prestosql/tests/JsonFunctionsTest.cpp
@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed 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.
+ */
+
+#include "velox/functions/prestosql/tests/FunctionBaseTest.h"
+#include "velox/functions/prestosql/types/JsonType.h"
+
+namespace facebook::velox::functions::prestosql {
+
+namespace {
+
+class JsonFunctionsTest : public functions::test::FunctionBaseTest {
+ public:
+  std::optional<int64_t> json_array_length(std::optional<std::string> json) {
+    return evaluateOnce<int64_t>("json_array_length(c0)", json);
+  }
+};
+
+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)"), 0);
+  EXPECT_EQ(json_array_length(R"("hello")"), 0);
+  EXPECT_EQ(json_array_length(R"("")"), 0);
+  EXPECT_EQ(json_array_length(R"(true)"), 0);
+  EXPECT_EQ(json_array_length(R"({"k1":"v1"})"), 0);
+  EXPECT_EQ(json_array_length(R"({"k1":[0,1,2]})"), 0);
+}
+
+} // namespace
+
+} // namespace facebook::velox::functions::prestosql
diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp
index a49465c327f88..52f8987398091 100644
--- a/velox/functions/sparksql/Register.cpp
+++ b/velox/functions/sparksql/Register.cpp
@@ -18,7 +18,7 @@
 #include "velox/functions/lib/Re2Functions.h"
 #include "velox/functions/lib/RegistrationHelpers.h"
 #include "velox/functions/prestosql/DateTimeFunctions.h"
-#include "velox/functions/prestosql/JsonExtractScalar.h"
+#include "velox/functions/prestosql/JsonFunctions.h"
 #include "velox/functions/prestosql/Rand.h"
 #include "velox/functions/prestosql/StringFunctions.h"
 #include "velox/functions/sparksql/ArraySort.h"