Merge pull request #48207 from bigo-sg/map_from_maps

vdimir · web-flow · commit 35572321a14d · 2023-04-04T11:12:05.000+02:00
diff --git a/docs/en/sql-reference/functions/tuple-map-functions.md b/docs/en/sql-reference/functions/tuple-map-functions.md
@@ -68,9 +68,9 @@ Result:
 
 ## mapFromArrays
 
-Merges an [Array](../../sql-reference/data-types/array.md) of keys and an [Array](../../sql-reference/data-types/array.md) of values into a [Map(key, value)](../../sql-reference/data-types/map.md).
+Merges an [Array](../../sql-reference/data-types/array.md) of keys and an [Array](../../sql-reference/data-types/array.md) of values into a [Map(key, value)](../../sql-reference/data-types/map.md). Notice that the second argument could also be a [Map](../../sql-reference/data-types/map.md), thus it is casted to an Array when executing.
 
-The function is a more convenient alternative to `CAST((key_array, value_array), 'Map(key_type, value_type)')`. For example, instead of writing `CAST((['aa', 'bb'], [4, 5]), 'Map(String, UInt32)')`, you can write `mapFromArrays(['aa', 'bb'], [4, 5])`.
+The function is a more convenient alternative to `CAST((key_array, value_array_or_map), 'Map(key_type, value_type)')`. For example, instead of writing `CAST((['aa', 'bb'], [4, 5]), 'Map(String, UInt32)')`, you can write `mapFromArrays(['aa', 'bb'], [4, 5])`.
   
 **Syntax**
 
@@ -82,25 +82,29 @@ Alias: `MAP_FROM_ARRAYS(keys, values)`
   
 **Arguments**
 -   `keys` — Given key array to create a map from. The nested type of array must be: [String](../../sql-reference/data-types/string.md), [Integer](../../sql-reference/data-types/int-uint.md), [LowCardinality](../../sql-reference/data-types/lowcardinality.md), [FixedString](../../sql-reference/data-types/fixedstring.md), [UUID](../../sql-reference/data-types/uuid.md), [Date](../../sql-reference/data-types/date.md), [DateTime](../../sql-reference/data-types/datetime.md), [Date32](../../sql-reference/data-types/date32.md), [Enum](../../sql-reference/data-types/enum.md)
--   `values`  - Given value array to create a map from.
+-   `values`  - Given value array or map to create a map from.
   
 **Returned value**
 
