Add IColumn::capacity

Signed-off-by: JaySon-Huang <[email protected]>

Author: JaySon-Huang

dbms/src/Columns/ColumnAggregateFunction.h

@@ -98,7 +98,7 @@ class ColumnAggregateFunction final : public COWPtrHelper<IColumn, ColumnAggrega
     {}
 
 public:
-    ~ColumnAggregateFunction();
+    ~ColumnAggregateFunction() override;
 
     void set(const AggregateFunctionPtr & func_) { func = func_; }
 
@@ -328,6 +328,8 @@ class ColumnAggregateFunction final : public COWPtrHelper<IColumn, ColumnAggrega
 
     void getPermutation(bool reverse, size_t limit, int nan_direction_hint, Permutation & res) const override;
 
+    size_t capacity() const override { return data.capacity(); }
+
     /** More efficient manipulation methods */
     Container & getData() { return data; }
 

dbms/src/Columns/ColumnArray.cpp

@@ -620,11 +620,19 @@ struct Less
 } // namespace
 
 
-void ColumnArray::reserve(size_t n)
+void ColumnArray::reserveWithStrategy(size_t n, ReserveStrategy strategy)
 {
-    getOffsets().reserve(n);
+    switch (strategy)
+    {
+    case ReserveStrategy::Default:
+        getOffsets().reserve(n);
+        break;
+    case ReserveStrategy::ScaleFactor1_5:
+        getOffsets().reserve(n / 2 * 3);
+        break;
+    }
     /// The average size of arrays is not taken into account here. Or it is considered to be no more than 1.
-    getData().reserve(n);
+    getData().reserveWithStrategy(n, strategy);
 }
 
 void ColumnArray::reserveAlign(size_t n, size_t alignment)

dbms/src/Columns/ColumnArray.h

@@ -214,7 +214,8 @@ class ColumnArray final : public COWPtrHelper<IColumn, ColumnArray>
     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;
     void getPermutation(bool reverse, size_t limit, int nan_direction_hint, Permutation & res) const override;
-    void reserve(size_t n) override;
+    size_t capacity() const override { return data->capacity(); }
+    void reserveWithStrategy(size_t n, ReserveStrategy strategy) override;
     void reserveAlign(size_t n, size_t alignment) override;
     size_t byteSize() const override;
     size_t byteSize(size_t offset, size_t limit) const override;

dbms/src/Columns/ColumnConst.h

@@ -261,6 +261,8 @@ class ColumnConst final : public COWPtrHelper<IColumn, ColumnConst>
     ColumnPtr permute(const Permutation & perm, size_t limit) const override;
     void getPermutation(bool reverse, size_t limit, int nan_direction_hint, Permutation & res) const override;
 
+    size_t capacity() const override { return data->capacity(); }
+
     size_t byteSize() const override { return data->byteSize() + sizeof(s); }
 
     size_t byteSize(size_t /*offset*/, size_t /*limit*/) const override { return byteSize(); }

dbms/src/Columns/ColumnDecimal.cpp

@@ -764,6 +764,20 @@ MutableColumnPtr ColumnDecimal<T>::cloneResized(size_t size) const
     return res;
 }
 
+template <typename T>
+void ColumnDecimal<T>::reserveWithStrategy(size_t n, IColumn::ReserveStrategy strategy)
+{
+    switch (strategy)
+    {
+    case IColumn::ReserveStrategy::Default:
+        data.reserve(n);
+        break;
+    case IColumn::ReserveStrategy::ScaleFactor1_5:
+        data.reserve_exact(n / 2 * 3);
+        break;
+    }
+}
+
 template <typename T>
 void ColumnDecimal<T>::insertData(const char * src [[maybe_unused]], size_t /*length*/)
 {

dbms/src/Columns/ColumnDecimal.h

@@ -113,8 +113,9 @@ class ColumnDecimal final : public COWPtrHelper<ColumnVectorHelper, ColumnDecima
     size_t byteSize() const override { return data.size() * sizeof(data[0]); }
     size_t byteSize(size_t /*offset*/, size_t limit) const override { return limit * sizeof(data[0]); }
     size_t allocatedBytes() const override { return data.allocated_bytes(); }
-    //void protect() override { data.protect(); }
-    void reserve(size_t n) override { data.reserve(n); }
+
+    size_t capacity() const override { return data.capacity(); }
+    void reserveWithStrategy(size_t n, IColumn::ReserveStrategy strategy) override;
     void reserveAlign(size_t n, size_t alignment) override { data.reserve(n, alignment); }
 
     void insertFrom(const IColumn & src, size_t n) override

dbms/src/Columns/ColumnFixedString.cpp

@@ -681,6 +681,19 @@ void ColumnFixedString::gather(ColumnGathererStream & gatherer)
     gatherer.gather(*this);
 }
 
+void ColumnFixedString::reserveWithStrategy(size_t size, IColumn::ReserveStrategy strategy)
+{
+    switch (strategy)
+    {
+    case ReserveStrategy::Default:
+        chars.reserve(n * size);
+        break;
+    case ReserveStrategy::ScaleFactor1_5:
+        chars.reserve_exact(n / 2 * 3 * size);
+        break;
+    }
+}
+
 void ColumnFixedString::getExtremes(Field & min, Field & max) const
 {
     min = String();

dbms/src/Columns/ColumnFixedString.h

@@ -193,6 +193,8 @@ class ColumnFixedString final : public COWPtrHelper<IColumn, ColumnFixedString>
 
     void getPermutation(bool reverse, size_t limit, int nan_direction_hint, Permutation & res) const override;
 
+    size_t capacity() const override { return chars.capacity() / n; }
+
     void insertRangeFrom(const IColumn & src, size_t start, size_t length) override;
 
     ColumnPtr filter(const IColumn::Filter & filt, ssize_t result_size_hint) const override;
@@ -223,7 +225,7 @@ class ColumnFixedString final : public COWPtrHelper<IColumn, ColumnFixedString>
 
     void gather(ColumnGathererStream & gatherer_stream) override;
 
-    void reserve(size_t size) override { chars.reserve(n * size); }
+    void reserveWithStrategy(size_t size, IColumn::ReserveStrategy strategy) override;
     void reserveAlign(size_t size, size_t alignment) override { chars.reserve(n * size, alignment); }
 
     void getExtremes(Field & min, Field & max) const override;

dbms/src/Columns/ColumnFunction.h

@@ -272,6 +272,9 @@ class ColumnFunction final : public COWPtrHelper<IColumn, ColumnFunction>
         throw Exception("getPermutation is not implemented for " + getName(), ErrorCodes::NOT_IMPLEMENTED);
     }
 
+    // A fake capacity for ColumnFunction, as it does not hold data in the same way as other columns.
+    size_t capacity() const override { return 0; }
+
     void gather(ColumnGathererStream &) override
     {
         throw Exception("Method gather is not supported for " + getName(), ErrorCodes::NOT_IMPLEMENTED);

dbms/src/Columns/ColumnNothing.h

@@ -29,9 +29,12 @@ class ColumnNothing final : public COWPtrHelper<IColumnDummy, ColumnNothing>
 
 public:
     const char * getFamilyName() const override { return "Nothing"; }
-    MutableColumnPtr cloneDummy(size_t s) const override { return ColumnNothing::create(s); };
+    MutableColumnPtr cloneDummy(size_t s) const override { return ColumnNothing::create(s); }
 
     bool canBeInsideNullable() const override { return true; }
+
+    // ColumnNothing does not hold any data, so capacity is 0.
+    size_t capacity() const override { return 1; }
 };
 
 } // namespace DB