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: #2294 Reviewed By: Yuhta Differential Revision: D38721693 Pulled By: kagamiori fbshipit-source-id: 39df39eea3f91e9903ffddcb0b6b0e3f69754b37
facebookincubator · Aug 16, 2022 · 8b17442 · 8b17442
1 parent 098cdf9
commit 8b17442
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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