Grow column bodies and wire strings as the data arrives

clickhouse/base/wire_format.h

@@ -1,7 +1,8 @@
 #pragma once
 
-#include <string>
+#include <algorithm>
 #include <cstdint>
+#include <string>
 
 namespace clickhouse {
 
@@ -15,6 +16,20 @@ class WireFormat {
     static bool ReadString(InputStream& input, std::string* value);
     static bool SkipString(InputStream& input);
     static bool ReadBytes(InputStream& input, void* buf, size_t len);
+
+    /// Replaces the contents of `container` with `count` elements read from `input`.
+    /// The container grows with the data instead of being sized to `count` up front, so
+    /// a count declared by the peer is never allocated before its bytes are read.
+    template <typename Container>
+    static bool ReadPODArray(InputStream& input, Container* container, size_t count);
+
+    /// Capacity to grow to once `loaded` of `count` elements are held: eight times `loaded`
+    /// and at least `first`, capped at `count`, without overflowing `Size`.
+    template <typename Size>
+    static constexpr Size NextReadCapacity(Size loaded, Size count, Size first) {
+        return loaded > count / 8 ? count : std::min(count, std::max(first, static_cast<Size>(loaded * 8)));
+    }
+
     static bool ReadUInt64(InputStream& input, uint64_t* value);
     static bool ReadVarint64(InputStream& output, uint64_t* value);
 
@@ -43,8 +58,7 @@ inline bool WireFormat::ReadString(InputStream& input, std::string* value) {
         if (len > 0x00FFFFFFULL) {
             return false;
         }
-        value->resize((size_t)len);
-        return ReadAll(input, value->data(), (size_t)len);
+        return ReadPODArray(input, value, (size_t)len);
     }
 
     return false;
@@ -54,6 +68,30 @@ inline bool WireFormat::ReadBytes(InputStream& input, void* buf, size_t len) {
     return ReadAll(input, buf, len);
 }
 
+template <typename Container>
+inline bool WireFormat::ReadPODArray(InputStream& input, Container* container, size_t count) {
+    using T                      = typename Container::value_type;
+    constexpr size_t first_chunk = std::max<size_t>(1, (1 << 20) / sizeof(T));
+
+    container->clear();
+    size_t loaded = 0;
+    while (loaded < count) {
+        // Reserving far ahead keeps reallocation copies small; resizing only by what is
+        // about to be read keeps the zero-filled part proportional to the data received.
+        if (container->capacity() <= loaded) {
+            container->reserve(NextReadCapacity(loaded, count, first_chunk));
+        }
+        const size_t chunk = std::min({count - loaded, container->capacity() - loaded, std::max(loaded, first_chunk)});
+        container->resize(loaded + chunk);
+        if (!ReadAll(input, container->data() + loaded, chunk * sizeof(T))) {
+            return false;
+        }
+        loaded += chunk;
+    }
+
+    return true;
+}
+
 inline bool WireFormat::ReadUInt64(InputStream& input, uint64_t* value) {
     return ReadVarint64(input, value);
 }

clickhouse/client.cpp

@@ -835,6 +835,10 @@ bool Client::Impl::ReadBlock(InputStream& input, Block* block) {
     if (!WireFormat::ReadUInt64(input, &num_rows)) {
         return false;
     }
+    if (static_cast<size_t>(num_columns) != num_columns || static_cast<size_t>(num_rows) != num_rows) {
+        throw ProtocolError("block declares " + std::to_string(num_columns) + " columns and " + std::to_string(num_rows) +
+                            " rows, more than size_t can hold");
+    }
 
     CreateColumnByTypeSettings create_column_settings;
     create_column_settings.low_cardinality_as_wrapped_column = options_.backward_compatibility_lowcardinality_as_wrapped_column;

clickhouse/columns/enum.cpp

@@ -81,8 +81,7 @@ void ColumnEnum<T>::Append(ColumnRef column) {
 
 template <typename T>
 bool ColumnEnum<T>::LoadBody(InputStream* input, size_t rows) {
-    data_.resize(rows);
-    return WireFormat::ReadBytes(*input, data_.data(), data_.size() * sizeof(T));
+    return WireFormat::ReadPODArray(*input, &data_, rows);
 }
 
 template <typename T>

clickhouse/columns/lowcardinality.cpp

@@ -452,7 +452,7 @@ auto Load(ColumnRef new_dictionary_column, InputStream& input, size_t rows) {
     if (number_of_rows != rows)
         throw AssertionError("LowCardinality column must be read in full.");
 
-    new_index_column->LoadBody(&input, number_of_rows);
+    if (!new_index_column->LoadBody(&input, number_of_rows)) throw ProtocolError("Failed to read index column.");
 
     if (auto nullable = new_dictionary_column->As<ColumnNullable>()) {
         nullable->Append(true);

clickhouse/columns/numeric.cpp

@@ -76,9 +76,7 @@ void ColumnVector<T>::Append(ColumnRef column) {
 
 template <typename T>
 bool ColumnVector<T>::LoadBody(InputStream* input, size_t rows) {
-    data_.resize(rows);
-
-    return WireFormat::ReadBytes(*input, data_.data(), data_.size() * sizeof(T));
+    return WireFormat::ReadPODArray(*input, &data_, rows);
 }
 
 template <typename T>

clickhouse/columns/string.cpp

@@ -8,6 +8,18 @@ namespace {
 
 constexpr size_t DEFAULT_BLOCK_SIZE = 4096;
 
+// Upper bound on what LoadBody allocates ahead of the data, since row counts and
+// string lengths come from the peer.
+constexpr size_t MAX_PREALLOCATED_BYTES = 1 << 20;
+
+// Grows `items` eightfold rather than twofold so that loading many rows copies little.
+void AppendItem(std::vector<std::string_view>& items, size_t rows, std::string_view item) {
+    if (items.size() == items.capacity()) {
+        items.reserve(clickhouse::WireFormat::NextReadCapacity(items.size(), rows, size_t(1)));
+    }
+    items.emplace_back(item);
+}
+
 template <typename Container>
 size_t ComputeTotalSize(const Container & strings, size_t begin = 0, size_t len = -1) {
     size_t result = 0;
@@ -78,12 +90,7 @@ void ColumnFixedString::Append(ColumnRef column) {
 }
 
 bool ColumnFixedString::LoadBody(InputStream * input, size_t rows) {
-    data_.resize(string_size_ * rows);
-    if (!WireFormat::ReadBytes(*input, &data_[0], data_.size())) {
-        return false;
-    }
-
-    return true;
+    return WireFormat::ReadPODArray(*input, &data_, string_size_ * rows);
 }
 
 void ColumnFixedString::SaveBody(OutputStream* output) {
@@ -134,6 +141,16 @@ struct ColumnString::Block
         data_(new CharT[capacity])
     {}
 
+    // Moves the data, so it must not be called once a string_view points into the block.
+    void Grow(size_t new_capacity) {
+        std::unique_ptr<CharT[]> new_data(new CharT[new_capacity]);
+        if (size > 0) {
+            memcpy(new_data.get(), data_.get(), size);
+        }
+        data_    = std::move(new_data);
+        capacity = new_capacity;
+    }
+
     inline auto GetAvailable() const {
         return capacity - size;
     }
@@ -163,7 +180,7 @@ struct ColumnString::Block
     }
 
     size_t size;
-    const size_t capacity;
+    size_t capacity;
     std::unique_ptr<CharT[]> data_;
 };
 
@@ -274,7 +291,7 @@ bool ColumnString::LoadBody(InputStream* input, size_t rows) {
     decltype(items_) new_items;
     decltype(blocks_) new_blocks;
 
-    new_items.reserve(rows);
+    new_items.reserve(std::min(rows, MAX_PREALLOCATED_BYTES / sizeof(std::string_view)));
 
     // Suboptimzal if the first row string is >DEFAULT_BLOCK_SIZE, but that must be a very rare case.
     Block * block = &new_blocks.emplace_back(DEFAULT_BLOCK_SIZE);
@@ -284,13 +301,31 @@ bool ColumnString::LoadBody(InputStream* input, size_t rows) {
         if (!WireFormat::ReadUInt64(*input, &len))
             return false;
 
+        if (len > MAX_PREALLOCATED_BYTES) {
+            if (len > SIZE_MAX) return false;
+
+            // A block of its own, grown as the bytes arrive.
+            block = &new_blocks.emplace_back(MAX_PREALLOCATED_BYTES);
+            while (block->size < len) {
+                if (block->GetAvailable() == 0)
+                    block->Grow(WireFormat::NextReadCapacity(block->capacity, static_cast<size_t>(len), size_t(0)));
+
+                const size_t chunk = std::min<size_t>(static_cast<size_t>(len) - block->size, block->GetAvailable());
+                if (!WireFormat::ReadBytes(*input, block->GetCurrentWritePos(), chunk)) return false;
+                block->size += chunk;
+            }
+
+            AppendItem(new_items, rows, std::string_view(block->data_.get(), block->size));
+            continue;
+        }
+
         if (len > block->GetAvailable())
             block = &new_blocks.emplace_back(std::max<size_t>(DEFAULT_BLOCK_SIZE, len));
 
         if (!WireFormat::ReadBytes(*input, block->GetCurrentWritePos(), len))
             return false;
 
-        new_items.emplace_back(block->ConsumeTailAsStringViewUnsafe(len));
+        AppendItem(new_items, rows, block->ConsumeTailAsStringViewUnsafe(len));
     }
 
     items_.swap(new_items);

ut/BUILD.bazel

@@ -51,6 +51,7 @@ cc_test(
         "types_ut.cpp",
         "utils_ut.cpp",
         # Test entry point and shared support code.
+        "allocation_tracker.cpp",
         "main.cpp",
         "tcp_server.cpp",
         "tcp_server.h",

ut/CMakeLists.txt

@@ -31,6 +31,7 @@ SET ( clickhouse-cpp-ut-src
 
     bignum_samples.cpp
     utils.cpp
+    allocation_tracker.cpp
     value_generators.cpp
     low_cardinality_nullable_tests.cpp
 

ut/allocation_tracker.cpp

@@ -0,0 +1,51 @@
+#include <clickhouse/columns/numeric.h>
+
+#include <cstdlib>
+#include <new>
+
+#include "utils.h"
+
+namespace {
+thread_local bool track_allocations    = false;
+thread_local size_t largest_allocation = 0;
+}  // namespace
+
+// The default operator new[], operator delete[], and nothrow forms forward to these.
+void* operator new(std::size_t size) {
+    if (track_allocations && size > largest_allocation) {
+        largest_allocation = size;
+    }
+    if (void* ptr = std::malloc(size ? size : 1)) {
+        return ptr;
+    }
+    throw std::bad_alloc();
+}
+
+void operator delete(void* ptr) noexcept {
+    std::free(ptr);
+}
+
+void operator delete(void* ptr, std::size_t) noexcept {
+    std::free(ptr);
+}
+
+LargestAllocationScope::LargestAllocationScope() {
+    largest_allocation = 0;
+    track_allocations  = true;
+}
+
+LargestAllocationScope::~LargestAllocationScope() {
+    track_allocations = false;
+}
+
+size_t LargestAllocationScope::Get() const {
+    return largest_allocation;
+}
+
+bool LargestAllocationScope::ObservesLibraryAllocations() {
+    constexpr size_t rows = 1 << 20;
+    LargestAllocationScope allocations;
+    clickhouse::ColumnUInt64 column;
+    column.Reserve(rows);
+    return allocations.Get() >= rows * sizeof(uint64_t);
+}

ut/client_protocol_ut.cpp

@@ -238,3 +238,29 @@ TEST(ClientProtocol, WellFormedSelectResponseSucceeds) {
     }));
     EXPECT_EQ(blocks, 1u);
 }
+
+TEST(ClientProtocol, RowCountBeyondSizeTThrowsProtocolError) {
+    if (sizeof(size_t) >= sizeof(uint64_t)) {
+        GTEST_SKIP() << "every uint64_t row count fits in size_t here";
+    }
+
+    std::vector<uint8_t> script         = kServerHello;
+    const std::vector<uint8_t> response = {
+        0x01,                                // packet type: ServerCodes::Data (varint)
+        0x01,                                // num_columns = 1 (varint)
+        0x80, 0x80, 0x80, 0x80, 0x80, 0x40,  // num_rows = 2^41 (varint), 0 once truncated to 32 bits
+        0x01, 0x63,                          // column name: "c"
+        0x05, 0x55, 0x49, 0x6E, 0x74, 0x38,  // column type: "UInt8"
+        0x05,                                // packet type: ServerCodes::EndOfStream (varint)
+    };
+    script.insert(script.end(), response.begin(), response.end());
+
+    Client client(ScriptedClientOptions(), std::make_unique<ScriptedSocketFactory>(script));
+
+    try {
+        client.Select("SELECT 1", [](const Block&) {});
+        FAIL() << "expected ProtocolError";
+    } catch (const ProtocolError& e) {
+        EXPECT_TRUE(std::string(e.what()).find("more than size_t can hold") != std::string::npos) << "unexpected message: " << e.what();
+    }
+}