Merge pull request #330 from rtcote/rtcote_add_numeric_methods

Add reserve, capacity methods to numeric columns
ClickHouse · Oct 9, 2023 · e750dd5 · e750dd5
2 parents 89cd7bc + 979bb9f
commit e750dd5
Show file tree

Hide file tree

Showing 5 changed files with 184 additions and 0 deletions.
diff --git a/clickhouse/columns/date.cpp b/clickhouse/columns/date.cpp
@@ -43,6 +43,18 @@ void ColumnDate::Append(ColumnRef column) {
     }
 }
 
+std::vector<uint16_t> &ColumnDate::GetRawVector() {
+    return data_->GetRawVector();
+}
+
+void ColumnDate::Reserve(size_t new_cap) {
+    data_->Reserve(new_cap);
+}
+
+size_t ColumnDate::Capacity() const {
+    return data_->Capacity();
+}
+
 bool ColumnDate::LoadBody(InputStream* input, size_t rows) {
     return data_->LoadBody(input, rows);
 }
@@ -110,6 +122,18 @@ void ColumnDate32::Append(ColumnRef column) {
     }
 }
 
+std::vector<int32_t> & ColumnDate32::GetRawVector() {
+    return data_->GetRawVector();
+}
+
+void ColumnDate32::Reserve(size_t new_cap) {
+    data_->Reserve(new_cap);
+}
+
+size_t ColumnDate32::Capacity() const {
+    return data_->Capacity();
+}
+
 void ColumnDate32::AppendRaw(int32_t value) {
     data_->Append(value);
 }
@@ -182,6 +206,10 @@ std::time_t ColumnDateTime::At(size_t n) const {
     return data_->At(n);
 }
 
+void ColumnDateTime::AppendRaw(uint32_t value) {
+    data_->Append(value);
+}
+
 std::string ColumnDateTime::Timezone() const {
     return type_->As<DateTimeType>()->Timezone();
 }
@@ -192,6 +220,18 @@ void ColumnDateTime::Append(ColumnRef column) {
     }
 }
 
+std::vector<uint32_t> & ColumnDateTime::GetRawVector() {
+    return data_->GetRawVector();
+}
+
+void ColumnDateTime::Reserve(size_t new_cap) {
+    data_->Reserve(new_cap);
+}
+
+size_t ColumnDateTime::Capacity() const {
+    return data_->Capacity();
+}
+
 bool ColumnDateTime::LoadBody(InputStream* input, size_t rows) {
     return data_->LoadBody(input, rows);
 }

