insert range selective for shuffle write

liuneng1994 · kyligence-git · commit aa13a91cdd67 · 2023-12-18T23:03:07.000Z
diff --git a/src/Columns/ColumnAggregateFunction.cpp b/src/Columns/ColumnAggregateFunction.cpp
@@ -524,6 +524,35 @@ void ColumnAggregateFunction::insertDefault()
     pushBackAndCreateState(data, arena, func.get());
 }
 
+void ColumnAggregateFunction::insertRangeSelective(
+    const IColumn & from, const IColumn::Selector & selector, size_t selector_start, size_t length)
+{
+    const ColumnAggregateFunction & from_concrete = static_cast<const ColumnAggregateFunction &>(from);
+    const auto & from_data = from_concrete.data;
+    if (!empty() && src.get() != &from_concrete)
+    {
+        ensureOwnership();
+        Arena & arena = createOrGetArena();
+        Arena * arena_ptr = &arena;
+        data.reserve(size() + length);
+        for (size_t i = 0; i < length; ++i)
+        {
+            pushBackAndCreateState(data, arena, func.get());
+            func->merge(data.back(), from_data[selector[selector_start + i]], arena_ptr);
+        }
+        return;
+    }
+    /// Keep shared ownership of aggregation states.
+    src = from_concrete.getPtr();
+
+    size_t old_size = data.size();
+    data.resize(old_size + length);
+    auto * data_start = data.data();
+    size_t element_size = sizeof(data[0]);
+    for (size_t i = 0; i < length; ++i)
+        memcpy(data_start + old_size + i, &from_concrete.data[selector[selector_start + i]], element_size);
+}
+
 StringRef ColumnAggregateFunction::serializeValueIntoArena(size_t n, Arena & arena, const char *& begin, const UInt8 *) const
 {
     WriteBufferFromArena out(arena, begin);
diff --git a/src/Columns/ColumnAggregateFunction.h b/src/Columns/ColumnAggregateFunction.h
@@ -184,6 +184,8 @@ class ColumnAggregateFunction final : public COWHelper<IColumn, ColumnAggregateF
 
     void insertRangeFrom(const IColumn & from, size_t start, size_t length) override;
 
+    void insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length) override;
+
     void popBack(size_t n) override;
 
     ColumnPtr filter(const Filter & filter, ssize_t result_size_hint) const override;
diff --git a/src/Columns/ColumnArray.cpp b/src/Columns/ColumnArray.cpp
@@ -547,6 +547,29 @@ void ColumnArray::insertRangeFrom(const IColumn & src, size_t start, size_t leng
     }
 }
 
+void ColumnArray::insertRangeSelective(const IColumn & src, const Selector & selector, size_t selector_start, size_t length){
+
+    const ColumnArray & src_concrete = static_cast<const ColumnArray &>(src);
+    const Offsets & src_offsets = src_concrete.getOffsets();
+    const IColumn & src_data = src_concrete.getData();
+    IColumn & cur_data = getData();
+    Offsets & cur_offsets = getOffsets();
+
+    size_t old_size = cur_offsets.size();
+    size_t cur_size = old_size + length;
+    cur_data.reserve(cur_size);
+    cur_offsets.resize(cur_size);
+
+    for (size_t i = 0; i < length; ++i)
+    {
+        size_t src_pos = selector[selector_start + i];
+        size_t offset = src_offsets[src_pos - 1];
+        size_t size = src_offsets[src_pos] - offset;
+        cur_data.insertRangeFrom(src_data, offset, size);
+        cur_offsets[old_size + i] = cur_offsets[old_size + i - 1] + size; // PaddedPODArray allows to use -1th element that will have value 0
+    }
+}
+
 
 ColumnPtr ColumnArray::filter(const Filter & filt, ssize_t result_size_hint) const
 {
diff --git a/src/Columns/ColumnArray.h b/src/Columns/ColumnArray.h
@@ -84,6 +84,7 @@ class ColumnArray final : public COWHelper<IColumn, ColumnArray>
     void updateWeakHash32(WeakHash32 & hash) const override;
     void updateHashFast(SipHash & hash) const override;
     void insertRangeFrom(const IColumn & src, size_t start, size_t length) override;
+    void insertRangeSelective(const IColumn & src, const Selector & selector, size_t selector_start, size_t length) override;
     void insert(const Field & x) override;
     void insertFrom(const IColumn & src_, size_t n) override;
     void insertDefault() override;
@@ -149,6 +150,8 @@ class ColumnArray final : public COWHelper<IColumn, ColumnArray>
 
     void gather(ColumnGathererStream & gatherer_stream) override;
 
+    bool canBeInsideNullable() const override { return true; }
+
     ColumnPtr compress() const override;
 
     void forEachSubcolumn(MutableColumnCallback callback) override
diff --git a/src/Columns/ColumnCompressed.h b/src/Columns/ColumnCompressed.h
@@ -85,6 +85,7 @@ class ColumnCompressed : public COWHelper<IColumn, ColumnCompressed>
     bool isDefaultAt(size_t) const override { throwMustBeDecompressed(); }
     void insert(const Field &) override { throwMustBeDecompressed(); }
     void insertRangeFrom(const IColumn &, size_t, size_t) override { throwMustBeDecompressed(); }
+    void insertRangeSelective(const IColumn &, const Selector &, size_t, size_t) override { throwMustBeDecompressed(); }
     void insertData(const char *, size_t) override { throwMustBeDecompressed(); }
     void insertDefault() override { throwMustBeDecompressed(); }
     void popBack(size_t) override { throwMustBeDecompressed(); }
diff --git a/src/Columns/ColumnConst.h b/src/Columns/ColumnConst.h
@@ -126,6 +126,11 @@ class ColumnConst final : public COWHelper<IColumn, ColumnConst>
         s += length;
     }
 
