Merge pull request #30056 from orloffv/orloffv.JSONExtractKeys

Add JSONExtractKeys function

src/Functions/FunctionsJSON.cpp

@@ -73,6 +73,7 @@ void registerFunctionsJSON(FunctionFactory & factory)
     factory.registerFunction<JSONOverloadResolver<NameJSONExtractRaw, JSONExtractRawImpl>>();
     factory.registerFunction<JSONOverloadResolver<NameJSONExtractArrayRaw, JSONExtractArrayRawImpl>>();
     factory.registerFunction<JSONOverloadResolver<NameJSONExtractKeysAndValuesRaw, JSONExtractKeysAndValuesRawImpl>>();
+    factory.registerFunction<JSONOverloadResolver<NameJSONExtractKeys, JSONExtractKeysImpl>>();
 }
 
 }

src/Functions/FunctionsJSON.h

@@ -433,6 +433,7 @@ struct NameJSONExtractKeysAndValues { static constexpr auto name{"JSONExtractKey
 struct NameJSONExtractRaw { static constexpr auto name{"JSONExtractRaw"}; };
 struct NameJSONExtractArrayRaw { static constexpr auto name{"JSONExtractArrayRaw"}; };
 struct NameJSONExtractKeysAndValuesRaw { static constexpr auto name{"JSONExtractKeysAndValuesRaw"}; };
+struct NameJSONExtractKeys { static constexpr auto name{"JSONExtractKeys"}; };
 
 
 template <typename JSONParser>
@@ -1353,4 +1354,37 @@ public:
     }
 };
 
+template <typename JSONParser>
+class JSONExtractKeysImpl
+{
+public:
+    using Element = typename JSONParser::Element;
+
+    static DataTypePtr getReturnType(const char *, const ColumnsWithTypeAndName &)
+    {
+        return std::make_unique<DataTypeArray>(std::make_shared<DataTypeString>());
+    }
+
+    static size_t getNumberOfIndexArguments(const ColumnsWithTypeAndName & arguments) { return arguments.size() - 1; }
+
+    bool insertResultToColumn(IColumn & dest, const Element & element, const std::string_view &)
+    {
+        if (!element.isObject())
+            return false;
+
+        auto object = element.getObject();
+
+        ColumnArray & col_res = assert_cast<ColumnArray &>(dest);
+        auto & col_key = assert_cast<ColumnString &>(col_res.getData());
+
+        for (const auto & [key, value] : object)
+        {
+            col_key.insertData(key.data(), key.size());
+        }
+
+        col_res.getOffsets().push_back(col_res.getOffsets().back() + object.size());
+        return true;
+    }
+};
+
 }

tests/fuzz/all.dict

@@ -650,6 +650,7 @@
 "JSONExtractInt"
 "JSONExtractKeysAndValues"
 "JSONExtractKeysAndValuesRaw"
+"JSONExtractKeys"
 "JSONExtractRaw"
 "JSONExtractString"
 "JSONExtractUInt"

tests/fuzz/dictionaries/functions.dict

@@ -31,6 +31,7 @@
 "toJSONString"
 "JSON_VALUE"
 "JSONExtractKeysAndValuesRaw"
+"JSONExtractKeys"
 "JSONExtractString"
 "JSONType"
 "JSONKey"

tests/queries/0_stateless/00918_json_functions.reference

@@ -203,6 +203,11 @@ Friday
 [('a','"hello"'),('b','[-100,200,300]')]
 [('a','"hello"'),('b','[-100,200,300]'),('c','{"d":[121,144]}')]
 [('d','[121,144]')]
+--JSONExtractKeys--
+['a','b']
+[]
+[]
+['d']
 --const/non-const mixed--
 a
 b

tests/queries/0_stateless/00918_json_functions.sql

@@ -231,6 +231,12 @@ SELECT JSONExtractKeysAndValuesRaw('{"a": "hello", "b": [-100, 200.0, 300]}');
 SELECT JSONExtractKeysAndValuesRaw('{"a": "hello", "b": [-100, 200.0, 300], "c":{"d":[121,144]}}');
 SELECT JSONExtractKeysAndValuesRaw('{"a": "hello", "b": [-100, 200.0, 300], "c":{"d":[121,144]}}', 'c');
 
+SELECT '--JSONExtractKeys--';
+SELECT JSONExtractKeys('{"a": "hello", "b": [-100, 200.0, 300]}');
+SELECT JSONExtractKeys('{"a": "hello", "b": [-100, 200.0, 300]}', 'b');
+SELECT JSONExtractKeys('{"a": "hello", "b": [-100, 200.0, 300]}', 'a');
+SELECT JSONExtractKeys('{"a": "hello", "b": [-100, 200.0, 300], "c":{"d":[121,144]}}', 'c');
+
 SELECT '--const/non-const mixed--';
 SELECT JSONExtractString('["a", "b", "c", "d", "e"]', idx) FROM (SELECT arrayJoin([1,2,3,4,5]) AS idx);
 SELECT JSONExtractString(json, 's') FROM (SELECT arrayJoin(['{"s":"u"}', '{"s":"v"}']) AS json);