[BugFix] fix Column::element_memory_usage (StarRocks#17184)

huandzh · Feb 2, 2023 · 187c52c · 187c52c
1 parent 4a7b324
commit 187c52c
Show file tree

Hide file tree

Showing 12 changed files with 141 additions and 10 deletions.
diff --git a/be/src/column/array_column.cpp b/be/src/column/array_column.cpp
@@ -536,8 +536,7 @@ size_t ArrayColumn::element_memory_usage(size_t from, size_t size) const {
     DCHECK_LE(from + size, this->size()) << "Range error";
     size_t start_offset = _offsets->get_data()[from];
     size_t elements_num = _offsets->get_data()[from + size] - start_offset;
-    return _elements->element_memory_usage(start_offset, elements_num) +
-           _offsets->Column::element_memory_usage(from, size);
+    return _elements->element_memory_usage(start_offset, elements_num) + _offsets->element_memory_usage(from, size);
 }
 
 void ArrayColumn::swap_column(Column& rhs) {

diff --git a/be/src/column/binary_column.h b/be/src/column/binary_column.h
@@ -277,6 +277,10 @@ class BinaryColumnBase final : public ColumnFactory<Column, BinaryColumnBase<T>>
         return _bytes.capacity() + _offsets.capacity() * sizeof(_offsets[0]) + _slices.capacity() * sizeof(_slices[0]);
     }
 
+    size_t element_memory_usage(size_t from, size_t size) const override {
+        return _offsets[from + size] - _offsets[from] + size * sizeof(T);
+    }
+
     void swap_column(Column& rhs) override {
         auto& r = down_cast<BinaryColumnBase<T>&>(rhs);
         using std::swap;

diff --git a/be/src/column/column.h b/be/src/column/column.h
@@ -379,7 +379,7 @@ class Column {
     virtual size_t memory_usage() const { return container_memory_usage() + element_memory_usage(); }
     virtual size_t container_memory_usage() const = 0;
     virtual size_t element_memory_usage() const { return element_memory_usage(0, size()); }
-    virtual size_t element_memory_usage(size_t from, size_t size) const { return 0; }
+    virtual size_t element_memory_usage(size_t from, size_t size) const = 0;
 
     virtual void swap_column(Column& rhs) = 0;
 

diff --git a/be/src/column/const_column.h b/be/src/column/const_column.h
@@ -201,7 +201,7 @@ class ConstColumn final : public ColumnFactory<Column, ConstColumn> {
 
     size_t element_memory_usage(size_t from, size_t size) const override {
         // const column has only one element
-        return element_memory_usage();
+        return size == 0 ? 0 : element_memory_usage();
     }
 
     void swap_column(Column& rhs) override {

diff --git a/be/src/column/fixed_length_column_base.h b/be/src/column/fixed_length_column_base.h
@@ -206,6 +206,7 @@ class FixedLengthColumnBase : public ColumnFactory<Column, FixedLengthColumnBase
     std::string debug_string() const override;
 
     size_t container_memory_usage() const override { return _data.capacity() * sizeof(ValueType); }
+    size_t element_memory_usage(size_t from, size_t size) const override { return sizeof(ValueType) * size; }
 
     void swap_column(Column& rhs) override {
         auto& r = down_cast<FixedLengthColumnBase&>(rhs);

diff --git a/be/src/column/map_column.cpp b/be/src/column/map_column.cpp
@@ -540,9 +540,10 @@ bool MapColumn::set_null(size_t idx) {
 
 size_t MapColumn::element_memory_usage(size_t from, size_t size) const {
     DCHECK_LE(from + size, this->size()) << "Range error";
-    return _keys->element_memory_usage(_offsets->get_data()[from], _offsets->get_data()[from + size]) +
-           _values->element_memory_usage(_offsets->get_data()[from], _offsets->get_data()[from + size]) +
-           _offsets->Column::element_memory_usage(from, size);
+    size_t start_offset = _offsets->get_data()[from];
+    size_t elements_num = _offsets->get_data()[from + size] - start_offset;
+    return _keys->element_memory_usage(start_offset, elements_num) +
+           _values->element_memory_usage(start_offset, elements_num) + _offsets->element_memory_usage(from, size);
 }
 
 void MapColumn::swap_column(Column& rhs) {

diff --git a/be/test/column/array_column_test.cpp b/be/test/column/array_column_test.cpp
@@ -1190,4 +1190,43 @@ PARALLEL_TEST(ArrayColumnTest, test_replicate) {
     ASSERT_EQ("[]", res->debug_item(6));
 }
 
+PARALLEL_TEST(ArrayColumnTest, test_element_memory_usage) {
+    auto offsets = UInt32Column::create();
+    auto elements = Int32Column::create();
+    auto column = ArrayColumn::create(elements, offsets);
+
+    // insert [],[1],[2, 3],[4, 5, 6]
+    offsets->append(0);
+
+    elements->append(1);
+    offsets->append(1);
+
+    elements->append(2);
+    elements->append(3);
+    offsets->append(3);
+
+    elements->append(4);
+    elements->append(5);
+    elements->append(6);
+    offsets->append(6);
+
+    ASSERT_EQ("[]", column->debug_item(0));
+    ASSERT_EQ("[1]", column->debug_item(1));
+    ASSERT_EQ("[2,3]", column->debug_item(2));
+    ASSERT_EQ("[4,5,6]", column->debug_item(3));
+
+    ASSERT_EQ(40, column->Column::element_memory_usage());
+
+    std::vector<size_t> element_mem_usages = {4, 8, 12, 16};
+    size_t element_num = element_mem_usages.size();
+    for (size_t start = 0; start < element_num; start++) {
+        size_t expected_usage = 0;
+        ASSERT_EQ(expected_usage, column->element_memory_usage(start, 0));
+        for (size_t size = 1; start + size <= element_num; size++) {
+            expected_usage += element_mem_usages[start + size - 1];
+            ASSERT_EQ(expected_usage, column->element_memory_usage(start, size));
+        }
+    }
+}
+
 } // namespace starrocks
diff --git a/be/test/column/binary_column_test.cpp b/be/test/column/binary_column_test.cpp
@@ -655,4 +655,25 @@ PARALLEL_TEST(BinaryColumnTest, test_replicate) {
     ASSERT_EQ("def", slices[4]);
 }
 
+PARALLEL_TEST(BinaryColumnTest, test_element_memory_usage) {
+    auto column = BinaryColumn::create();
+    column->append("");
+    column->append("1");
+    column->append("23");
+    column->append("456");
+
+    ASSERT_EQ(22, column->Column::element_memory_usage());
+
+    std::vector<size_t> element_mem_usages = {4, 5, 6, 7};
+    size_t element_num = element_mem_usages.size();
+    for (size_t start = 0; start < element_num; start++) {
+        size_t expected_usage = 0;
+        ASSERT_EQ(0, column->element_memory_usage(start, 0));
+        for (size_t size = 1; start + size <= element_num; size++) {
+            expected_usage += element_mem_usages[start + size - 1];
+            ASSERT_EQ(expected_usage, column->element_memory_usage(start, size));
+        }
+    }
+}
+
 } // namespace starrocks
diff --git a/be/test/column/const_column_test.cpp b/be/test/column/const_column_test.cpp
@@ -354,4 +354,22 @@ PARALLEL_TEST(ConstColumnTest, test_replicate) {
     ASSERT_EQ(1, c2->get(6).get_int32());
 }
 
+PARALLEL_TEST(ConstColumnTest, test_element_memory_usage) {
+    auto create_const_column = [](int32_t value, size_t size) {
+        auto c = Int32Column::create();
+        c->append_numbers(&value, sizeof(value));
+        return ConstColumn::create(c, size);
+    };
+
+    auto column = create_const_column(1, 10);
+    ASSERT_EQ(4, column->element_memory_usage());
+
+    for (size_t start = 0; start < column->size(); start++) {
+        ASSERT_EQ(0, column->element_memory_usage(start, 0));
+        for (size_t size = 1; start + size <= column->size(); size++) {
+            ASSERT_EQ(4, column->element_memory_usage(start, size));
+        }
+    }
+}
+
 } // namespace starrocks
diff --git a/be/test/column/map_column_test.cpp b/be/test/column/map_column_test.cpp
@@ -1281,4 +1281,28 @@ PARALLEL_TEST(MapColumnTest, test_euqals) {
     ASSERT_FALSE(lhs->equals(2, *rhs, 2));
     ASSERT_FALSE(lhs->equals(3, *rhs, 3));
 }
+
+PARALLEL_TEST(MapColumnTest, test_element_memory_usage) {
+    auto column = MapColumn::create(NullableColumn::create(Int32Column::create(), NullColumn::create()),
+                                    NullableColumn::create(Int32Column::create(), NullColumn::create()),
+                                    UInt32Column::create());
+
+    // {}, {1:2},{3:4,5:6}
+    column->append_datum(DatumMap{});
+    column->append_datum(DatumMap{{1, 2}});
+    column->append_datum(DatumMap{{3, 4}, {5, 6}});
+
+    ASSERT_EQ(42, column->Column::element_memory_usage());
+
+    std::vector<size_t> element_mem_usages = {4, 14, 24};
+    size_t element_num = element_mem_usages.size();
+    for (size_t start = 0; start < element_num; start++) {
+        size_t expected_usage = 0;
+        ASSERT_EQ(0, column->element_memory_usage(start, 0));
+        for (size_t size = 1; start + size <= element_num; size++) {
+            expected_usage += element_mem_usages[start + size - 1];
+            ASSERT_EQ(expected_usage, column->element_memory_usage(start, size));
+        }
+    }
+}
 } // namespace starrocks
diff --git a/be/test/column/struct_column_test.cpp b/be/test/column/struct_column_test.cpp
@@ -550,4 +550,28 @@ TEST(StructColumnTest, test_unnamed) {
     ASSERT_EQ("{1,'smith'}, {2,'cruise'}", column->debug_string());
 }
 
+TEST(StructColumnTest, test_element_memory_usage) {
+    auto id = NullableColumn::create(UInt64Column::create(), NullColumn::create());
+    auto name = NullableColumn::create(BinaryColumn::create(), NullColumn::create());
+    Columns fields{id, name};
+    auto column = StructColumn::create(fields);
+
+    column->append_datum(DatumStruct{uint64_t(1), Slice("2")});
+    column->append_datum(DatumStruct{uint64_t(1), Slice("4")});
+    column->append_datum(DatumStruct{uint64_t(1), Slice("6")});
+
+    ASSERT_EQ(45, column->Column::element_memory_usage());
+
+    std::vector<size_t> element_mem_usages = {15, 15, 15};
+    size_t element_num = element_mem_usages.size();
+    for (size_t start = 0; start < element_num; start++) {
+        size_t expected_usage = 0;
+        ASSERT_EQ(0, column->element_memory_usage(start, 0));
+        for (size_t size = 1; start + size <= element_num; size++) {
+            expected_usage += element_mem_usages[start + size - 1];
+            ASSERT_EQ(expected_usage, column->element_memory_usage(start, size));
+        }
+    }
+}
+
 } // namespace starrocks
diff --git a/be/test/exec/query_cache/query_cache_test.cpp b/be/test/exec/query_cache/query_cache_test.cpp
@@ -110,18 +110,18 @@ TEST_F(QueryCacheTest, testCacheManager) {
         cache_mgr->populate(strings::Substitute("key_$0", i), create_cache_value(96));
     }
 
-    ASSERT_EQ(cache_mgr->memory_usage(), 960);
+    ASSERT_EQ(cache_mgr->memory_usage(), 1360);
     for (auto i = 0; i < 10; ++i) {
         auto status = cache_mgr->probe(strings::Substitute("key_$0", i));
         ASSERT_TRUE(status.ok());
     }
 
-    ASSERT_EQ(cache_mgr->memory_usage(), 960);
+    ASSERT_EQ(cache_mgr->memory_usage(), 1360);
     for (auto i = 10; i < 20; ++i) {
         auto status = cache_mgr->probe(strings::Substitute("key_$0", i));
         ASSERT_FALSE(status.ok());
     }
-    ASSERT_EQ(cache_mgr->memory_usage(), 960);
+    ASSERT_EQ(cache_mgr->memory_usage(), 1360);
 
     for (auto i = 20; i < 30; ++i) {
         auto status = cache_mgr->populate(strings::Substitute("key_$0", i), create_cache_value(100));