+    void insertRangeSelective(const IColumn & /*src*/, const Selector & /*selector*/, size_t /*selector_start*/, size_t length) override
+    {
+        s += length;
+    }
+
     void insert(const Field &) override
     {
         ++s;
diff --git a/src/Columns/ColumnDecimal.cpp b/src/Columns/ColumnDecimal.cpp
@@ -359,6 +359,18 @@ void ColumnDecimal<T>::insertRangeFrom(const IColumn & src, size_t start, size_t
     memcpy(data.data() + old_size, &src_vec.data[start], length * sizeof(data[0]));
 }
 
+template <is_decimal T>
+void ColumnDecimal<T>::insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length)
+{
+    size_t old_size = data.size();
+    data.resize(old_size + length);
+    const auto & src_data = (static_cast<const Self &>(src)).getData();
+    for (size_t i = 0; i < length; ++i)
+    {
+        data[old_size + i] = src_data[selector[selector_start + i]];
+    }
+}
+
 template <is_decimal T>
 ColumnPtr ColumnDecimal<T>::filter(const IColumn::Filter & filt, ssize_t result_size_hint) const
 {
diff --git a/src/Columns/ColumnDecimal.h b/src/Columns/ColumnDecimal.h
@@ -62,6 +62,7 @@ class ColumnDecimal final : public COWHelper<ColumnVectorHelper, ColumnDecimal<T
     void insertManyDefaults(size_t length) override { data.resize_fill(data.size() + length); }
     void insert(const Field & x) override { data.push_back(x.get<T>()); }
     void insertRangeFrom(const IColumn & src, size_t start, size_t length) override;
+    void insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length) override;
 
     void popBack(size_t n) override
     {
diff --git a/src/Columns/ColumnFixedString.cpp b/src/Columns/ColumnFixedString.cpp
@@ -214,6 +214,22 @@ void ColumnFixedString::insertRangeFrom(const IColumn & src, size_t start, size_
     memcpy(chars.data() + old_size, &src_concrete.chars[start * n], length * n);
 }
 
+void ColumnFixedString::insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length)
+{
+    const ColumnFixedString & src_concrete = static_cast<const ColumnFixedString &>(src);
+
+    size_t old_size = chars.size();
+    chars.resize(old_size + length * n);
+    auto * cur_data_end = chars.data() + old_size;
+    auto * src_data_start = src_concrete.chars.data();
+
+    for (size_t i = 0; i < length; ++i)
+    {
+        size_t src_pos = selector[selector_start + i];
+        memcpySmallAllowReadWriteOverflow15(cur_data_end + i * n, src_data_start+ n * src_pos, n);
+    }
+}
+
 ColumnPtr ColumnFixedString::filter(const IColumn::Filter & filt, ssize_t result_size_hint) const
 {
     size_t col_size = size();
diff --git a/src/Columns/ColumnFixedString.h b/src/Columns/ColumnFixedString.h
@@ -154,6 +154,8 @@ class ColumnFixedString final : public COWHelper<ColumnVectorHelper, ColumnFixed
 
     void insertRangeFrom(const IColumn & src, size_t start, size_t length) override;
 
+    void insertRangeSelective(const IColumn & src, const Selector & selector, size_t selector_start, size_t length) override;
+
     ColumnPtr filter(const IColumn::Filter & filt, ssize_t result_size_hint) const override;
 
     void expand(const IColumn::Filter & mask, bool inverted) override;
diff --git a/src/Columns/ColumnMap.cpp b/src/Columns/ColumnMap.cpp
@@ -148,6 +148,11 @@ void ColumnMap::insertRangeFrom(const IColumn & src, size_t start, size_t length
         start, length);
 }
 
