diff --git a/velox/functions/sparksql/JsonArrayLength.h b/velox/functions/sparksql/JsonArrayLength.h new file mode 100644 index 000000000000..813b801d73f8 --- /dev/null +++ b/velox/functions/sparksql/JsonArrayLength.h @@ -0,0 +1,53 @@ +/* + * 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. + */ + +#pragma once + +#include "velox/functions/Macros.h" +#include "velox/functions/prestosql/types/JsonType.h" + +namespace facebook::velox::functions::sparksql { + +template +struct JsonArrayLengthFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE bool call(int32_t& len, const arg_type& json) { + simdjson::ondemand::document jsonDoc; + + simdjson::padded_string paddedJson(json.data(), json.size()); + if (simdjsonParse(paddedJson).get(jsonDoc)) { + return false; + } + if (jsonDoc.type().error()) { + return false; + } + + if (jsonDoc.type() != simdjson::ondemand::json_type::array) { + return false; + } + + size_t numElements; + if (jsonDoc.count_elements().get(numElements)) { + return false; + } + + len = numElements; + return true; + } +}; +} // namespace facebook::velox::functions::sparksql + diff --git a/velox/functions/sparksql/registration/RegisterJson.cpp b/velox/functions/sparksql/registration/RegisterJson.cpp index 7f41807ceb94..142cbdd15860 100644 --- a/velox/functions/sparksql/registration/RegisterJson.cpp +++ b/velox/functions/sparksql/registration/RegisterJson.cpp @@ -15,6 +15,7 @@ */ #include "velox/functions/lib/RegistrationHelpers.h" #include "velox/functions/sparksql/GetJsonObject.h" +#include "velox/functions/sparksql/JsonArrayLength.h" #include "velox/functions/sparksql/JsonObjectKeys.h" namespace facebook::velox::functions::sparksql { @@ -24,6 +25,8 @@ void registerJsonFunctions(const std::string& prefix) { {prefix + "get_json_object"}); registerFunction, Varchar>( {prefix + "json_object_keys"}); + registerFunction( + {prefix + "json_array_length"}); } } // namespace facebook::velox::functions::sparksql