From e94b81460674008ade2dddfb60d27f6c7f7c9dc1 Mon Sep 17 00:00:00 2001
From: Nikolai Kochetov <nik-kochetov@yandex-team.ru>
Date: Wed, 16 Sep 2020 16:07:32 +0300
Subject: [PATCH] Update arrayConcat

---
 src/Functions/GatherUtils/GatherUtils.h |  2 +-
 src/Functions/GatherUtils/concat.cpp    | 33 ++++++++++++++++++-------
 src/Functions/array/arrayConcat.cpp     |  7 ++----
 3 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/src/Functions/GatherUtils/GatherUtils.h b/src/Functions/GatherUtils/GatherUtils.h
index f383b9f4e7f..6c31898b202 100644
--- a/src/Functions/GatherUtils/GatherUtils.h
+++ b/src/Functions/GatherUtils/GatherUtils.h
@@ -41,7 +41,7 @@ std::unique_ptr<IArraySource> createArraySource(const ColumnArray & col, bool is
 std::unique_ptr<IValueSource> createValueSource(const IColumn & col, bool is_const, size_t total_rows);
 std::unique_ptr<IArraySink> createArraySink(ColumnArray & col, size_t column_size);
 
-void concat(const std::vector<std::unique_ptr<IArraySource>> & sources, IArraySink & sink);
+ColumnArray::MutablePtr concat(const std::vector<std::unique_ptr<IArraySource>> & sources);
 
 ColumnArray::MutablePtr sliceFromLeftConstantOffsetUnbounded(IArraySource & src, size_t offset);
 ColumnArray::MutablePtr sliceFromLeftConstantOffsetBounded(IArraySource & src, size_t offset, ssize_t length);
diff --git a/src/Functions/GatherUtils/concat.cpp b/src/Functions/GatherUtils/concat.cpp
index 0962baed603..e116052c637 100644
--- a/src/Functions/GatherUtils/concat.cpp
+++ b/src/Functions/GatherUtils/concat.cpp
@@ -16,37 +16,52 @@ namespace ErrorCodes
 namespace GatherUtils
 {
 
-struct ArrayConcat : public ArraySinkSourceSelector<ArrayConcat>
+struct ArrayConcat : public ArraySinkSelector<ArrayConcat>
 {
     using Sources = std::vector<std::unique_ptr<IArraySource>>;
 
-    template <typename Source, typename Sink>
-    static void selectSourceSink(Source &&, Sink && sink, const Sources & sources)
+    template <typename Source>
+    static void selectImpl(Source &&, const Sources & sources, ColumnArray::MutablePtr & result)
     {
         using SourceType = typename std::decay<Source>::type;
+        using Sink = typename SourceType::SinkType;
+        result = ColumnArray::create(source.createValuesColumn());
+        Sink sink(result->getData(), result->getOffsets(), source.getColumnSize());
+
         concat<SourceType, Sink>(sources, sink);
     }
 
-    template <typename Source, typename Sink>
-    static void selectSourceSink(ConstSource<Source> &&, Sink && sink, const Sources & sources)
+    template <typename Source>
+    static void selectImpl(ConstSource<Source> &&, const Sources & sources, ColumnArray::MutablePtr & result)
     {
         using SourceType = typename std::decay<Source>::type;
+        using Sink = typename SourceType::SinkType;
+        result = ColumnArray::create(source.createValuesColumn());
+        Sink sink(result->getData(), result->getOffsets(), source.getColumnSize());
+
         concat<SourceType, Sink>(sources, sink);
     }
 
-    template <typename Source, typename Sink>
-    static void selectSourceSink(ConstSource<Source> &, Sink && sink, const Sources & sources)
+    template <typename Source>
+    static void selectImpl(ConstSource<Source> &, const Sources & sources, ColumnArray::MutablePtr & result)
     {
         using SourceType = typename std::decay<Source>::type;
+        using Sink = typename SourceType::SinkType;
+        result = ColumnArray::create(source.createValuesColumn());
+        Sink sink(result->getData(), result->getOffsets(), source.getColumnSize());
+
         concat<SourceType, Sink>(sources, sink);
     }
 };
 
-void concat(const std::vector<std::unique_ptr<IArraySource>> & sources, IArraySink & sink)
+ColumnArray::MutablePtr concat(const std::vector<std::unique_ptr<IArraySource>> & sources)
 {
     if (sources.empty())
         throw Exception("Concat function should get at least 1 ArraySource", ErrorCodes::LOGICAL_ERROR);
-    return ArrayConcat::select(*sources.front(), sink, sources);
+
+    ColumnArray::MutablePtr res;
+    ArrayConcat::select(*sources.front(), sources, res);
+    return res;
 }
 
 }
diff --git a/src/Functions/array/arrayConcat.cpp b/src/Functions/array/arrayConcat.cpp
index 23e6c7a28a8..bea797c1979 100644
--- a/src/Functions/array/arrayConcat.cpp
+++ b/src/Functions/array/arrayConcat.cpp
@@ -59,8 +59,6 @@ public:
             return;
         }
 
-        auto result_column = return_type->createColumn();
-
         size_t rows = input_rows_count;
         size_t num_args = arguments.size();
 
@@ -95,10 +93,9 @@ public:
                 throw Exception{"Arguments for function " + getName() + " must be arrays.", ErrorCodes::LOGICAL_ERROR};
         }
 
-        auto sink = GatherUtils::createArraySink(typeid_cast<ColumnArray &>(*result_column), rows);
-        GatherUtils::concat(sources, *sink);
+        auto sink = GatherUtils::concat(sources);
 
-        block.getByPosition(result).column = std::move(result_column);
+        block.getByPosition(result).column = std::move(sink);
     }
 
     bool useDefaultImplementationForConstants() const override { return true; }