+void ColumnMap::insertRangeSelective(const IColumn & src, const Selector & selector, size_t selector_start, size_t length)
+{
+    nested->insertRangeSelective(assert_cast<const ColumnMap &>(src).getNestedColumn(), selector, selector_start, length);
+}
+
 ColumnPtr ColumnMap::filter(const Filter & filt, ssize_t result_size_hint) const
 {
     auto filtered = nested->filter(filt, result_size_hint);
diff --git a/src/Columns/ColumnMap.h b/src/Columns/ColumnMap.h
@@ -65,13 +65,15 @@ class ColumnMap final : public COWHelper<IColumn, ColumnMap>
     void updateWeakHash32(WeakHash32 & hash) const override;
     void updateHashFast(SipHash & hash) const override;
     void insertRangeFrom(const IColumn & src, size_t start, size_t length) override;
+    void insertRangeSelective(const IColumn & src, const Selector & selector, size_t selector_start, size_t length) override;
     ColumnPtr filter(const Filter & filt, ssize_t result_size_hint) const override;
     void expand(const Filter & mask, bool inverted) override;
     ColumnPtr permute(const Permutation & perm, size_t limit) const override;
     ColumnPtr index(const IColumn & indexes, size_t limit) const override;
     ColumnPtr replicate(const Offsets & offsets) const override;
     MutableColumns scatter(ColumnIndex num_columns, const Selector & selector) const override;
     void gather(ColumnGathererStream & gatherer_stream) override;
+    bool canBeInsideNullable() const override { return true; }
     int compareAt(size_t n, size_t m, const IColumn & rhs, int nan_direction_hint) const override;
     void compareColumn(const IColumn & rhs, size_t rhs_row_num,
                        PaddedPODArray<UInt64> * row_indexes, PaddedPODArray<Int8> & compare_results,
diff --git a/src/Columns/ColumnNullable.cpp b/src/Columns/ColumnNullable.cpp
@@ -12,6 +12,7 @@
 #include <Columns/ColumnConst.h>
 #include <Columns/ColumnString.h>
 #include <Columns/ColumnCompressed.h>
+#include <Columns/ColumnsCommon.h>
 #include <Columns/ColumnLowCardinality.h>
 #include <Processors/Transforms/ColumnGathererTransform.h>
 
@@ -245,6 +246,21 @@ void ColumnNullable::insertRangeFrom(const IColumn & src, size_t start, size_t l
     getNestedColumn().insertRangeFrom(*nullable_col.nested_column, start, length);
 }
 
+void ColumnNullable::insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length)
+{
+    const ColumnNullable & nullable_col = static_cast<const ColumnNullable &>(src);
+    getNestedColumn().insertRangeSelective(*nullable_col.nested_column, selector, selector_start, length);
+
+    if (!memoryIsZero(nullable_col.getNullMapData().data(), 0, nullable_col.size()))
+    {
+        getNullMapColumn().insertRangeSelective(*nullable_col.null_map, selector, selector_start, length);
+    }
+    else
+    {
+        getNullMapColumn().insertManyDefaults(length);
+    }
+}
+
 void ColumnNullable::insert(const Field & x)
 {
     if (x.isNull())
diff --git a/src/Columns/ColumnNullable.h b/src/Columns/ColumnNullable.h
@@ -67,6 +67,7 @@ class ColumnNullable final : public COWHelper<IColumn, ColumnNullable>
     const char * deserializeAndInsertFromArena(const char * pos) override;
     const char * skipSerializedInArena(const char * pos) const override;
     void insertRangeFrom(const IColumn & src, size_t start, size_t length) override;
+    void insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length) override;
     void insert(const Field & x) override;
     void insertFrom(const IColumn & src, size_t n) override;
 
@@ -129,6 +130,8 @@ class ColumnNullable final : public COWHelper<IColumn, ColumnNullable>
 
     void gather(ColumnGathererStream & gatherer_stream) override;
 
+    bool canBeInsideNullable() const override { return true; }
+
     ColumnPtr compress() const override;
 
     void forEachSubcolumn(MutableColumnCallback callback) override
diff --git a/src/Columns/ColumnString.cpp b/src/Columns/ColumnString.cpp
@@ -589,5 +589,41 @@ void ColumnString::validate() const
                         "ColumnString validation failed: size mismatch (internal logical error) {} != {}",
                         last_offset, chars.size());
 }
