Merge pull request #26663 from ildus/mappopulateseries_for_maps

Add Map type support in mapPopulateSeries function

docs/en/sql-reference/functions/tuple-map-functions.md

@@ -11,7 +11,7 @@ Arranges `key:value` pairs into [Map(key, value)](../../sql-reference/data-types
 
 **Syntax**
 
-``` sql
+```sql
 map(key1, value1[, key2, value2, ...])
 ```
 
@@ -30,7 +30,7 @@ Type: [Map(key, value)](../../sql-reference/data-types/map.md).
 
 Query:
 
-``` sql
+```sql
 SELECT map('key1', number, 'key2', number * 2) FROM numbers(3);
 ```
 
@@ -46,7 +46,7 @@ Result:
 
 Query:
 
-``` sql
+```sql
 CREATE TABLE table_map (a Map(String, UInt64)) ENGINE = MergeTree() ORDER BY a;
 INSERT INTO table_map SELECT map('key1', number, 'key2', number * 2) FROM numbers(3);
 SELECT a['key2'] FROM table_map;
@@ -54,7 +54,7 @@ SELECT a['key2'] FROM table_map;
 
 Result:
 
-``` text
+```text
 ┌─arrayElement(a, 'key2')─┐
 │                       0 │
 │                       2 │
@@ -72,7 +72,7 @@ Collect all the keys and sum corresponding values.
 
 **Syntax**
 
-``` sql
+```sql
 mapAdd(arg1, arg2 [, ...])
 ```
 
@@ -88,13 +88,13 @@ Arguments are [maps](../../sql-reference/data-types/map.md) or [tuples](../../sq
 
 Query with a tuple map:
 
-``` sql
+```sql
 SELECT mapAdd(([toUInt8(1), 2], [1, 1]), ([toUInt8(1), 2], [1, 1])) as res, toTypeName(res) as type;
 ```
 
 Result:
 
-``` text
+```text
 ┌─res───────────┬─type───────────────────────────────┐
 │ ([1,2],[2,2]) │ Tuple(Array(UInt8), Array(UInt64)) │
 └───────────────┴────────────────────────────────────┘
@@ -102,7 +102,16 @@ Result:
 
 Query with `Map` type:
 
-``` sql
+```sql
+SELECT mapAdd(map(1,1), map(1,1));
+```
+
+Result:
+
+```text
+┌─mapAdd(map(1, 1), map(1, 1))─┐
+│ {1:2}                        │
+└──────────────────────────────┘
 ```
 
 ## mapSubtract {#function-mapsubtract}
@@ -111,21 +120,21 @@ Collect all the keys and subtract corresponding values.
 
 **Syntax**
 
-``` sql
+```sql
 mapSubtract(Tuple(Array, Array), Tuple(Array, Array) [, ...])
 ```
 
 **Arguments**
 
-Arguments are [tuples](../../sql-reference/data-types/tuple.md#tuplet1-t2) of two [arrays](../../sql-reference/data-types/array.md#data-type-array), where items in the first array represent keys, and the second array contains values for the each key. All key arrays should have same type, and all value arrays should contain items which are promote to the one type ([Int64](../../sql-reference/data-types/int-uint.md#int-ranges), [UInt64](../../sql-reference/data-types/int-uint.md#uint-ranges) or [Float64](../../sql-reference/data-types/float.md#float32-float64)). The common promoted type is used as a type for the result array.
+Arguments are [maps](../../sql-reference/data-types/map.md) or [tuples](../../sql-reference/data-types/tuple.md#tuplet1-t2) of two [arrays](../../sql-reference/data-types/array.md#data-type-array), where items in the first array represent keys, and the second array contains values for the each key. All key arrays should have same type, and all value arrays should contain items which are promote to the one type ([Int64](../../sql-reference/data-types/int-uint.md#int-ranges), [UInt64](../../sql-reference/data-types/int-uint.md#uint-ranges) or [Float64](../../sql-reference/data-types/float.md#float32-float64)). The common promoted type is used as a type for the result array.
 
 **Returned value**
 
--   Returns one [tuple](../../sql-reference/data-types/tuple.md#tuplet1-t2), where the first array contains the sorted keys and the second array contains values.
+-   Depending on the arguments returns one [map](../../sql-reference/data-types/map.md) or [tuple](../../sql-reference/data-types/tuple.md#tuplet1-t2), where the first array contains the sorted keys and the second array contains values.
 
 **Example**
 
-Query:
+Query with a tuple map:
 
 ```sql
 SELECT mapSubtract(([toUInt8(1), 2], [toInt32(1), 1]), ([toUInt8(1), 2], [toInt32(2), 1])) as res, toTypeName(res) as type;
@@ -139,32 +148,54 @@ Result:
 └────────────────┴───────────────────────────────────┘
 ```
 
+Query with `Map` type:
+
+```sql
+SELECT mapSubtract(map(1,1), map(1,1));
+```
+
+Result:
+
+```text
+┌─mapSubtract(map(1, 1), map(1, 1))─┐
+│ {1:0}                             │
+└───────────────────────────────────┘
+```
+
 ## mapPopulateSeries {#function-mappopulateseries}
 
 Fills missing keys in the maps (key and value array pair), where keys are integers. Also, it supports specifying the max key, which is used to extend the keys array.
+Arguments are [maps](../../sql-reference/data-types/map.md) or two [arrays](../../sql-reference/data-types/array.md#data-type-array), where the first array represent keys, and the second array contains values for the each key.
+
+For array arguments the number of elements in `keys` and `values` must be the same for each row.
 
 **Syntax**
 
-``` sql
+```sql
 mapPopulateSeries(keys, values[, max])
+mapPopulateSeries(map[, max])
 ```
 
-Generates a map, where keys are a series of numbers, from minimum to maximum keys (or `max` argument if it specified) taken from `keys` array with a step size of one, and corresponding values taken from `values` array. If the value is not specified for the key, then it uses the default value in the resulting map. For repeated keys, only the first value (in order of appearing) gets associated with the key.
-
-The number of elements in `keys` and `values` must be the same for each row.
+Generates a map (a tuple with two arrays or a value of `Map` type, depending on the arguments), where keys are a series of numbers, from minimum to maximum keys (or `max` argument if it specified) taken from the map with a step size of one, and corresponding values. If the value is not specified for the key, then it uses the default value in the resulting map. For repeated keys, only the first value (in order of appearing) gets associated with the key.
 
 **Arguments**
 
+Mapped arrays:
+
 -   `keys` — Array of keys. [Array](../../sql-reference/data-types/array.md#data-type-array)([Int](../../sql-reference/data-types/int-uint.md#uint-ranges)).
 -   `values` — Array of values. [Array](../../sql-reference/data-types/array.md#data-type-array)([Int](../../sql-reference/data-types/int-uint.md#uint-ranges)).
 
+or
+
+-   `map` — Map with integer keys. [Map](../../sql-reference/data-types/map.md).
+
 **Returned value**
 
--  Returns a [tuple](../../sql-reference/data-types/tuple.md#tuplet1-t2) of two [arrays](../../sql-reference/data-types/array.md#data-type-array): keys in sorted order, and values the corresponding keys.
+-  Depending on the arguments returns a [map](../../sql-reference/data-types/map.md) or a [tuple](../../sql-reference/data-types/tuple.md#tuplet1-t2) of two [arrays](../../sql-reference/data-types/array.md#data-type-array): keys in sorted order, and values the corresponding keys.
 
 **Example**
 
-Query:
+Query with mapped arrays:
 
 ```sql
 select mapPopulateSeries([1,2,4], [11,22,44], 5) as res, toTypeName(res) as type;
@@ -178,13 +209,27 @@ Result:
 └──────────────────────────────┴───────────────────────────────────┘
 ```
 
+Query with `Map` type:
+
+```sql
+SELECT mapPopulateSeries(map(1, 10, 5, 20), 6);
+```
+
+Result:
+
+```text
+┌─mapPopulateSeries(map(1, 10, 5, 20), 6)─┐
+│ {1:10,2:0,3:0,4:0,5:20,6:0}             │
+└─────────────────────────────────────────┘
+```
+
 ## mapContains {#mapcontains}
 
 Determines  whether the `map` contains the `key` parameter.
 
 **Syntax**
 
-``` sql
+```sql
 mapContains(map, key)
 ```
 

src/Functions/array/mapPopulateSeries.cpp

@@ -1,4 +1,5 @@
 #include <Columns/ColumnArray.h>
+#include <Columns/ColumnMap.h>
 #include <Columns/ColumnTuple.h>
 #include <Columns/ColumnVector.h>
 #include <DataTypes/DataTypeArray.h>
@@ -7,6 +8,7 @@
 #include <Functions/FunctionHelpers.h>
 #include <Functions/IFunction.h>
 #include "Core/ColumnWithTypeAndName.h"
+#include "DataTypes/DataTypeMap.h"
 #include "DataTypes/IDataType.h"
 
 namespace DB
@@ -32,85 +34,211 @@ private:
     bool isVariadic() const override { return true; }
     bool useDefaultImplementationForConstants() const override { return true; }
 
-    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
+    void checkTypes(const DataTypePtr & key_type, const DataTypePtr max_key_type) const
+    {
+        WhichDataType which_key(key_type);
+        if (!(which_key.isInt() || which_key.isUInt()))
+        {
+            throw Exception(
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Keys for {} function should be of integer type (signed or unsigned)", getName());
+        }
+
+        if (max_key_type)
+        {
+            WhichDataType which_max_key(max_key_type);
+
+            if (which_max_key.isNullable())
+                throw Exception(
+                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
+                    "Max key argument in arguments of function " + getName() + " can not be Nullable");
+
+            if (key_type->getTypeId() != max_key_type->getTypeId())
+                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Max key type in {} should be same as keys type", getName());
+        }
+    }
+
+    DataTypePtr getReturnTypeForTuple(const DataTypes & arguments) const
     {
         if (arguments.size() < 2)
-            throw Exception{getName() + " accepts at least two arrays for key and value", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH};
+            throw Exception(
+                ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function {} accepts at least two arrays for key and value", getName());
 
         if (arguments.size() > 3)
-            throw Exception{"too many arguments in " + getName() + " call", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH};
+            throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Too many arguments in {} call", getName());
 
         const DataTypeArray * key_array_type = checkAndGetDataType<DataTypeArray>(arguments[0].get());
         const DataTypeArray * val_array_type = checkAndGetDataType<DataTypeArray>(arguments[1].get());
 
         if (!key_array_type || !val_array_type)
-            throw Exception{getName() + " accepts two arrays for key and value", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT};
+            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Function {} accepts two arrays for key and value", getName());
 
-        DataTypePtr keys_type = key_array_type->getNestedType();
-        WhichDataType which_key(keys_type);
-        if (!(which_key.isNativeInt() || which_key.isNativeUInt()))
-        {
-            throw Exception(
-                "Keys for " + getName() + " should be of native integer type (signed or unsigned)", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
-        }
+        const auto & key_type = key_array_type->getNestedType();
 
         if (arguments.size() == 3)
-        {
-            DataTypePtr max_key_type = arguments[2];
-            WhichDataType which_max_key(max_key_type);
-
-            if (which_max_key.isNullable())
-                throw Exception(
-                    "Max key argument in arguments of function " + getName() + " can not be Nullable",
-                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
-
-            if (keys_type->getTypeId() != max_key_type->getTypeId())
-                throw Exception("Max key type in " + getName() + " should be same as keys type", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
-        }
+            this->checkTypes(key_type, arguments[2]);
+        else
+            this->checkTypes(key_type, nullptr);
 
         return std::make_shared<DataTypeTuple>(DataTypes{arguments[0], arguments[1]});
     }
 
-    template <typename KeyType, typename ValType>
-    ColumnPtr execute2(ColumnPtr key_column, ColumnPtr val_column, ColumnPtr max_key_column, const DataTypeTuple & res_type) const
+    DataTypePtr getReturnTypeForMap(const DataTypes & arguments) const
     {
-        MutableColumnPtr res_tuple = res_type.createColumn();
+        const auto * map = assert_cast<const DataTypeMap *>(arguments[0].get());
+        if (arguments.size() == 1)
+            this->checkTypes(map->getKeyType(), nullptr);
+        else if (arguments.size() == 2)
+            this->checkTypes(map->getKeyType(), arguments[1]);
+        else
+            throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Too many arguments in {} call", getName());
 
-        auto * to_tuple = assert_cast<ColumnTuple *>(res_tuple.get());
-        auto & to_keys_arr = assert_cast<ColumnArray &>(to_tuple->getColumn(0));
-        auto & to_keys_data = to_keys_arr.getData();
-        auto & to_keys_offsets = to_keys_arr.getOffsets();
+        return std::make_shared<DataTypeMap>(map->getKeyType(), map->getValueType());
+    }
 
-        auto & to_vals_arr = assert_cast<ColumnArray &>(to_tuple->getColumn(1));
-        auto & to_values_data = to_vals_arr.getData();
+    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
+    {
+        if (arguments.empty())
+            throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, getName() + " accepts at least one map or two arrays");
 
-        bool max_key_is_const = false, key_is_const = false, val_is_const = false;
+        if (arguments[0]->getTypeId() == TypeIndex::Array)
+            return getReturnTypeForTuple(arguments);
+        else if (arguments[0]->getTypeId() == TypeIndex::Map)
+            return getReturnTypeForMap(arguments);
+        else
+            throw Exception(
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
+                "Function {} only accepts one map or arrays, but got {}",
+                getName(),
+                arguments[0]->getName());
+    }
 
-        const auto * keys_array = checkAndGetColumn<ColumnArray>(key_column.get());
-        if (!keys_array)
+    // Struct holds input and output columns references,
+    // Both arrays and maps have similar columns to work with but extracted differently
+    template <typename KeyType, typename ValType>
+    struct ColumnsInOut
+    {
+        // inputs
+        const PaddedPODArray<KeyType> & in_keys_data;
+        const PaddedPODArray<ValType> & in_vals_data;
+        const IColumn::Offsets & in_key_offsets;
+        const IColumn::Offsets & in_val_offsets;
+        size_t row_count;
+        bool key_is_const;
+        bool val_is_const;
+
+        // outputs
+        PaddedPODArray<KeyType> & out_keys_data;
+        PaddedPODArray<ValType> & out_vals_data;
+
+        IColumn::Offsets & out_keys_offsets;
+        // with map argument this field will not be used
+        IColumn::Offsets * out_vals_offsets;
+    };
+
+    template <typename KeyType, typename ValType>
+    ColumnsInOut<KeyType, ValType> getInOutDataFromArrays(MutableColumnPtr & res_column, ColumnPtr * arg_columns) const
+    {
+        auto * out_tuple = assert_cast<ColumnTuple *>(res_column.get());
+        auto & out_keys_array = assert_cast<ColumnArray &>(out_tuple->getColumn(0));
+        auto & out_vals_array = assert_cast<ColumnArray &>(out_tuple->getColumn(1));
+
+        const auto * key_column = arg_columns[0].get();
+        const auto * in_keys_array = checkAndGetColumn<ColumnArray>(key_column);
+
+        bool key_is_const = false, val_is_const = false;
+
+        if (!in_keys_array)
         {
-            const ColumnConst * const_array = checkAndGetColumnConst<ColumnArray>(key_column.get());
+            const ColumnConst * const_array = checkAndGetColumnConst<ColumnArray>(key_column);
             if (!const_array)
-                throw Exception("Expected array column, found " + key_column->getName(), ErrorCodes::ILLEGAL_COLUMN);
+                throw Exception(
+                    ErrorCodes::ILLEGAL_COLUMN, "Expected array column in function {}, found {}", getName(), key_column->getName());
 
-            keys_array = checkAndGetColumn<ColumnArray>(const_array->getDataColumnPtr().get());
+            in_keys_array = checkAndGetColumn<ColumnArray>(const_array->getDataColumnPtr().get());
             key_is_const = true;
         }
 
-        const auto * values_array = checkAndGetColumn<ColumnArray>(val_column.get());
-        if (!values_array)
+        const auto * val_column = arg_columns[1].get();
+        const auto * in_values_array = checkAndGetColumn<ColumnArray>(val_column);
+        if (!in_values_array)
         {
-            const ColumnConst * const_array = checkAndGetColumnConst<ColumnArray>(val_column.get());
+            const ColumnConst * const_array = checkAndGetColumnConst<ColumnArray>(val_column);
             if (!const_array)
-                throw Exception("Expected array column, found " + val_column->getName(), ErrorCodes::ILLEGAL_COLUMN);
+                throw Exception(
+                    ErrorCodes::ILLEGAL_COLUMN, "Expected array column in function {}, found {}", getName(), val_column->getName());
 
-            values_array = checkAndGetColumn<ColumnArray>(const_array->getDataColumnPtr().get());
+            in_values_array = checkAndGetColumn<ColumnArray>(const_array->getDataColumnPtr().get());
             val_is_const = true;
         }
 
-        if (!keys_array || !values_array)
+        if (!in_keys_array || !in_values_array)
             /* something went wrong */
-            throw Exception{"Illegal columns in arguments of function " + getName(), ErrorCodes::ILLEGAL_COLUMN};
+            throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Illegal columns in arguments of function " + getName());
+
+        const auto & in_keys_data = assert_cast<const ColumnVector<KeyType> &>(in_keys_array->getData()).getData();
+        const auto & in_values_data = assert_cast<const ColumnVector<ValType> &>(in_values_array->getData()).getData();
+        const auto & in_keys_offsets = in_keys_array->getOffsets();
+        const auto & in_vals_offsets = in_values_array->getOffsets();
+
+        auto & out_keys_data = assert_cast<ColumnVector<KeyType> &>(out_keys_array.getData()).getData();
+        auto & out_vals_data = assert_cast<ColumnVector<ValType> &>(out_vals_array.getData()).getData();
+        auto & out_keys_offsets = out_keys_array.getOffsets();
+
+        size_t row_count = key_is_const ? in_values_array->size() : in_keys_array->size();
+        IColumn::Offsets * out_vals_offsets = &out_vals_array.getOffsets();
+
+        return {
+            in_keys_data,
+            in_values_data,
+            in_keys_offsets,
+            in_vals_offsets,
+            row_count,
+            key_is_const,
+            val_is_const,
+            out_keys_data,
+            out_vals_data,
+            out_keys_offsets,
+            out_vals_offsets};
+    }
+
+    template <typename KeyType, typename ValType>
+    ColumnsInOut<KeyType, ValType> getInOutDataFromMap(MutableColumnPtr & res_column, ColumnPtr * arg_columns) const
+    {
+        const auto * in_map = assert_cast<const ColumnMap *>(arg_columns[0].get());
+        const auto & in_nested_array = in_map->getNestedColumn();
+        const auto & in_nested_tuple = in_map->getNestedData();
+        const auto & in_keys_data = assert_cast<const ColumnVector<KeyType> &>(in_nested_tuple.getColumn(0)).getData();
+        const auto & in_vals_data = assert_cast<const ColumnVector<ValType> &>(in_nested_tuple.getColumn(1)).getData();
+        const auto & in_keys_offsets = in_nested_array.getOffsets();
+
+        auto * out_map = assert_cast<ColumnMap *>(res_column.get());
+        auto & out_nested_array = out_map->getNestedColumn();
+        auto & out_nested_tuple = out_map->getNestedData();
+        auto & out_keys_data = assert_cast<ColumnVector<KeyType> &>(out_nested_tuple.getColumn(0)).getData();
+        auto & out_vals_data = assert_cast<ColumnVector<ValType> &>(out_nested_tuple.getColumn(1)).getData();
+        auto & out_keys_offsets = out_nested_array.getOffsets();
+
+        return {
+            in_keys_data,
+            in_vals_data,
+            in_keys_offsets,
+            in_keys_offsets,
+            in_nested_array.size(),
+            false,
+            false,
+            out_keys_data,
+            out_vals_data,
+            out_keys_offsets,
+            nullptr};
+    }
+
+    template <typename KeyType, typename ValType>
+    ColumnPtr execute2(ColumnPtr * arg_columns, ColumnPtr max_key_column, const DataTypePtr & res_type) const
+    {
+        MutableColumnPtr res_column = res_type->createColumn();
+        bool max_key_is_const = false;
+        auto columns = res_column->getDataType() == TypeIndex::Tuple ? getInOutDataFromArrays<KeyType, ValType>(res_column, arg_columns)
+                                                                     : getInOutDataFromMap<KeyType, ValType>(res_column, arg_columns);
 
         KeyType max_key_const{0};
 
@@ -121,49 +249,43 @@ private:
             max_key_is_const = true;
         }
 
-        auto & keys_data = assert_cast<const ColumnVector<KeyType> &>(keys_array->getData()).getData();
-        auto & values_data = assert_cast<const ColumnVector<ValType> &>(values_array->getData()).getData();
-
-        // Original offsets
-        const IColumn::Offsets & key_offsets = keys_array->getOffsets();
-        const IColumn::Offsets & val_offsets = values_array->getOffsets();
-
         IColumn::Offset offset{0};
-        size_t row_count = key_is_const ? values_array->size() : keys_array->size();
-
         std::map<KeyType, ValType> res_map;
 
         //Iterate through two arrays and fill result values.
-        for (size_t row = 0; row < row_count; ++row)
+        for (size_t row = 0; row < columns.row_count; ++row)
         {
-            size_t key_offset = 0, val_offset = 0, array_size = key_offsets[0], val_array_size = val_offsets[0];
+            size_t key_offset = 0, val_offset = 0, items_count = columns.in_key_offsets[0], val_array_size = columns.in_val_offsets[0];
 
             res_map.clear();
 
-            if (!key_is_const)
+            if (!columns.key_is_const)
             {
-                key_offset = row > 0 ? key_offsets[row - 1] : 0;
-                array_size = key_offsets[row] - key_offset;
+                key_offset = row > 0 ? columns.in_key_offsets[row - 1] : 0;
+                items_count = columns.in_key_offsets[row] - key_offset;
             }
 
-            if (!val_is_const)
+            if (!columns.val_is_const)
             {
-                val_offset = row > 0 ? val_offsets[row - 1] : 0;
-                val_array_size = val_offsets[row] - val_offset;
+                val_offset = row > 0 ? columns.in_val_offsets[row - 1] : 0;
+                val_array_size = columns.in_val_offsets[row] - val_offset;
             }
 
-            if (array_size != val_array_size)
-                throw Exception("Key and value array should have same amount of elements", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+            if (items_count != val_array_size)
+                throw Exception(
+                    ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
+                    "Key and value array should have same amount of elements in function {}",
+                    getName());
 
-            if (array_size == 0)
+            if (items_count == 0)
             {
-                to_keys_offsets.push_back(offset);
+                columns.out_keys_offsets.push_back(offset);
                 continue;
             }
 
-            for (size_t i = 0; i < array_size; ++i)
+            for (size_t i = 0; i < items_count; ++i)
             {
-                res_map.insert({keys_data[key_offset + i], values_data[val_offset + i]});
+                res_map.insert({columns.in_keys_data[key_offset + i], columns.in_vals_data[val_offset + i]});
             }
 
             auto min_key = res_map.begin()->first;
@@ -184,7 +306,7 @@ private:
                 /* no need to add anything, max key is less that first key */
                 if (max_key < min_key)
                 {
-                    to_keys_offsets.push_back(offset);
+                    columns.out_keys_offsets.push_back(offset);
                     continue;
                 }
             }
@@ -197,16 +319,16 @@ private:
             KeyType key;
             for (key = min_key;; ++key)
             {
-                to_keys_data.insert(key);
+                columns.out_keys_data.push_back(key);
 
                 auto it = res_map.find(key);
                 if (it != res_map.end())
                 {
-                    to_values_data.insert(it->second);
+                    columns.out_vals_data.push_back(it->second);
                 }
                 else
                 {
-                    to_values_data.insertDefault();
+                    columns.out_vals_data.push_back(0);
                 }
 
                 ++offset;
@@ -214,80 +336,112 @@ private:
                     break;
             }
 
-            to_keys_offsets.push_back(offset);
+            columns.out_keys_offsets.push_back(offset);
         }
 
-        to_vals_arr.getOffsets().insert(to_keys_offsets.begin(), to_keys_offsets.end());
-        return res_tuple;
+        if (columns.out_vals_offsets)
+            columns.out_vals_offsets->insert(columns.out_keys_offsets.begin(), columns.out_keys_offsets.end());
+
+        return res_column;
     }
 
     template <typename KeyType>
-    ColumnPtr execute1(ColumnPtr key_column, ColumnPtr val_column, ColumnPtr max_key_column, const DataTypeTuple & res_type) const
+    ColumnPtr execute1(ColumnPtr * arg_columns, ColumnPtr max_key_column, const DataTypePtr & res_type, const DataTypePtr & val_type) const
     {
-        const auto & val_type = (assert_cast<const DataTypeArray *>(res_type.getElements()[1].get()))->getNestedType();
         switch (val_type->getTypeId())
         {
             case TypeIndex::Int8:
-                return execute2<KeyType, Int8>(key_column, val_column, max_key_column, res_type);
+                return execute2<KeyType, Int8>(arg_columns, max_key_column, res_type);
             case TypeIndex::Int16:
-                return execute2<KeyType, Int16>(key_column, val_column, max_key_column, res_type);
+                return execute2<KeyType, Int16>(arg_columns, max_key_column, res_type);
             case TypeIndex::Int32:
-                return execute2<KeyType, Int32>(key_column, val_column, max_key_column, res_type);
+                return execute2<KeyType, Int32>(arg_columns, max_key_column, res_type);
             case TypeIndex::Int64:
-                return execute2<KeyType, Int64>(key_column, val_column, max_key_column, res_type);
+                return execute2<KeyType, Int64>(arg_columns, max_key_column, res_type);
+            case TypeIndex::Int128:
+                return execute2<KeyType, Int128>(arg_columns, max_key_column, res_type);
+            case TypeIndex::Int256:
+                return execute2<KeyType, Int256>(arg_columns, max_key_column, res_type);
             case TypeIndex::UInt8:
-                return execute2<KeyType, UInt8>(key_column, val_column, max_key_column, res_type);
+                return execute2<KeyType, UInt8>(arg_columns, max_key_column, res_type);
             case TypeIndex::UInt16:
-                return execute2<KeyType, UInt16>(key_column, val_column, max_key_column, res_type);
+                return execute2<KeyType, UInt16>(arg_columns, max_key_column, res_type);
             case TypeIndex::UInt32:
-                return execute2<KeyType, UInt32>(key_column, val_column, max_key_column, res_type);
+                return execute2<KeyType, UInt32>(arg_columns, max_key_column, res_type);
             case TypeIndex::UInt64:
-                return execute2<KeyType, UInt64>(key_column, val_column, max_key_column, res_type);
+                return execute2<KeyType, UInt64>(arg_columns, max_key_column, res_type);
+            case TypeIndex::UInt128:
+                return execute2<KeyType, UInt128>(arg_columns, max_key_column, res_type);
+            case TypeIndex::UInt256:
+                return execute2<KeyType, UInt256>(arg_columns, max_key_column, res_type);
             default:
-                throw Exception{"Illegal columns in arguments of function " + getName(), ErrorCodes::ILLEGAL_COLUMN};
+                throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Illegal columns in arguments of function " + getName());
         }
     }
 
     ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t) const override
     {
-        auto col1 = arguments[0];
-        auto col2 = arguments[1];
-
-        const auto * k = assert_cast<const DataTypeArray *>(col1.type.get());
-        const auto * v = assert_cast<const DataTypeArray *>(col2.type.get());
-
-        /* determine output type */
-        const DataTypeTuple & res_type = DataTypeTuple(
-            DataTypes{std::make_shared<DataTypeArray>(k->getNestedType()), std::make_shared<DataTypeArray>(v->getNestedType())});
-
+        DataTypePtr res_type, key_type, val_type;
         ColumnPtr max_key_column = nullptr;
+        ColumnPtr arg_columns[] = {arguments[0].column, nullptr};
 
-        if (arguments.size() == 3)
+        if (arguments[0].type->getTypeId() == TypeIndex::Array)
         {
-            /* max key provided */
-            max_key_column = arguments[2].column;
+            key_type = assert_cast<const DataTypeArray *>(arguments[0].type.get())->getNestedType();
+            val_type = assert_cast<const DataTypeArray *>(arguments[1].type.get())->getNestedType();
+            res_type = getReturnTypeImpl(DataTypes{arguments[0].type, arguments[1].type});
+
+            arg_columns[1] = arguments[1].column;
+            if (arguments.size() == 3)
+            {
+                /* max key provided */
+                max_key_column = arguments[2].column;
+            }
+        }
+        else
+        {
+            assert(arguments[0].type->getTypeId() == TypeIndex::Map);
+
+            const auto * map_type = assert_cast<const DataTypeMap *>(arguments[0].type.get());
+            res_type = getReturnTypeImpl(DataTypes{arguments[0].type});
+            key_type = map_type->getKeyType();
+            val_type = map_type->getValueType();
+
+            if (arguments.size() == 2)
+            {
+                /* max key provided */
+                max_key_column = arguments[1].column;
+            }
         }
 
-        switch (k->getNestedType()->getTypeId())
+        switch (key_type->getTypeId())
         {
             case TypeIndex::Int8:
-                return execute1<Int8>(col1.column, col2.column, max_key_column, res_type);
+                return execute1<Int8>(arg_columns, max_key_column, res_type, val_type);
             case TypeIndex::Int16:
-                return execute1<Int16>(col1.column, col2.column, max_key_column, res_type);
+                return execute1<Int16>(arg_columns, max_key_column, res_type, val_type);
             case TypeIndex::Int32:
-                return execute1<Int32>(col1.column, col2.column, max_key_column, res_type);
+                return execute1<Int32>(arg_columns, max_key_column, res_type, val_type);
             case TypeIndex::Int64:
-                return execute1<Int64>(col1.column, col2.column, max_key_column, res_type);
+                return execute1<Int64>(arg_columns, max_key_column, res_type, val_type);
+            case TypeIndex::Int128:
+                return execute1<Int128>(arg_columns, max_key_column, res_type, val_type);
+            case TypeIndex::Int256:
+                return execute1<Int256>(arg_columns, max_key_column, res_type, val_type);
             case TypeIndex::UInt8:
-                return execute1<UInt8>(col1.column, col2.column, max_key_column, res_type);
+                return execute1<UInt8>(arg_columns, max_key_column, res_type, val_type);
             case TypeIndex::UInt16:
-                return execute1<UInt16>(col1.column, col2.column, max_key_column, res_type);
+                return execute1<UInt16>(arg_columns, max_key_column, res_type, val_type);
             case TypeIndex::UInt32:
-                return execute1<UInt32>(col1.column, col2.column, max_key_column, res_type);
+                return execute1<UInt32>(arg_columns, max_key_column, res_type, val_type);
             case TypeIndex::UInt64:
-                return execute1<UInt64>(col1.column, col2.column, max_key_column, res_type);
+                return execute1<UInt64>(arg_columns, max_key_column, res_type, val_type);
+            case TypeIndex::UInt128:
+                return execute1<UInt128>(arg_columns, max_key_column, res_type, val_type);
+            case TypeIndex::UInt256:
+                return execute1<UInt256>(arg_columns, max_key_column, res_type, val_type);
             default:
-                throw Exception{"Illegal columns in arguments of function " + getName(), ErrorCodes::ILLEGAL_COLUMN};
+                throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Illegal columns in arguments of function " + getName());
         }
     }
 };
@@ -296,5 +450,4 @@ void registerFunctionMapPopulateSeries(FunctionFactory & factory)
 {
     factory.registerFunction<FunctionMapPopulateSeries>();
 }
-
 }

tests/queries/0_stateless/01925_map_populate_series_on_map.reference

@@ -0,0 +1,67 @@
+-- { echo }
+drop table if exists map_test;
+set allow_experimental_map_type = 1;
+create table map_test engine=TinyLog() as (select (number + 1) as n, map(1, 1, number,2) as m from numbers(1, 5));
+select mapPopulateSeries(m) from map_test;
+{1:1}
+{1:1,2:2}
+{1:1,2:0,3:2}
+{1:1,2:0,3:0,4:2}
+{1:1,2:0,3:0,4:0,5:2}
+select mapPopulateSeries(m, toUInt64(3)) from map_test;
+{1:1,2:0,3:0}
+{1:1,2:2,3:0}
+{1:1,2:0,3:2}
+{1:1,2:0,3:0}
+{1:1,2:0,3:0}
+select mapPopulateSeries(m, toUInt64(10)) from map_test;
+{1:1,2:0,3:0,4:0,5:0,6:0,7:0,8:0,9:0,10:0}
+{1:1,2:2,3:0,4:0,5:0,6:0,7:0,8:0,9:0,10:0}
+{1:1,2:0,3:2,4:0,5:0,6:0,7:0,8:0,9:0,10:0}
+{1:1,2:0,3:0,4:2,5:0,6:0,7:0,8:0,9:0,10:0}
+{1:1,2:0,3:0,4:0,5:2,6:0,7:0,8:0,9:0,10:0}
+select mapPopulateSeries(m, 1000) from map_test; -- { serverError 43 }
+select mapPopulateSeries(m, n) from map_test;
+{1:1,2:0}
+{1:1,2:2,3:0}
+{1:1,2:0,3:2,4:0}
+{1:1,2:0,3:0,4:2,5:0}
+{1:1,2:0,3:0,4:0,5:2,6:0}
+drop table map_test;
+select mapPopulateSeries(map(toUInt8(1), toUInt8(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(UInt8,UInt8)
+select mapPopulateSeries(map(toUInt16(1), toUInt16(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(UInt16,UInt16)
+select mapPopulateSeries(map(toUInt32(1), toUInt32(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(UInt32,UInt32)
+select mapPopulateSeries(map(toUInt64(1), toUInt64(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(UInt64,UInt64)
+select mapPopulateSeries(map(toUInt128(1), toUInt128(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(UInt128,UInt128)
+select mapPopulateSeries(map(toUInt256(1), toUInt256(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(UInt256,UInt256)
+select mapPopulateSeries(map(toInt8(1), toInt8(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(Int16,Int16)
+select mapPopulateSeries(map(toInt16(1), toInt16(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(Int16,Int16)
+select mapPopulateSeries(map(toInt32(1), toInt32(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(Int32,Int32)
+select mapPopulateSeries(map(toInt64(1), toInt64(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(Int64,Int64)
+select mapPopulateSeries(map(toInt128(1), toInt128(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(Int128,Int128)
+select mapPopulateSeries(map(toInt256(1), toInt256(1), 2, 1)) as res, toTypeName(res);
+{1:1,2:1}	Map(Int256,Int256)
+select mapPopulateSeries(map(toInt8(-10), toInt8(1), 2, 1)) as res, toTypeName(res);
+{-10:1,-9:0,-8:0,-7:0,-6:0,-5:0,-4:0,-3:0,-2:0,-1:0,0:0,1:0,2:1}	Map(Int16,Int16)
+select mapPopulateSeries(map(toInt16(-10), toInt16(1), 2, 1)) as res, toTypeName(res);
+{-10:1,-9:0,-8:0,-7:0,-6:0,-5:0,-4:0,-3:0,-2:0,-1:0,0:0,1:0,2:1}	Map(Int16,Int16)
+select mapPopulateSeries(map(toInt32(-10), toInt32(1), 2, 1)) as res, toTypeName(res);
+{-10:1,-9:0,-8:0,-7:0,-6:0,-5:0,-4:0,-3:0,-2:0,-1:0,0:0,1:0,2:1}	Map(Int32,Int32)
+select mapPopulateSeries(map(toInt64(-10), toInt64(1), 2, 1)) as res, toTypeName(res);
+{-10:1,-9:0,-8:0,-7:0,-6:0,-5:0,-4:0,-3:0,-2:0,-1:0,0:0,1:0,2:1}	Map(Int64,Int64)
+select mapPopulateSeries(map(toInt64(-10), toInt64(1), 2, 1), toInt64(-5)) as res, toTypeName(res);
+{-10:1,-9:0,-8:0,-7:0,-6:0,-5:0}	Map(Int64,Int64)
+select mapPopulateSeries(); -- { serverError 42 }
+select mapPopulateSeries('asdf'); -- { serverError 43 }
+select mapPopulateSeries(map('1', 1, '2', 1)) as res, toTypeName(res); -- { serverError 43 }

tests/queries/0_stateless/01925_map_populate_series_on_map.sql

@@ -0,0 +1,36 @@
+-- { echo }
+drop table if exists map_test;
+set allow_experimental_map_type = 1;
+create table map_test engine=TinyLog() as (select (number + 1) as n, map(1, 1, number,2) as m from numbers(1, 5));
+
+select mapPopulateSeries(m) from map_test;
+select mapPopulateSeries(m, toUInt64(3)) from map_test;
+select mapPopulateSeries(m, toUInt64(10)) from map_test;
+select mapPopulateSeries(m, 1000) from map_test; -- { serverError 43 }
+select mapPopulateSeries(m, n) from map_test;
+
+drop table map_test;
+
+select mapPopulateSeries(map(toUInt8(1), toUInt8(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toUInt16(1), toUInt16(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toUInt32(1), toUInt32(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toUInt64(1), toUInt64(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toUInt128(1), toUInt128(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toUInt256(1), toUInt256(1), 2, 1)) as res, toTypeName(res);
+
+select mapPopulateSeries(map(toInt8(1), toInt8(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toInt16(1), toInt16(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toInt32(1), toInt32(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toInt64(1), toInt64(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toInt128(1), toInt128(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toInt256(1), toInt256(1), 2, 1)) as res, toTypeName(res);
+
+select mapPopulateSeries(map(toInt8(-10), toInt8(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toInt16(-10), toInt16(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toInt32(-10), toInt32(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toInt64(-10), toInt64(1), 2, 1)) as res, toTypeName(res);
+select mapPopulateSeries(map(toInt64(-10), toInt64(1), 2, 1), toInt64(-5)) as res, toTypeName(res);
+
+select mapPopulateSeries(); -- { serverError 42 }
+select mapPopulateSeries('asdf'); -- { serverError 43 }
+select mapPopulateSeries(map('1', 1, '2', 1)) as res, toTypeName(res); -- { serverError 43 }

tests/testflows/extended_precision_data_types/snapshots/common.py.tests.snapshot

@@ -1045,6 +1045,7 @@ a
 
 mapPopulateSeries_with_Int128_on_a_table = r"""
 a
+([1,2,3,4,5],[1,2,3,0,0])
 """
 
 mapContains_with_Int128_on_a_table = r"""
@@ -1575,6 +1576,7 @@ a
 
 mapPopulateSeries_with_Int256_on_a_table = r"""
 a
+([1,2,3,4,5],[1,2,3,0,0])
 """
 
 mapContains_with_Int256_on_a_table = r"""
@@ -2105,6 +2107,7 @@ a
 
 mapPopulateSeries_with_UInt128_on_a_table = r"""
 a
+([1,2,3,4,5],[1,2,3,0,0])
 """
 
 mapContains_with_UInt128_on_a_table = r"""
@@ -2635,6 +2638,7 @@ a
 
 mapPopulateSeries_with_UInt256_on_a_table = r"""
 a
+([1,2,3,4,5],[1,2,3,0,0])
 """
 
 mapContains_with_UInt256_on_a_table = r"""

tests/testflows/extended_precision_data_types/tests/array_tuple_map.py

@@ -393,9 +393,13 @@ def map_func(self, data_type, node=None):
         execute_query(f"SELECT * FROM {table_name} ORDER BY a ASC")
 
     with Scenario(f"mapPopulateSeries with {data_type}"):
-        node.query(f"SELECT mapPopulateSeries([1,2,3], [{to_data_type(data_type,1)},"
-            f"{to_data_type(data_type,2)}, {to_data_type(data_type,3)}], 5)",
-            exitcode = 44, message='Exception:')
+        sql = (f"SELECT mapPopulateSeries([1,2,3], [{to_data_type(data_type,1)},"
+            f"{to_data_type(data_type,2)}, {to_data_type(data_type,3)}], 5)")
+
+        exitcode, message = 0, None
+        if data_type.startswith("Decimal"):
+            exitcode, message = 44, "Exception:"
+        node.query(sql, exitcode=exitcode, message=message)
 
     with Scenario(f"mapPopulateSeries with {data_type} on a table"):
         table_name = get_table_name()
@@ -403,9 +407,13 @@ def map_func(self, data_type, node=None):
         table(name = table_name, data_type = f'Tuple(Array({data_type}), Array({data_type}))')
 
         with When("I insert the output into a table"):
-            node.query(f"INSERT INTO {table_name} SELECT mapPopulateSeries([1,2,3],"
-                f"[{to_data_type(data_type,1)}, {to_data_type(data_type,2)}, {to_data_type(data_type,3)}], 5)",
-                exitcode = 44, message='Exception:')
+            sql = (f"INSERT INTO {table_name} SELECT mapPopulateSeries([1,2,3],"
+                f"[{to_data_type(data_type,1)}, {to_data_type(data_type,2)}, {to_data_type(data_type,3)}], 5)")
+
+            exitcode, message = 0, None
+            if data_type.startswith("Decimal"):
+                exitcode, message = 44, "Exception:"
+            node.query(sql, exitcode=exitcode, message=message)
 
         execute_query(f"SELECT * FROM {table_name} ORDER BY a ASC")