ClickHouse · evillique · Apr 14, 2023 · Apr 12, 2023
diff --git a/docs/en/sql-reference/functions/json-functions.md b/docs/en/sql-reference/functions/json-functions.md
@@ -206,6 +206,7 @@ Examples:
 ``` sql
 SELECT JSONExtract('{"a": "hello", "b": [-100, 200.0, 300]}', 'Tuple(String, Array(Float64))') = ('hello',[-100,200,300])
 SELECT JSONExtract('{"a": "hello", "b": [-100, 200.0, 300]}', 'Tuple(b Array(Float64), a String)') = ([-100,200,300],'hello')
+SELECT JSONExtract('{"a": "hello", "b": "world"}', 'Map(String, String)') = map('a',  'hello', 'b', 'world');
 SELECT JSONExtract('{"a": "hello", "b": [-100, 200.0, 300]}', 'b', 'Array(Nullable(Int8))') = [-100, NULL, NULL]
 SELECT JSONExtract('{"a": "hello", "b": [-100, 200.0, 300]}', 'b', 4, 'Nullable(Int64)') = NULL
 SELECT JSONExtract('{"passed": true}', 'passed', 'UInt8') = 1

diff --git a/src/Functions/FunctionsJSON.h b/src/Functions/FunctionsJSON.h
@@ -23,6 +23,7 @@
 #include <Columns/ColumnTuple.h>
 
 #include <DataTypes/DataTypeArray.h>
+#include <DataTypes/DataTypeMap.h>
 #include <DataTypes/DataTypeEnum.h>
 #include <DataTypes/DataTypeFactory.h>
 #include <DataTypes/DataTypeFixedString.h>
@@ -1190,6 +1191,46 @@ struct JSONExtractTree
         std::unordered_map<std::string_view, size_t> name_to_index_map;
     };
 
+    class MapNode : public Node
+    {
+    public:
+        MapNode(std::unique_ptr<Node> key_, std::unique_ptr<Node> value_) : key(std::move(key_)), value(std::move(value_)) { }
+
+        bool insertResultToColumn(IColumn & dest, const Element & element) override
+        {
+            if (!element.isObject())
+                return false;
+
+            ColumnMap & map_col = assert_cast<ColumnMap &>(dest);
+            auto & offsets = map_col.getNestedColumn().getOffsets();
+            auto & tuple_col = map_col.getNestedData();
+            auto & key_col = tuple_col.getColumn(0);
+            auto & value_col = tuple_col.getColumn(1);
+            size_t old_size = tuple_col.size();
+
+            auto object = element.getObject();
+            auto it = object.begin();
+            for (; it != object.end(); ++it)
+            {
+                auto pair = *it;
+
+                /// Insert key
+                key_col.insertData(pair.first.data(), pair.first.size());
+
+                /// Insert value
+                if (!value->insertResultToColumn(value_col, pair.second))
+                    value_col.insertDefault();
+            }
+
+            offsets.push_back(old_size + object.size());
+            return true;
+        }
+
+    private:
+        std::unique_ptr<Node> key;
+        std::unique_ptr<Node> value;
+    };
+
     static std::unique_ptr<Node> build(const char * function_name, const DataTypePtr & type)
     {
         switch (type->getTypeId())
@@ -1252,6 +1293,20 @@ struct JSONExtractTree
                     elements.emplace_back(build(function_name, tuple_element));
                 return std::make_unique<TupleNode>(std::move(elements), tuple.haveExplicitNames() ? tuple.getElementNames() : Strings{});
             }
+            case TypeIndex::Map:
+            {
+                const auto & map_type = static_cast<const DataTypeMap &>(*type);
+                const auto & key_type = map_type.getKeyType();
+                if (!isString(removeLowCardinality(key_type)))
+                    throw Exception(
+                        ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
+                        "Function {} doesn't support the return type schema: {} with key type not String",
+                        String(function_name),
+                        type->getName());
+
+                const auto & value_type = map_type.getValueType();
+                return std::make_unique<MapNode>(build(function_name, key_type), build(function_name, value_type));
+            }
             default:
                 throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
                                 "Function {} doesn't support the return type schema: {}",

diff --git a/tests/queries/0_stateless/00918_json_functions.reference b/tests/queries/0_stateless/00918_json_functions.reference
@@ -108,6 +108,12 @@ true	Bool
 123456789012	UInt64
 0	UInt64
 0	Int8
+{'a':'hello','b':'world'}
+{'a':'hello','b':'world'}
+{'a':('hello',100),'b':('world',200)}
+{'a':[100,200],'b':[-100,200,300]}
+{'a':{'c':'hello'},'b':{'d':'world'}}
+{'c':'hello'}
 --JSONExtractKeysAndValues--
 [('a','hello'),('b','[-100,200,300]')]
 [('b',[-100,200,300])]
@@ -152,6 +158,7 @@ e
 u
 v
 --show error: type should be const string
+--show error: key of map type should be String
 --allow_simdjson=0--
 --JSONLength--
 2
@@ -217,6 +224,12 @@ Friday
 (3,0)
 (3,5)
 (3,0)
+{'a':'hello','b':'world'}
+{'a':'hello','b':'world'}
+{'a':('hello',100),'b':('world',200)}
+{'a':[100,200],'b':[-100,200,300]}
+{'a':{'c':'hello'},'b':{'d':'world'}}
+{'c':'hello'}
 --JSONExtractKeysAndValues--
 [('a','hello'),('b','[-100,200,300]')]
 [('b',[-100,200,300])]
@@ -266,3 +279,4 @@ u
 v
 --show error: type should be const string
 --show error: index type should be integer
+--show error: key of map type should be String
diff --git a/tests/queries/0_stateless/00918_json_functions.sql b/tests/queries/0_stateless/00918_json_functions.sql
@@ -123,6 +123,13 @@ SELECT JSONExtract('{"a": "123456789012.345"}', 'a', 'UInt64') as a, toTypeName(
 SELECT JSONExtract('{"a": "-2000.22"}', 'a', 'UInt64') as a, toTypeName(a);
 SELECT JSONExtract('{"a": "-2000.22"}', 'a', 'Int8') as a, toTypeName(a);
 
+SELECT JSONExtract('{"a": "hello", "b": "world"}', 'Map(String, String)');
+SELECT JSONExtract('{"a": "hello", "b": "world"}', 'Map(LowCardinality(String), String)');
+SELECT JSONExtract('{"a": ["hello", 100.0], "b": ["world", 200]}', 'Map(String, Tuple(String, Float64))');
+SELECT JSONExtract('{"a": [100.0, 200], "b": [-100, 200.0, 300]}', 'Map(String, Array(Float64))');
+SELECT JSONExtract('{"a": {"c": "hello"}, "b": {"d": "world"}}', 'Map(String, Map(String, String))');
+SELECT JSONExtract('{"a": {"c": "hello"}, "b": {"d": "world"}}', 'a',  'Map(String, String)');
+
 SELECT '--JSONExtractKeysAndValues--';
 SELECT JSONExtractKeysAndValues('{"a": "hello", "b": [-100, 200.0, 300]}', 'String');
 SELECT JSONExtractKeysAndValues('{"a": "hello", "b": [-100, 200.0, 300]}', 'Array(Float64)');
@@ -166,8 +173,11 @@ SELECT JSONExtractString('["a", "b", "c", "d", "e"]', idx) FROM (SELECT arrayJoi
 SELECT JSONExtractString(json, 's') FROM (SELECT arrayJoin(['{"s":"u"}', '{"s":"v"}']) AS json);
 
 SELECT '--show error: type should be const string';
-SELECT JSONExtractKeysAndValues([], JSONLength('^?V{LSwp')); -- { serverError 44 }
-WITH '{"i": 1, "f": 1.2}' AS json SELECT JSONExtract(json, 'i', JSONType(json, 'i')); -- { serverError 44 }
+SELECT JSONExtractKeysAndValues([], JSONLength('^?V{LSwp')); -- { serverError ILLEGAL_COLUMN }
+WITH '{"i": 1, "f": 1.2}' AS json SELECT JSONExtract(json, 'i', JSONType(json, 'i')); -- { serverError ILLEGAL_COLUMN }
+
+SELECT '--show error: key of map type should be String';
+SELECT JSONExtract('{"a": [100.0, 200], "b": [-100, 200.0, 300]}', 'Map(Int64, Array(Float64))'); -- { serverError ILLEGAL_TYPE_OF_ARGUMENT }
 
 
 SELECT '--allow_simdjson=0--';
@@ -247,6 +257,13 @@ SELECT JSONExtract('{"a":3}', 'Tuple(Int, Int)');
 SELECT JSONExtract('[3,5,7]', 'Tuple(Int, Int)');
 SELECT JSONExtract('[3]', 'Tuple(Int, Int)');
 
+SELECT JSONExtract('{"a": "hello", "b": "world"}', 'Map(String, String)');
+SELECT JSONExtract('{"a": "hello", "b": "world"}', 'Map(LowCardinality(String), String)');
+SELECT JSONExtract('{"a": ["hello", 100.0], "b": ["world", 200]}', 'Map(String, Tuple(String, Float64))');
+SELECT JSONExtract('{"a": [100.0, 200], "b": [-100, 200.0, 300]}', 'Map(String, Array(Float64))');
+SELECT JSONExtract('{"a": {"c": "hello"}, "b": {"d": "world"}}', 'Map(String, Map(String, String))');
+SELECT JSONExtract('{"a": {"c": "hello"}, "b": {"d": "world"}}', 'a',  'Map(String, String)');
+
 SELECT '--JSONExtractKeysAndValues--';
 SELECT JSONExtractKeysAndValues('{"a": "hello", "b": [-100, 200.0, 300]}', 'String');
 SELECT JSONExtractKeysAndValues('{"a": "hello", "b": [-100, 200.0, 300]}', 'Array(Float64)');
@@ -295,8 +312,11 @@ SELECT JSONExtractString('["a", "b", "c", "d", "e"]', idx) FROM (SELECT arrayJoi
 SELECT JSONExtractString(json, 's') FROM (SELECT arrayJoin(['{"s":"u"}', '{"s":"v"}']) AS json);
 
 SELECT '--show error: type should be const string';
-SELECT JSONExtractKeysAndValues([], JSONLength('^?V{LSwp')); -- { serverError 44 }
-WITH '{"i": 1, "f": 1.2}' AS json SELECT JSONExtract(json, 'i', JSONType(json, 'i')); -- { serverError 44 }
+SELECT JSONExtractKeysAndValues([], JSONLength('^?V{LSwp')); -- { serverError ILLEGAL_COLUMN }
+WITH '{"i": 1, "f": 1.2}' AS json SELECT JSONExtract(json, 'i', JSONType(json, 'i')); -- { serverError ILLEGAL_COLUMN }
 
 SELECT '--show error: index type should be integer';
-SELECT JSONExtract('[]', JSONExtract('0', 'UInt256'), 'UInt256'); -- { serverError 43 }
+SELECT JSONExtract('[]', JSONExtract('0', 'UInt256'), 'UInt256'); -- { serverError ILLEGAL_TYPE_OF_ARGUMENT }
+
+SELECT '--show error: key of map type should be String';
+SELECT JSONExtract('{"a": [100.0, 200], "b": [-100, 200.0, 300]}', 'Map(Int64, Array(Float64))'); -- { serverError ILLEGAL_TYPE_OF_ARGUMENT }