+void ColumnString::insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length)
+{
+    const ColumnString & src_concrete = static_cast<const ColumnString &>(src);
+    const Offsets & src_offsets = src_concrete.getOffsets();
+    auto * src_data_start = src_concrete.chars.data();
+
+    Offsets & cur_offsets = getOffsets();
+
+    if (length == 0)
+        return;
+
+    size_t old_offset_size = cur_offsets.size();
+    cur_offsets.resize(old_offset_size + length);
+
+    size_t old_chars_size = chars.size();
+    size_t new_chars_size = old_chars_size;
+    for (size_t i = 0; i < length; ++i)
+    {
+        new_chars_size += src_concrete.sizeAt(selector[selector_start + i]);
+    }
+    chars.resize(new_chars_size);
+
+    size_t cur_offset = cur_offsets[old_offset_size - 1];
+
+    auto * cur_chars_start = chars.data(); // realloc memory is not allowed in the following
+    for (size_t i = 0; i < length; ++i)
+    {
+        size_t src_pos = selector[selector_start + i];
+        size_t offset = src_offsets[src_pos - 1];
+        const size_t size_to_append = src_offsets[src_pos] - offset; /// -1th index is Ok, see PaddedPODArray.
+
+        memcpySmallAllowReadWriteOverflow15(cur_chars_start + cur_offset, src_data_start + offset, size_to_append);
+        cur_offset += size_to_append;
+        cur_offsets[old_offset_size + i] = cur_offset;
+    }
+}
 
 }
diff --git a/src/Columns/ColumnString.h b/src/Columns/ColumnString.h
@@ -195,6 +195,8 @@ class ColumnString final : public COWHelper<IColumn, ColumnString>
 
     void insertRangeFrom(const IColumn & src, size_t start, size_t length) override;
 
+    void insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length) override;
+
     ColumnPtr filter(const Filter & filt, ssize_t result_size_hint) const override;
 
     void expand(const Filter & mask, bool inverted) override;
diff --git a/src/Columns/ColumnTuple.cpp b/src/Columns/ColumnTuple.cpp
@@ -233,6 +233,18 @@ void ColumnTuple::insertRangeFrom(const IColumn & src, size_t start, size_t leng
             start, length);
 }
 
+void ColumnTuple::insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length)
+{
+    const ColumnTuple & src_concrete = static_cast<const ColumnTuple &>(src);
+
+    const size_t tuple_size = columns.size();
+    if (src_concrete.columns.size() != tuple_size)
+        throw Exception(ErrorCodes::CANNOT_INSERT_VALUE_OF_DIFFERENT_SIZE_INTO_TUPLE, "Cannot insert value of different size into tuple");
+
+    for (size_t i = 0; i < tuple_size; ++i)
+        columns[i]->insertRangeSelective(*src_concrete.columns[i], selector, selector_start, length);
+}
+
 ColumnPtr ColumnTuple::filter(const Filter & filt, ssize_t result_size_hint) const
 {
     const size_t tuple_size = columns.size();
diff --git a/src/Columns/ColumnTuple.h b/src/Columns/ColumnTuple.h
@@ -68,13 +68,15 @@ class ColumnTuple final : public COWHelper<IColumn, ColumnTuple>
     void updateWeakHash32(WeakHash32 & hash) const override;
     void updateHashFast(SipHash & hash) const override;
     void insertRangeFrom(const IColumn & src, size_t start, size_t length) override;
+    void insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length) override;
     ColumnPtr filter(const Filter & filt, ssize_t result_size_hint) const override;
     void expand(const Filter & mask, bool inverted) override;
     ColumnPtr permute(const Permutation & perm, size_t limit) const override;
     ColumnPtr index(const IColumn & indexes, size_t limit) const override;
     ColumnPtr replicate(const Offsets & offsets) const override;
     MutableColumns scatter(ColumnIndex num_columns, const Selector & selector) const override;
     void gather(ColumnGathererStream & gatherer_stream) override;
