Merge pull request #67707 from bigo-sg/opt_array_map

Optimize function `array()` and `map()`

src/Functions/array/array.cpp

@@ -1,11 +1,15 @@
-#include <Functions/IFunction.h>
+#include <Columns/ColumnArray.h>
-#include <Functions/FunctionFactory.h>
+#include <Columns/ColumnDecimal.h>
+#include <Columns/ColumnFixedString.h>
+#include <Columns/ColumnNullable.h>
+#include <Columns/ColumnString.h>
+#include <Core/Settings.h>
 #include <DataTypes/DataTypeArray.h>
 #include <DataTypes/getLeastSupertype.h>
-#include <Columns/ColumnArray.h>
+#include <Functions/FunctionFactory.h>
-#include <Core/Settings.h>
+#include <Functions/IFunction.h>
-#include <Interpreters/castColumn.h>
 #include <Interpreters/Context.h>
+#include <Interpreters/castColumn.h>
 
 
 namespace DB
@@ -44,11 +48,13 @@ public:
 
     ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override
     {
-        size_t num_elements = arguments.size();
+        const size_t num_elements = arguments.size();
 
         if (num_elements == 0)
+        {
             /// We should return constant empty array.
             return result_type->createColumnConstWithDefaultValue(input_rows_count);
+        }
 
         const DataTypePtr & elem_type = static_cast<const DataTypeArray &>(*result_type).getNestedType();
 
@@ -60,7 +66,6 @@ public:
 
         Columns columns_holder(num_elements);
         ColumnRawPtrs column_ptrs(num_elements);
-
         for (size_t i = 0; i < num_elements; ++i)
         {
             const auto & arg = arguments[i];
@@ -77,35 +82,199 @@ public:
         }
 
         /// Create and fill the result array.
-
         auto out = ColumnArray::create(elem_type->createColumn());
         IColumn & out_data = out->getData();
         IColumn::Offsets & out_offsets = out->getOffsets();
 
-        out_data.reserve(input_rows_count * num_elements);
+        /// Fill out_offsets
-        out_offsets.resize(input_rows_count);
+        out_offsets.resize_exact(input_rows_count);
-
         IColumn::Offset current_offset = 0;
         for (size_t i = 0; i < input_rows_count; ++i)
         {
-            for (size_t j = 0; j < num_elements; ++j)
-                out_data.insertFrom(*column_ptrs[j], i);
-
             current_offset += num_elements;
             out_offsets[i] = current_offset;
         }
 
+        /// Fill out_data
+        out_data.reserve(input_rows_count * num_elements);
+        if (num_elements == 1)
+            out_data.insertRangeFrom(*column_ptrs[0], 0, input_rows_count);
+        else
+            execute(column_ptrs, out_data, input_rows_count);
         return out;
     }
 
 private:
