Use size of filtered column to reserve memory for filter when result_size_hint < 0 (#6746)

dbms/src/Columns/ColumnAggregateFunction.cpp

@@ -14,6 +14,7 @@
 
 #include <AggregateFunctions/AggregateFunctionState.h>
 #include <Columns/ColumnAggregateFunction.h>
+#include <Columns/ColumnsCommon.h>
 #include <Common/HashTable/Hash.h>
 #include <Common/SipHash.h>
 #include <Common/typeid_cast.h>
@@ -143,7 +144,11 @@ ColumnPtr ColumnAggregateFunction::filter(const Filter & filter, ssize_t result_
     auto & res_data = res->getData();
 
     if (result_size_hint)
-        res_data.reserve(result_size_hint > 0 ? result_size_hint : size);
+    {
+        if (result_size_hint < 0)
+            result_size_hint = countBytesInFilter(filter);
+        res_data.reserve(result_size_hint);
+    }
 
     for (size_t i = 0; i < size; ++i)
         if (filter[i])

dbms/src/Columns/ColumnArray.h

@@ -87,6 +87,7 @@ class ColumnArray final : public COWPtrHelper<IColumn, ColumnArray>
     void insertFrom(const IColumn & src_, size_t n) override;
     void insertDefault() override;
     void popBack(size_t n) override;
+    /// TODO: If result_size_hint < 0, makes reserve() using size of filtered column, not source column to avoid some OOM issues.
     ColumnPtr filter(const Filter & filt, ssize_t result_size_hint) const override;
     ColumnPtr permute(const Permutation & perm, size_t limit) const override;
     int compareAt(size_t n, size_t m, const IColumn & rhs_, int nan_direction_hint) const override;

dbms/src/Columns/ColumnDecimal.cpp

@@ -306,7 +306,11 @@ ColumnPtr ColumnDecimal<T>::filter(const IColumn::Filter & filt, ssize_t result_
     Container & res_data = res->getData();
 
     if (result_size_hint)
-        res_data.reserve(result_size_hint > 0 ? result_size_hint : size);
+    {
+        if (result_size_hint < 0)
+            result_size_hint = countBytesInFilter(filt);
+        res_data.reserve(result_size_hint);
+    }
 
     const UInt8 * filt_pos = filt.data();
     const UInt8 * filt_end = filt_pos + size;

dbms/src/Columns/ColumnFixedString.cpp

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include <Columns/ColumnFixedString.h>
+#include <Columns/ColumnsCommon.h>
 #include <Common/Arena.h>
 #include <Common/HashTable/Hash.h>
 #include <Common/SipHash.h>
@@ -205,7 +206,11 @@ ColumnPtr ColumnFixedString::filter(const IColumn::Filter & filt, ssize_t result
     auto res = ColumnFixedString::create(n);
 
     if (result_size_hint)
-        res->chars.reserve(result_size_hint > 0 ? result_size_hint * n : chars.size());
+    {
+        if (result_size_hint < 0)
+            result_size_hint = countBytesInFilter(filt);
+        res->chars.reserve(result_size_hint * n);
+    }
 
     const UInt8 * filt_pos = &filt[0];
     const UInt8 * filt_end = filt_pos + col_size;

dbms/src/Columns/ColumnVector.cpp

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include <Columns/ColumnVector.h>
+#include <Columns/ColumnsCommon.h>
 #include <Common/Arena.h>
 #include <Common/Exception.h>
 #include <Common/HashTable/Hash.h>
@@ -213,7 +214,11 @@ ColumnPtr ColumnVector<T>::filter(const IColumn::Filter & filt, ssize_t result_s
     Container & res_data = res->getData();
 
     if (result_size_hint)
-        res_data.reserve(result_size_hint > 0 ? result_size_hint : size);
+    {
+        if (result_size_hint < 0)
+            result_size_hint = countBytesInFilter(filt);
+        res_data.reserve(result_size_hint);
+    }
 
     const UInt8 * filt_pos = &filt[0];
     const UInt8 * filt_end = filt_pos + size;

dbms/src/Columns/ColumnsCommon.cpp

@@ -122,9 +122,9 @@ struct ResultOffsetsBuilder
         : res_offsets(*res_offsets_)
     {}
 
-    void reserve(ssize_t result_size_hint, size_t src_size)
+    void reserve(size_t result_size_hint)
     {
-        res_offsets.reserve(result_size_hint > 0 ? result_size_hint : src_size);
+        res_offsets.reserve(result_size_hint);
     }
 
     void insertOne(size_t array_size)
@@ -165,7 +165,7 @@ struct ResultOffsetsBuilder
 struct NoResultOffsetsBuilder
 {
     explicit NoResultOffsetsBuilder(IColumn::Offsets *) {}
-    void reserve(ssize_t, size_t) {}
+    void reserve(size_t) {}
     void insertOne(size_t) {}
 
     template <size_t SIMD_BYTES>
@@ -196,11 +196,12 @@ void filterArraysImplGeneric(
 
     if (result_size_hint)
     {
-        result_offsets_builder.reserve(result_size_hint, size);
-
         if (result_size_hint < 0)
-            res_elems.reserve(src_elems.size());
-        else if (result_size_hint < 1000000000 && src_elems.size() < 1000000000) /// Avoid overflow.
+            result_size_hint = countBytesInFilter(filt);
+
+        result_offsets_builder.reserve(result_size_hint);
+
+        if (result_size_hint < 1000000000 && src_elems.size() < 1000000000) /// Avoid overflow.
             res_elems.reserve((result_size_hint * src_elems.size() + size - 1) / size);
     }
 

dbms/src/Columns/IColumn.h

@@ -216,7 +216,7 @@ class IColumn : public COWPtr<IColumn>
       * Is used in WHERE and HAVING operations.
       * If result_size_hint > 0, then makes advance reserve(result_size_hint) for the result column;
       *  if 0, then don't makes reserve(),
-      *  otherwise (i.e. < 0), makes reserve() using size of source column.
+      *  otherwise (i.e. < 0), makes reserve() using size of filtered column.
       */
     using Filter = PaddedPODArray<UInt8>;
     virtual Ptr filter(const Filter & filt, ssize_t result_size_hint) const = 0;