+    bool canBeInsideNullable() const override { return true; }
     int compareAt(size_t n, size_t m, const IColumn & rhs, int nan_direction_hint) const override;
     void compareColumn(const IColumn & rhs, size_t rhs_row_num,
                        PaddedPODArray<UInt64> * row_indexes, PaddedPODArray<Int8> & compare_results,
diff --git a/src/Columns/ColumnVector.cpp b/src/Columns/ColumnVector.cpp
@@ -633,6 +633,18 @@ inline void doFilterAligned(const UInt8 *& filt_pos, const UInt8 *& filt_end_ali
 }
 )
 
+template <typename T>
+void ColumnVector<T>::insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length)
+{
+size_t old_size = data.size();
+data.resize(old_size + length);
+const auto & src_data = (static_cast<const Self &>(src)).getData();
+for (size_t i = 0; i < length; ++i)
+{
+    data[old_size + i] = src_data[selector[selector_start + i]];
+}
+}
+
 template <typename T>
 ColumnPtr ColumnVector<T>::filter(const IColumn::Filter & filt, ssize_t result_size_hint) const
 {
diff --git a/src/Columns/ColumnVector.h b/src/Columns/ColumnVector.h
@@ -341,6 +341,8 @@ class ColumnVector final : public COWHelper<ColumnVectorHelper, ColumnVector<T>>
         return this->template scatterImpl<Self>(num_columns, selector);
     }
 
+    void insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length) override;
+
     void gather(ColumnGathererStream & gatherer_stream) override;
 
     bool canBeInsideNullable() const override { return true; }
diff --git a/src/Columns/IColumn.h b/src/Columns/IColumn.h
@@ -379,6 +379,18 @@ class IColumn : public COW<IColumn>
     using Selector = PaddedPODArray<ColumnIndex>;
     [[nodiscard]] virtual std::vector<MutablePtr> scatter(ColumnIndex num_columns, const Selector & selector) const = 0;
 
+    /// This function will get row index from selector and append the data to this column.
+    /// This function will handle indexes start from input 'selector_start' and will append 'size' times
+    /// For example:
+    ///      input selector: [1, 2, 3, 4, 5, 6]
+    ///      selector_start: 2
+    ///      length: 3
+    /// This function will copy the [3, 4, 5] row of src to this column.
+    virtual void insertRangeSelective(const IColumn & /*src*/, const Selector & /*selector*/, size_t /*selector_start*/, size_t /*length*/)
+    {
+        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "insertRangeSelective is not implemented for {}", getName());
+    }
+
     /// Insert data from several other columns according to source mask (used in vertical merge).
     /// For now it is a helper to de-virtualize calls to insert*() functions inside gather loop
     /// (descendants should call gatherer_stream.gather(*this) to implement this function.)

Original file line number	Diff line number	Diff line change
`@@ -126,6 +126,11 @@ class ColumnConst final : public COWHelper<IColumn, ColumnConst>`
`126`	`126`	`s += length;`
`127`	`127`	`}`
`128`	`128`
	`129`	`+ void insertRangeSelective(const IColumn & /src/, const Selector & /selector/, size_t /selector_start/, size_t length) override`
	`130`	`+ {`
	`131`	`+ s += length;`
	`132`	`+ }`
	`133`	`+`
`129`	`134`	`void insert(const Field &) override`
`130`	`135`	`{`
`131`	`136`	`++s;`
Original file line number	Diff line number	Diff line change
`@@ -62,6 +62,7 @@ class ColumnDecimal final : public COWHelper<ColumnVectorHelper, ColumnDecimal<T`
`62`	`62`	`void insertManyDefaults(size_t length) override { data.resize_fill(data.size() + length); }`
`63`	`63`	`void insert(const Field & x) override { data.push_back(x.get<T>()); }`
`64`	`64`	`void insertRangeFrom(const IColumn & src, size_t start, size_t length) override;`
	`65`	`+ void insertRangeSelective(const IColumn & src, const IColumn::Selector & selector, size_t selector_start, size_t length) override;`
`65`	`66`
`66`	`67`	`void popBack(size_t n) override`
`67`	`68`	`{`