-- A map whose keys and values are constructed from the key and value arrays
+- A map whose keys and values are constructed from the key array and value array/map.
   
 **Example**
 
 Query:
 
 ```sql
 select mapFromArrays(['a', 'b', 'c'], [1, 2, 3])
-```
-  
-```text
+
 ┌─mapFromArrays(['a', 'b', 'c'], [1, 2, 3])─┐
 │ {'a':1,'b':2,'c':3}                       │
 └───────────────────────────────────────────┘
-```  
+
+SELECT mapFromArrays([1, 2, 3], map('a', 1, 'b', 2, 'c', 3))
+
+┌─mapFromArrays([1, 2, 3], map('a', 1, 'b', 2, 'c', 3))─┐
+│ {1:('a',1),2:('b',2),3:('c',3)}                       │
+└───────────────────────────────────────────────────────┘
+```
 
 ## mapAdd
 
diff --git a/src/Functions/map.cpp b/src/Functions/map.cpp
@@ -174,23 +174,31 @@ class FunctionMapFromArrays : public IFunction
                 getName(),
                 arguments.size());
 
-        const auto * keys_type = checkAndGetDataType<DataTypeArray>(arguments[0].get());
-        if (!keys_type)
+        /// The first argument should always be Array.
+        /// Because key type can not be nested type of Map, which is Tuple
+        DataTypePtr key_type;
+        if (const auto * keys_type = checkAndGetDataType<DataTypeArray>(arguments[0].get()))
+            key_type = keys_type->getNestedType();
+        else
             throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "First argument for function {} must be an Array", getName());
 
-        const auto * values_type = checkAndGetDataType<DataTypeArray>(arguments[1].get());
-        if (!values_type)
-            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Second argument for function {} must be an Array", getName());
+        DataTypePtr value_type;
+        if (const auto * value_array_type = checkAndGetDataType<DataTypeArray>(arguments[1].get()))
+            value_type = value_array_type->getNestedType();
+        else if (const auto * value_map_type = checkAndGetDataType<DataTypeMap>(arguments[1].get()))
+            value_type = std::make_shared<DataTypeTuple>(value_map_type->getKeyValueTypes());
+        else
+            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Second argument for function {} must be Array or Map", getName());
 
-        DataTypes key_value_types{keys_type->getNestedType(), values_type->getNestedType()};
+        DataTypes key_value_types{key_type, value_type};
         return std::make_shared<DataTypeMap>(key_value_types);
     }
 
     ColumnPtr executeImpl(
         const ColumnsWithTypeAndName & arguments, const DataTypePtr & /* result_type */, size_t /* input_rows_count */) const override
     {
-        ColumnPtr holder_keys;
         bool is_keys_const = isColumnConst(*arguments[0].column);
+        ColumnPtr holder_keys;
         const ColumnArray * col_keys;
         if (is_keys_const)
         {
@@ -202,24 +210,26 @@ class FunctionMapFromArrays : public IFunction
             col_keys = checkAndGetColumn<ColumnArray>(arguments[0].column.get());
         }
 
-        ColumnPtr holder_values;
+        if (!col_keys)
+            throw Exception(ErrorCodes::ILLEGAL_COLUMN, "The first argument of function {} must be Array", getName());
+
         bool is_values_const = isColumnConst(*arguments[1].column);
-        const ColumnArray * col_values;
+        ColumnPtr holder_values;
         if (is_values_const)
-        {
             holder_values = arguments[1].column->convertToFullColumnIfConst();
-            col_values = checkAndGetColumn<ColumnArray>(holder_values.get());
-        }
         else
-        {
-            col_values = checkAndGetColumn<ColumnArray>(arguments[1].column.get());
-        }
+            holder_values = arguments[1].column;
 
-        if (!col_keys || !col_values)
-            throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Arguments of function {} must be array", getName());
+        const ColumnArray * col_values;
+        if (const auto * col_values_array = checkAndGetColumn<ColumnArray>(holder_values.get()))
+            col_values = col_values_array;
+        else if (const auto * col_values_map = checkAndGetColumn<ColumnMap>(holder_values.get()))
+            col_values = &col_values_map->getNestedColumn();
+        else
+            throw Exception(ErrorCodes::ILLEGAL_COLUMN, "The second arguments of function {} must be Array or Map", getName());
 
         if (!col_keys->hasEqualOffsets(*col_values))
-            throw Exception(ErrorCodes::SIZES_OF_ARRAYS_DONT_MATCH, "Array arguments for function {} must have equal sizes", getName());
+            throw Exception(ErrorCodes::SIZES_OF_ARRAYS_DONT_MATCH, "Two arguments for function {} must have equal sizes", getName());
 
         const auto & data_keys = col_keys->getDataPtr();
         const auto & data_values = col_values->getDataPtr();
diff --git a/tests/queries/0_stateless/01651_map_functions.reference b/tests/queries/0_stateless/01651_map_functions.reference
@@ -33,3 +33,6 @@
 {'aa':4,'bb':5}
 {'aa':4,'bb':5}
 {'aa':4,'bb':5}
+{'aa':('a',4),'bb':('b',5)}
+{'aa':('a',4),'bb':('b',5)}
+{'aa':('a',4),'bb':('b',5)}
diff --git a/tests/queries/0_stateless/01651_map_functions.sql b/tests/queries/0_stateless/01651_map_functions.sql
@@ -39,3 +39,8 @@ select mapFromArrays(['aa', 'bb'], 5); -- { serverError ILLEGAL_TYPE_OF_ARGUMENT
 select mapFromArrays(['aa', 'bb'], [4, 5], [6, 7]); -- { serverError NUMBER_OF_ARGUMENTS_DOESNT_MATCH }
 select mapFromArrays(['aa', 'bb'], [4, 5, 6]); -- { serverError SIZES_OF_ARRAYS_DONT_MATCH }
 select mapFromArrays([[1,2], [3,4]], [4, 5, 6]); -- { serverError BAD_ARGUMENTS }
+
+select mapFromArrays(['aa', 'bb'], map('a', 4, 'b', 5));
+select mapFromArrays(['aa', 'bb'], materialize(map('a', 4, 'b', 5))) from numbers(2);
+select mapFromArrays(map('a', 4, 'b', 4), ['aa', 'bb']) from numbers(2); -- { serverError ILLEGAL_TYPE_OF_ARGUMENT }
+select mapFromArrays(['aa', 'bb'], map('a', 4)); -- { serverError SIZES_OF_ARRAYS_DONT_MATCH }