+    bool execute(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
+    {
+        return executeNumber<UInt8>(columns, out_data, input_rows_count) || executeNumber<UInt16>(columns, out_data, input_rows_count)
+            || executeNumber<UInt32>(columns, out_data, input_rows_count) || executeNumber<UInt64>(columns, out_data, input_rows_count)
+            || executeNumber<UInt128>(columns, out_data, input_rows_count) || executeNumber<UInt256>(columns, out_data, input_rows_count)
+            || executeNumber<Int8>(columns, out_data, input_rows_count) || executeNumber<Int16>(columns, out_data, input_rows_count)
+            || executeNumber<Int32>(columns, out_data, input_rows_count) || executeNumber<Int64>(columns, out_data, input_rows_count)
+            || executeNumber<Int128>(columns, out_data, input_rows_count) || executeNumber<Int256>(columns, out_data, input_rows_count)
+            || executeNumber<Float32>(columns, out_data, input_rows_count) || executeNumber<Float64>(columns, out_data, input_rows_count)
+            || executeNumber<Decimal32>(columns, out_data, input_rows_count)
+            || executeNumber<Decimal64>(columns, out_data, input_rows_count)
+            || executeNumber<Decimal128>(columns, out_data, input_rows_count)
+            || executeNumber<Decimal256>(columns, out_data, input_rows_count)
+            || executeNumber<DateTime64>(columns, out_data, input_rows_count) || executeString(columns, out_data, input_rows_count)
+            || executeNullable(columns, out_data, input_rows_count) || executeTuple(columns, out_data, input_rows_count)
+            || executeFixedString(columns, out_data, input_rows_count) || executeGeneric(columns, out_data, input_rows_count);
+    }
+
+    template <typename T>
+    bool executeNumber(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
+    {
+        using Container = ColumnVectorOrDecimal<T>::Container;
+        std::vector<const Container *> containers(columns.size(), nullptr);
+        for (size_t i = 0; i < columns.size(); ++i)
+        {
+            const ColumnVectorOrDecimal<T> * concrete_column = checkAndGetColumn<ColumnVectorOrDecimal<T>>(columns[i]);
+            if (!concrete_column)
+                return false;
+
+            containers[i] = &concrete_column->getData();
+        }
+
+        ColumnVectorOrDecimal<T> & concrete_out_data = assert_cast<ColumnVectorOrDecimal<T> &>(out_data);
+        Container & out_container = concrete_out_data.getData();
+        out_container.resize_exact(columns.size() * input_rows_count);
+
+        for (size_t row_i = 0; row_i < input_rows_count; ++row_i)
+        {
+            const size_t base = row_i * columns.size();
+            for (size_t col_i = 0; col_i < columns.size(); ++col_i)
+                out_container[base + col_i] = (*containers[col_i])[row_i];
+        }
+        return true;
+    }
+
+    bool executeString(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
+    {
+        size_t total_bytes = 0;
+        std::vector<const ColumnString *> concrete_columns(columns.size(), nullptr);
+        for (size_t i = 0; i < columns.size(); ++i)
+        {
+            const ColumnString * concrete_column = checkAndGetColumn<ColumnString>(columns[i]);
+            if (!concrete_column)
+                return false;
+
+            total_bytes += concrete_column->getChars().size();
+            concrete_columns[i] = concrete_column;
+        }
+
+        ColumnString & concrete_out_data = assert_cast<ColumnString &>(out_data);
+        auto & out_chars = concrete_out_data.getChars();
+        auto & out_offsets = concrete_out_data.getOffsets();
+        out_chars.resize_exact(total_bytes);
+        out_offsets.resize_exact(input_rows_count * columns.size());
+
+        size_t cur_out_offset = 0;
+        for (size_t row_i = 0; row_i < input_rows_count; ++row_i)
+        {
+            const size_t base = row_i * columns.size();
+            for (size_t col_i = 0; col_i < columns.size(); ++col_i)
+            {
+                StringRef ref = concrete_columns[col_i]->getDataAt(row_i);
+                memcpySmallAllowReadWriteOverflow15(&out_chars[cur_out_offset], ref.data, ref.size);
+                out_chars[cur_out_offset + ref.size] = 0;
+
+                cur_out_offset += ref.size + 1;
+                out_offsets[base + col_i] = cur_out_offset;
+            }
+        }
+        return true;
+    }
+
+    bool executeFixedString(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
+    {
+        std::vector<const ColumnFixedString *> concrete_columns(columns.size(), nullptr);
+        for (size_t i = 0; i < columns.size(); ++i)
+        {
+            const ColumnFixedString * concrete_column = checkAndGetColumn<ColumnFixedString>(columns[i]);
+            if (!concrete_column)
+                return false;
+
+            concrete_columns[i] = concrete_column;
+        }
+
+        ColumnFixedString & concrete_out_data = assert_cast<ColumnFixedString &>(out_data);
+        auto & out_chars = concrete_out_data.getChars();
+
+        const size_t n = concrete_out_data.getN();
+        size_t total_bytes = n * columns.size() * input_rows_count;
+        out_chars.resize_exact(total_bytes);
+
+        size_t curr_out_offset = 0;
+        for (size_t row_i = 0; row_i < input_rows_count; ++row_i)
+        {
+            for (size_t col_i = 0; col_i < columns.size(); ++col_i)
+            {
+                StringRef ref = concrete_columns[col_i]->getDataAt(row_i);
+                memcpySmallAllowReadWriteOverflow15(&out_chars[curr_out_offset], ref.data, n);
+                curr_out_offset += n;
+            }
+        }
+        return true;
+    }
+
+    bool executeNullable(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
+    {
+        ColumnRawPtrs null_maps(columns.size(), nullptr);
+        ColumnRawPtrs nested_columns(columns.size(), nullptr);
+        for (size_t i = 0; i < columns.size(); ++i)
+        {
+            const ColumnNullable * concrete_column = checkAndGetColumn<ColumnNullable>(columns[i]);
+            if (!concrete_column)
+                return false;
+
+            null_maps[i] = &concrete_column->getNullMapColumn();
+            nested_columns[i] = &concrete_column->getNestedColumn();
+        }
+
+        ColumnNullable & concrete_out_data = assert_cast<ColumnNullable &>(out_data);
+        auto & out_null_map = concrete_out_data.getNullMapColumn();
+        auto & out_nested_column = concrete_out_data.getNestedColumn();
+        execute(null_maps, out_null_map, input_rows_count);
+        execute(nested_columns, out_nested_column, input_rows_count);
+        return true;
+    }
+
+    bool executeTuple(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
+    {
+        ColumnTuple * concrete_out_data = typeid_cast<ColumnTuple *>(&out_data);
+        if (!concrete_out_data)
+            return false;
+
+        const size_t tuple_size = concrete_out_data->tupleSize();
+        for (size_t i = 0; i < tuple_size; ++i)
+        {
+            ColumnRawPtrs elem_columns(columns.size(), nullptr);
+            for (size_t j = 0; j < columns.size(); ++j)
+            {
+                const ColumnTuple * concrete_column = assert_cast<const ColumnTuple *>(columns[j]);
+                elem_columns[j] = &concrete_column->getColumn(i);
+            }
+            execute(elem_columns, concrete_out_data->getColumn(i), input_rows_count);
+        }
+        return true;
+    }
+
+    bool executeGeneric(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
+    {
+        for (size_t i = 0; i < input_rows_count; ++i)
+            for (const auto * column : columns)
+                out_data.insertFrom(*column, i);
+        return true;
+    }
+
+
     String getName() const override
     {
         return name;
     }
 
-    bool addField(DataTypePtr type_res, const Field & f, Array & arr) const;
-
     bool use_variant_as_common_type = false;
 };
 

src/Functions/map.cpp

@@ -2,6 +2,8 @@
 #include <Columns/ColumnMap.h>
 #include <Columns/ColumnNullable.h>
 #include <Columns/ColumnsCommon.h>
+#include <Columns/ColumnsNumber.h>
+#include <Core/Settings.h>
 #include <DataTypes/DataTypeArray.h>
 #include <DataTypes/DataTypeMap.h>
 #include <DataTypes/DataTypeTuple.h>
@@ -13,7 +15,6 @@
 #include <Interpreters/Context.h>
 #include <Interpreters/castColumn.h>
 #include <Common/HashTable/HashSet.h>
-#include <Core/Settings.h>
 
 
 namespace DB
@@ -36,11 +37,18 @@ class FunctionMap : public IFunction
 public:
     static constexpr auto name = "map";
 
-    explicit FunctionMap(bool use_variant_as_common_type_) : use_variant_as_common_type(use_variant_as_common_type_) {}
+    explicit FunctionMap(ContextPtr context_)
+        : context(context_)
+        , use_variant_as_common_type(
+              context->getSettingsRef().allow_experimental_variant_type && context->getSettingsRef().use_variant_as_common_type)
+        , function_array(FunctionFactory::instance().get("array", context))
+        , function_map_from_arrays(FunctionFactory::instance().get("mapFromArrays", context))
+    {
+    }
 
     static FunctionPtr create(ContextPtr context)
     {
-        return std::make_shared<FunctionMap>(context->getSettingsRef().allow_experimental_variant_type && context->getSettingsRef().use_variant_as_common_type);
+        return std::make_shared<FunctionMap>(context);
     }
 
     String getName() const override
@@ -101,62 +109,38 @@ public:
     ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override
     {
         size_t num_elements = arguments.size();
-
         if (num_elements == 0)
             return result_type->createColumnConstWithDefaultValue(input_rows_count);
 
+        ColumnsWithTypeAndName key_args;
+        ColumnsWithTypeAndName value_args;
+        for (size_t i = 0; i < num_elements; i += 2)
+        {
+            key_args.emplace_back(arguments[i]);
+            value_args.emplace_back(arguments[i+1]);
+        }
+
         const auto & result_type_map = static_cast<const DataTypeMap &>(*result_type);
         const DataTypePtr & key_type = result_type_map.getKeyType();
         const DataTypePtr & value_type = result_type_map.getValueType();
+        const DataTypePtr & key_array_type = std::make_shared<DataTypeArray>(key_type);
+        const DataTypePtr & value_array_type = std::make_shared<DataTypeArray>(value_type);
 
-        Columns columns_holder(num_elements);
+        /// key_array = array(args[0], args[2]...)
-        ColumnRawPtrs column_ptrs(num_elements);
+        ColumnPtr key_array = function_array->build(key_args)->execute(key_args, key_array_type, input_rows_count);
+        /// value_array = array(args[1], args[3]...)
+        ColumnPtr value_array = function_array->build(value_args)->execute(value_args, value_array_type, input_rows_count);
 
-        for (size_t i = 0; i < num_elements; ++i)
+        /// result = mapFromArrays(key_array, value_array)
-        {
+        ColumnsWithTypeAndName map_args{{key_array, key_array_type, ""}, {value_array, value_array_type, ""}};
-            const auto & arg = arguments[i];
+        return function_map_from_arrays->build(map_args)->execute(map_args, result_type, input_rows_count);
-            const auto to_type = i % 2 == 0 ? key_type : value_type;
-
-            ColumnPtr preprocessed_column = castColumn(arg, to_type);
-            preprocessed_column = preprocessed_column->convertToFullColumnIfConst();
-
-            columns_holder[i] = std::move(preprocessed_column);
-            column_ptrs[i] = columns_holder[i].get();
-        }
-
-        /// Create and fill the result map.
-
-        MutableColumnPtr keys_data = key_type->createColumn();
-        MutableColumnPtr values_data = value_type->createColumn();
-        MutableColumnPtr offsets = DataTypeNumber<IColumn::Offset>().createColumn();
-
-        size_t total_elements = input_rows_count * num_elements / 2;
-        keys_data->reserve(total_elements);
-        values_data->reserve(total_elements);
-        offsets->reserve(input_rows_count);
-
-        IColumn::Offset current_offset = 0;
-        for (size_t i = 0; i < input_rows_count; ++i)
-        {
-            for (size_t j = 0; j < num_elements; j += 2)
-            {
-                keys_data->insertFrom(*column_ptrs[j], i);
-                values_data->insertFrom(*column_ptrs[j + 1], i);
-            }
-
-            current_offset += num_elements / 2;
-            offsets->insert(current_offset);
-        }
-
-        auto nested_column = ColumnArray::create(
-            ColumnTuple::create(Columns{std::move(keys_data), std::move(values_data)}),
-            std::move(offsets));
-
-        return ColumnMap::create(nested_column);
     }
 
 private:
+    ContextPtr context;
     bool use_variant_as_common_type = false;
+    FunctionOverloadResolverPtr function_array;
+    FunctionOverloadResolverPtr function_map_from_arrays;
 };
 
 /// mapFromArrays(keys, values) is a function that allows you to make key-value pair from a pair of arrays or maps
@@ -173,6 +157,7 @@ public:
     bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
     bool useDefaultImplementationForNulls() const override { return false; }
     bool useDefaultImplementationForConstants() const override { return true; }
+    bool useDefaultImplementationForLowCardinalityColumns() const override { return false; }
 
     DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
     {