From 1b9620001b31256ffb6e8fee7b521efa8688fbee Mon Sep 17 00:00:00 2001
From: Vitaly Baranov <vitlibar@clickhouse.com>
Date: Thu, 8 Feb 2024 15:24:01 +0100
Subject: [PATCH] Move ColumnArray-specific code out of the ActionsDAG.

---
 src/Columns/ColumnArray.cpp     | 15 +++++++++++++++
 src/Columns/ColumnArray.h       |  4 ++++
 src/Interpreters/ActionsDAG.cpp | 17 ++++++++++++-----
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/Columns/ColumnArray.cpp b/src/Columns/ColumnArray.cpp
index fddfa2ac6b2..6f60ec0e642 100644
--- a/src/Columns/ColumnArray.cpp
+++ b/src/Columns/ColumnArray.cpp
@@ -554,6 +554,21 @@ void ColumnArray::insertRangeFrom(const IColumn & src, size_t start, size_t leng
 }
 
 
+MutableColumnPtr ColumnArray::getDataInRange(size_t start, size_t length) const
+{
+    if (start + length > getOffsets().size())
+        throw Exception(ErrorCodes::PARAMETER_OUT_OF_BOUND, "Parameter out of bound in ColumnArray::getDataPtrForRange method. "
+            "[start({}) + length({}) > offsets.size({})]", start, length, getOffsets().size());
+
+    size_t start_offset = offsetAt(start);
+    size_t end_offset = offsetAt(start + length);
+
+    auto res = getData().cloneEmpty();
+    res->insertRangeFrom(getData(), start_offset, end_offset - start_offset);
+    return res;
+}
+
+
 ColumnPtr ColumnArray::filter(const Filter & filt, ssize_t result_size_hint) const
 {
     if (typeid_cast<const ColumnUInt8 *>(data.get()))
diff --git a/src/Columns/ColumnArray.h b/src/Columns/ColumnArray.h
index 407f44a6f3c..64c8801b0c3 100644
--- a/src/Columns/ColumnArray.h
+++ b/src/Columns/ColumnArray.h
@@ -143,6 +143,10 @@ public:
     const ColumnPtr & getOffsetsPtr() const { return offsets; }
     ColumnPtr & getOffsetsPtr() { return offsets; }
 
+    /// Returns a copy of the data column's part corresponding to a specified range of rows.
+    /// For example, `getDataInRange(0, size())` is the same as `getDataPtr()->clone()`.
+    MutableColumnPtr getDataInRange(size_t start, size_t length) const;
+
     MutableColumns scatter(ColumnIndex num_columns, const Selector & selector) const override
     {
         return scatterImpl<ColumnArray>(num_columns, selector);
diff --git a/src/Interpreters/ActionsDAG.cpp b/src/Interpreters/ActionsDAG.cpp
index a1858916ca7..b3f3f8da26d 100644
--- a/src/Interpreters/ActionsDAG.cpp
+++ b/src/Interpreters/ActionsDAG.cpp
@@ -628,17 +628,24 @@ static ColumnWithTypeAndName executeActionForPartialResult(const ActionsDAG::Nod
             if (!array)
                 throw Exception(ErrorCodes::TYPE_MISMATCH,
                                 "ARRAY JOIN of not array nor map: {}", node->result_name);
-            res_column.column = array->getDataPtr();
+
+            ColumnPtr data;
             if (input_rows_count < array->size())
-                res_column.column = res_column.column->cloneResized(array->getOffsets()[input_rows_count - 1]);
+                data = array->getDataInRange(0, input_rows_count);
+            else
+                data = array->getDataPtr();
+
+            res_column.column = data;
             break;
         }
 
         case ActionsDAG::ActionType::COLUMN:
         {
-            res_column.column = node->column;
-            if (input_rows_count < res_column.column->size())
-                res_column.column = res_column.column->cloneResized(input_rows_count);
+            auto column = node->column;
+            if (input_rows_count < column->size())
+                column = column->cloneResized(input_rows_count);
+
+            res_column.column = column;
             break;
         }