diff --git a/clickhouse/columns/date.h b/clickhouse/columns/date.h
@@ -31,6 +31,15 @@ class ColumnDate : public Column {
     /// Appends content of given column to the end of current one.
     void Append(ColumnRef column) override;
 
+    /// Get Raw Vector Contents
+    std::vector<uint16_t> & GetRawVector();
+
+    /// Increase the capacity of the column
+    void Reserve(size_t new_cap);
+
+    /// Returns the capacity of the column
+    size_t Capacity() const;
+
     /// Loads column data from input stream.
     bool LoadBody(InputStream* input, size_t rows) override;
 
@@ -79,6 +88,15 @@ class ColumnDate32 : public Column {
     void AppendRaw(int32_t value);
     int32_t RawAt(size_t n) const;
 
+    /// Get Raw Vector Contents
+    std::vector<int32_t> & GetRawVector();
+
+    /// Increase the capacity of the column
+    void Reserve(size_t new_cap);
+
+    /// Returns the capacity of the column
+    size_t Capacity() const;
+
     /// Loads column data from input stream.
     bool LoadBody(InputStream* input, size_t rows) override;
 
@@ -121,9 +139,21 @@ class ColumnDateTime : public Column {
     std::time_t At(size_t n) const;
     inline std::time_t operator [] (size_t n) const { return At(n); }
 
+    /// Append raw as UNIX epoch seconds in uint32
+    void AppendRaw(uint32_t value);
+
     /// Timezone associated with a data column.
     std::string Timezone() const;
 
+    /// Get Raw Vector Contents
+    std::vector<uint32_t> & GetRawVector();
+
+    /// Increase the capacity of the column
+    void Reserve(size_t new_cap);
+
+    /// Returns the capacity of the column
+    size_t Capacity() const;
+
 public:
     /// Appends content of given column to the end of current one.
     void Append(ColumnRef column) override;

diff --git a/clickhouse/columns/numeric.cpp b/clickhouse/columns/numeric.cpp
@@ -38,6 +38,21 @@ void ColumnVector<T>::Erase(size_t pos, size_t count) {
     data_.erase(data_.begin() + begin, data_.begin() + last);
 }
 
+template <typename T>
+std::vector<T> & ColumnVector<T>::GetRawVector() {
+    return data_;
+}
+
+template <typename T>
+void ColumnVector<T>::Reserve(size_t new_cap) {
+    data_.reserve(new_cap);
+}
+
+template <typename T>
+size_t ColumnVector<T>::Capacity() const {
+    return data_.capacity();
+}
+
 template <typename T>
 void ColumnVector<T>::Clear() {
     data_.clear();

diff --git a/clickhouse/columns/numeric.h b/clickhouse/columns/numeric.h
@@ -30,6 +30,15 @@ class ColumnVector : public Column {
 
     void Erase(size_t pos, size_t count = 1);
 
+    /// Get Raw Vector Contents
+    std::vector<T> & GetRawVector();
+
+    /// Increase the capacity of the column
+    void Reserve(size_t new_cap);
+
+    /// Returns the capacity of the column
+    size_t Capacity() const;
+
 public:
     /// Appends content of given column to the end of current one.
     void Append(ColumnRef column) override;

diff --git a/ut/columns_ut.cpp b/ut/columns_ut.cpp
@@ -965,3 +965,93 @@ TEST(ColumnsCase, ColumnMapT_Wrap) {
     EXPECT_EQ("123", map_view.At(1));
     EXPECT_EQ("abc", map_view.At(2));
 }
+
+TEST(ColumnsCase, ReservedAndCapacity) {
+    std::vector<std::shared_ptr<Column>> columns;
+
+    #define RaC_TEST_CASE(test_id_in, column_type) \
+        case test_id_in: { \
+            columns.push_back(std::make_shared<column_type>()); \
+            auto column = std::static_pointer_cast<column_type>(columns[test_id_in]); \
+            column->Reserve(100u); \
+            ASSERT_EQ(column->Capacity(), 100u); \
+            ASSERT_EQ(columns[test_id_in]->Size(), 0u); \
+            break; \
+        }
+
+    for (uint8_t rac_test_id = 0; rac_test_id < 14; ++rac_test_id) {
+        switch (rac_test_id) {
+            RaC_TEST_CASE( 0, ColumnUInt8);
+            RaC_TEST_CASE( 1, ColumnUInt16);
+            RaC_TEST_CASE( 2, ColumnUInt32);
+            RaC_TEST_CASE( 3, ColumnUInt64);
+            RaC_TEST_CASE( 4, ColumnInt8);
+            RaC_TEST_CASE( 5, ColumnInt16);
+            RaC_TEST_CASE( 6, ColumnInt32);
+            RaC_TEST_CASE( 7, ColumnInt64);
+            RaC_TEST_CASE( 8, ColumnInt128);
+            RaC_TEST_CASE( 9, ColumnFloat32);
+            RaC_TEST_CASE(10, ColumnFloat64);
+            RaC_TEST_CASE(11, ColumnDate);
+            RaC_TEST_CASE(12, ColumnDate32);
+            RaC_TEST_CASE(13, ColumnDateTime);
+            default: {
+                EXPECT_NE(0, 0);
+                break;
+            }
+        }
+    }
+}
+
+TEST(ColumnsCase, RawVector) {
+    std::vector<std::shared_ptr<Column>> columns;
+
+    #define RV_TEST_CASE(test_id_in, column_type) \
+        case test_id_in: { \
+            columns.push_back(std::make_shared<column_type>()); \
+            auto column = std::static_pointer_cast<column_type>(columns[test_id_in]); \
+            column->Append(10u); \
+            column->Append(20u); \
+            ASSERT_EQ(columns[test_id_in]->Size(), 2u); \
+            auto column_v = column->GetRawVector(); \
+            ASSERT_EQ(static_cast<uint32_t>(column_v[0]), 10u); \
+            ASSERT_EQ(static_cast<uint32_t>(column_v[1]), 20u); \
+            break; \
+        }
+
+    #define RV_TEST_CASE_D(test_id_in, column_type) \
+        case test_id_in: { \
+            columns.push_back(std::make_shared<column_type>()); \
+            auto column = std::static_pointer_cast<column_type>(columns[test_id_in]); \
+            column->AppendRaw(10u); \
+            column->AppendRaw(20u); \
+            ASSERT_EQ(columns[test_id_in]->Size(), 2u); \
+            auto column_v = column->GetRawVector(); \
+            ASSERT_EQ(static_cast<uint32_t>(column_v[0]), 10u); \
+            ASSERT_EQ(static_cast<uint32_t>(column_v[1]), 20u); \
+            break; \
+        }
+
+    for (uint8_t rv_test_id = 0; rv_test_id < 14; ++rv_test_id) {
+        switch (rv_test_id) {
+            RV_TEST_CASE(    0, ColumnUInt8);
+            RV_TEST_CASE(    1, ColumnUInt16);
+            RV_TEST_CASE(    2, ColumnUInt32);
+            RV_TEST_CASE(    3, ColumnUInt64);
+            RV_TEST_CASE(    4, ColumnInt8);
+            RV_TEST_CASE(    5, ColumnInt16);
+            RV_TEST_CASE(    6, ColumnInt32);
+            RV_TEST_CASE(    7, ColumnInt64);
+            RV_TEST_CASE(    8, ColumnInt128);
+            RV_TEST_CASE(    9, ColumnFloat32);
+            RV_TEST_CASE(   10, ColumnFloat64);
+            RV_TEST_CASE_D( 11, ColumnDate);
+            RV_TEST_CASE_D( 12, ColumnDate32);
+            RV_TEST_CASE_D( 13, ColumnDateTime);
+            default: {
+                EXPECT_NE(0, 0);
+                break;
+            }
+        }
+    }
+}