add arrayZip function

dbms/src/Functions/array/arrayZipWithIndex.cpp

@@ -0,0 +1,103 @@
+#include <Columns/ColumnTuple.h>
+#include <Columns/ColumnArray.h>
+#include <DataTypes/DataTypeTuple.h>
+#include <DataTypes/DataTypeArray.h>
+#include <Functions/FunctionFactory.h>
+#include <Functions/FunctionHelpers.h>
+#include <IO/WriteHelpers.h>
+
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int SIZES_OF_ARRAYS_DOESNT_MATCH;
+    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
+}
+
+/// arrayZip(['a', 'b', 'c'], ['d', 'e', 'f']) = [('a', 'd'), ('b', 'e'), ('c', 'f')]
+class FunctionArrayZip : public IFunction
+{
+public:
+    static constexpr auto name = "arrayZip";
+    static FunctionPtr create(const Context &) { return std::make_shared<FunctionArrayZip>(); }
+
+    String getName() const override
+    {
+        return name;
+    }
+
+    bool isVariadic() const override { return true; }
+    size_t getNumberOfArguments() const override { return 0; }
+    bool useDefaultImplementationForConstants() const override { return true; }
+
+    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    {
+        if (arguments.size() < 1)
+            throw Exception("Function " + getName() + " needs at least one argument; passed " + toString(arguments.size()) + "."
+                , ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+
+        DataTypes arguments_types;
+        for (size_t index = 0; index < arguments.size(); ++index)
+        {
+            const DataTypeArray * array_type = checkAndGetDataType<DataTypeArray>(arguments[index].type.get());
+
+            if (!array_type)
+                throw Exception(
+                    "Argument " + toString(index + 1) + " of function must be array. Found " + arguments[0].type->getName() + " instead.",
+                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+
+            arguments_types.emplace_back(array_type->getNestedType());
+        }
+
+        return std::make_shared<DataTypeArray>(std::make_shared<DataTypeTuple>(arguments_types));
+    }
+
+    void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t /*input_rows_count*/) override
+    {
+        auto first_argument = block.getByPosition(arguments[0]);
+        const auto & first_array_column = checkAndGetColumn<ColumnArray>(first_argument.column.get());
+
+        for (size_t index = 1; index < arguments.size(); ++index)
+        {
+            const auto & argument_type_and_column = block.getByPosition(arguments[index]);
+            const auto & argument_array_column = checkAndGetColumn<ColumnArray>(argument_type_and_column.column.get());
+
+            if (argument_array_column->size() != first_array_column->size())
+                throw Exception("The argument 1 and argument " + toString(index + 1) + " of function have different array sizes",
+                                ErrorCodes::SIZES_OF_ARRAYS_DOESNT_MATCH);
+
+            if (argument_array_column->getData().size() != first_array_column->getData().size())
+                throw Exception("The argument 1 and argument " + toString(index + 1) + " of function have different array sizes",
+                                ErrorCodes::SIZES_OF_ARRAYS_DOESNT_MATCH);
+        }
+
+        Columns res_tuple_columns(arguments.size());
+        res_tuple_columns[0] = first_array_column->getDataPtr();
+        for (size_t argument_index = 1; argument_index < arguments.size(); ++argument_index)
+        {
+            const auto & argument_type_and_column = block.getByPosition(arguments[argument_index]);
+            const auto & argument_array_column = checkAndGetColumn<ColumnArray>(argument_type_and_column.column.get());
+
+            for (size_t row_index = 0; row_index < first_array_column->size(); ++row_index)
+                if (first_array_column->getOffsets()[row_index] != argument_array_column->getOffsets()[row_index])
+                    throw Exception(
+                        "The argument 1 and argument " + toString(argument_index + 1) + " of function have different array sizes",
+                        ErrorCodes::SIZES_OF_ARRAYS_DOESNT_MATCH);
+
+            res_tuple_columns[argument_index] = argument_array_column->getDataPtr();
+        }
+
+        block.getByPosition(result).column = ColumnArray::create(
+            ColumnTuple::create(res_tuple_columns), first_array_column->getOffsetsPtr());
+    }
+};
+
+void registerFunctionArrayZip(FunctionFactory & factory)
+{
+    factory.registerFunction<FunctionArrayZip>();
+}
+
+}
+

dbms/src/Functions/array/registerFunctionsArray.cpp

@@ -32,6 +32,7 @@ void registerFunctionArrayUniq(FunctionFactory &);
 void registerFunctionArrayDistinct(FunctionFactory &);
 void registerFunctionArrayFlatten(FunctionFactory &);
 void registerFunctionArrayWithConstant(FunctionFactory &);
+void registerFunctionArrayZip(FunctionFactory &);
 
 void registerFunctionsArray(FunctionFactory & factory)
 {
@@ -64,6 +65,7 @@ void registerFunctionsArray(FunctionFactory & factory)
     registerFunctionArrayDistinct(factory);
     registerFunctionArrayFlatten(factory);
     registerFunctionArrayWithConstant(factory);
+    registerFunctionArrayZip(factory);
 }
 
 }

dbms/tests/queries/0_stateless/01045_array_zip.reference

@@ -0,0 +1,2 @@
+[('a','d'),('b','e'),('c','f')]
+[('a','d','g'),('b','e','h'),('c','f','i')]

dbms/tests/queries/0_stateless/01045_array_zip.sql

@@ -0,0 +1,9 @@
+SELECT arrayZip(['a', 'b', 'c'], ['d', 'e', 'f']);
+
+SELECT arrayZip(['a', 'b', 'c'], ['d', 'e', 'f'], ['g', 'h', 'i']);
+
+SELECT arrayZip(); -- { serverError 42 }
+
+SELECT arrayZip('a', 'b', 'c'); -- { serverError 43 }
+
+SELECT arrayZip(['a', 'b', 'c'], ['d', 'e', 'f', 'd']); -- { serverError 190 }

docs/en/query_language/functions/array_functions.md

@@ -806,7 +806,7 @@ Function:
 
 The flattened array contains all the elements from all source arrays.
 
-**Syntax** 
+**Syntax**
 
 ```sql
 flatten(array_of_arrays)
@@ -866,4 +866,38 @@ Result:
 └────────────────────────────────────────────┘
 ```
 
+## arrayZip {#arrayzip}
+
+Combine multiple Array type columns into one Array[Tuple(...)] column
+
+**Syntax**
+
+```sql
+arrayZip(arr1, arr2, ..., arrN)
+```
+
+**Parameters**
+
+`arr` — Any number of [array](../../data_types/array.md) type columns to combine.
+
+**Returned value**
+
+The result of Array[Tuple(...)] type after the combination of these arrays
+
+**Example**
+
+Query:
+
+```sql
+SELECT arrayZip(['a', 'b', 'c'], ['d', 'e', 'f']);
+```
+
+Result:
+
+```text
+┌─arrayZipWithIndex(['a', 'b', 'c'], ['d', 'e', 'f'])─┐
+│ [('a','d'),('b','e'),('c','f')]                     │
+└─────────────────────────────────────────────────────┘
+```
+
 [Original article](https://clickhouse.yandex/docs/en/query_language/functions/array_functions/) <!--hide-->