diff --git a/be/src/olap/rowset/column_reader.cpp b/be/src/olap/rowset/column_reader.cpp index a6bc0857d5..5524a4a158 100644 --- a/be/src/olap/rowset/column_reader.cpp +++ b/be/src/olap/rowset/column_reader.cpp @@ -206,7 +206,13 @@ OLAPStatus StringColumnDirectReader::next_vector( string_buffer_size += length; } - char* string_buffer = reinterpret_cast(mem_pool->allocate(string_buffer_size)); + uint8_t* allocated_mem; + res = mem_pool->allocate_safely(string_buffer_size, allocated_mem); + if (OLAP_SUCCESS != res) { + return res; + } + char* string_buffer = reinterpret_cast(allocated_mem); + for (int i = 0; i < size; ++i) { length = _values[i].size; if (UNLIKELY(length == 0)) { @@ -239,7 +245,13 @@ OLAPStatus StringColumnDirectReader::next_vector( } } - char* string_buffer = reinterpret_cast(mem_pool->allocate(string_buffer_size)); + uint8_t* allocated_mem; + res = mem_pool->allocate_safely(string_buffer_size, allocated_mem); + if (OLAP_SUCCESS != res) { + return res; + } + char* string_buffer = reinterpret_cast(allocated_mem); + for (int i = 0; i < size; ++i) { if (!is_null[i]) { length = _values[i].size; diff --git a/be/src/runtime/mem_pool.h b/be/src/runtime/mem_pool.h index a9da9cf0df..1933cdfeec 100644 --- a/be/src/runtime/mem_pool.h +++ b/be/src/runtime/mem_pool.h @@ -30,6 +30,7 @@ #include "gutil/dynamic_annotations.h" #include "util/bit_util.h" #include "runtime/memory/chunk.h" +#include "olap/olap_define.h" namespace doris { @@ -110,6 +111,11 @@ public: return allocate(size, DEFAULT_ALIGNMENT); } + /// Same as Allocate() excpect add a check when return a nullptr + OLAPStatus allocate_safely(int64_t size, uint8_t*& ret) { + return allocate_safely(size, DEFAULT_ALIGNMENT, ret); + } + /// Same as Allocate() except the mem limit is checked before the allocation and /// this call will fail (returns NULL) if it does. /// The caller must handle the NULL case. This should be used for allocations @@ -242,6 +248,16 @@ private: return result; } + template + OLAPStatus ALWAYS_INLINE allocate_safely(int64_t size, int alignment, uint8_t*& ret) { + uint8_t* result = allocate(size, alignment); + if (result == nullptr) { + return OLAP_ERR_MALLOC_ERROR; + } + ret = result; + return OLAP_SUCCESS; + } + private: /// chunk from which we served the last Allocate() call; /// always points to the last chunk that contains allocated data;