branch-4.1: [refactor](variant) Make Variant metadata key registration allocation-safe

be/src/core/value/variant/variant_batch_builder.cpp

@@ -22,11 +22,11 @@
 #include <algorithm>
 #include <cstdint>
 #include <cstring>
-#include <deque>
 #include <functional>
 #include <limits>
 #include <ranges>
 #include <string>
+#include <type_traits>
 #include <utility>
 
 #include "common/exception.h"
@@ -36,6 +36,10 @@
 #include "core/value/variant/variant_parquet_encoding.h"
 #include "util/utf8_check.h"
 
+#ifdef BE_TEST
+#include "util/debug_points.h"
+#endif
+
 namespace doris {
 namespace {
 
@@ -130,9 +134,30 @@ void append_unsigned(String& output, uint64_t value, uint8_t width) {
     }
 }
 
-template <typename Container>
-void release_batch_container(Container& container) noexcept {
-    Container().swap(container);
+template <typename T>
+void release_batch_vector(DorisVector<T>& buffer) noexcept {
+    static_assert(noexcept(DorisVector<T>()));
+    static_assert(noexcept(std::declval<DorisVector<T>&>().swap(buffer)));
+    DorisVector<T>().swap(buffer);
+}
+
+using MetadataKeyBytes = DorisVector<char>;
+using MetadataKeys = DorisVector<MetadataKeyBytes>;
+
+// A key map entry points at the bytes owned by one MetadataKeyBytes. Growing MetadataKeys moves
+// the inner vectors, so that move must only transfer the byte buffer. For example, reallocating
+// 1024 keys to 2048 entries must not copy any key bytes or leave a map entry pointing at freed
+// storage.
+static_assert(std::is_nothrow_move_constructible_v<MetadataKeyBytes>,
+              "Variant metadata key bytes must move without allocating");
+
+// An empty key owns no bytes, but its map entry still needs a stable, non-null address for both
+// hashing and comparison. One sentinel byte serves every empty key.
+constexpr char EMPTY_METADATA_KEY = '\0';
+
+StringRef metadata_key_ref(const MetadataKeyBytes& bytes) noexcept {
+    DCHECK(bytes.empty() || bytes.data() != nullptr);
+    return {bytes.empty() ? &EMPTY_METADATA_KEY : bytes.data(), bytes.size()};
 }
 
 } // namespace
@@ -141,9 +166,13 @@ struct VariantMetadataBuilder::Impl {
     using KeyMapValue = std::pair<const StringRef, uint32_t>;
     using KeyMap = phmap::flat_hash_map<StringRef, uint32_t, StringRefHash, std::equal_to<>,
                                         CustomStdAllocator<KeyMapValue>>;
-    using Keys = std::deque<PaddedPODArray<char>, CustomStdAllocator<PaddedPODArray<char>>>;
+    // register_key() publishes a map entry before appending the key, so a failed map allocation
+    // must not leave a half-inserted entry behind. phmap only moves entries after allocating a
+    // larger table, and moving this value type cannot throw, so a rejected growth leaves the map
+    // unchanged.
+    static_assert(std::is_nothrow_move_constructible_v<KeyMapValue>);
 
-    Keys keys;
+    MetadataKeys keys;
     KeyMap key_to_temporary_id;
     DorisVector<uint32_t> temporary_to_final_id;
     PaddedPODArray<char> encoded;
@@ -177,16 +206,33 @@ uint32_t VariantMetadataBuilder::register_key(StringRef key) {
     }
 
     const auto temporary_id = static_cast<uint32_t>(_impl->keys.size());
-    _impl->keys.emplace_back();
+
+    // Prepare every fallible step first: copy the bytes into a local owner and make room in the
+    // outer key vector. Only then publish the map entry and append the owner, so a rejected
+    // allocation leaves both containers untouched instead of needing a catch block.
+    MetadataKeyBytes owned_key;
     if (key.size != 0) {
-        _impl->keys.back().insert(key.data, key.data + key.size);
+        owned_key.insert(owned_key.end(), key.data, key.data + key.size);
+    }
+    // reserve() allocates exactly the requested capacity, so reserving size+1 every time would
+    // reallocate and move all previous keys on every registration. Grow geometrically only when
+    // the vector is full; the append below then needs no allocation.
+    if (_impl->keys.size() == _impl->keys.capacity()) {
+        _impl->keys.reserve(_impl->keys.capacity() == 0 ? 1 : _impl->keys.capacity() * 2);
     }
-    const StringRef owned_key {_impl->keys.back().data(), _impl->keys.back().size()};
+
     const size_t old_capacity = _impl->key_to_temporary_id.capacity();
-    const auto [unused, inserted] = _impl->key_to_temporary_id.emplace(owned_key, temporary_id);
-    _impl->key_capacity_growths += old_capacity != _impl->key_to_temporary_id.capacity();
+    // The entry is visible while `owned_key` still owns the bytes; the move below transfers that
+    // same buffer into `keys`, so the stored StringRef stays valid.
+    const auto [unused, inserted] =
+            _impl->key_to_temporary_id.emplace(metadata_key_ref(owned_key), temporary_id);
     static_cast<void>(unused);
     DCHECK(inserted);
+
+    // Commit: the outer vector has capacity for one more key and the move only transfers the byte
+    // buffer, so nothing below this point can allocate or throw.
+    _impl->keys.emplace_back(std::move(owned_key));
+    _impl->key_capacity_growths += old_capacity != _impl->key_to_temporary_id.capacity();
     return temporary_id;
 }
 
@@ -211,8 +257,9 @@ void VariantMetadataBuilder::_abort_row() noexcept {
     DCHECK(_impl->row_active);
     DCHECK_LE(_impl->row_start_key_count, _impl->keys.size());
     while (_impl->keys.size() > _impl->row_start_key_count) {
-        const StringRef key {_impl->keys.back().data(), _impl->keys.back().size()};
-        const size_t erased = _impl->key_to_temporary_id.erase(key);
+        // Erase the map entry before releasing the bytes it points at; neither step allocates.
+        const size_t erased =
+                _impl->key_to_temporary_id.erase(metadata_key_ref(_impl->keys.back()));
         DCHECK_EQ(erased, 1);
         _impl->keys.pop_back();
     }
@@ -224,12 +271,13 @@ void VariantMetadataBuilder::_reserve_keys(size_t count) {
         throw Exception(ErrorCode::INVALID_ARGUMENT,
                         "Cannot reserve Variant metadata keys after seal");
     }
+    _impl->keys.reserve(count);
     _impl->key_to_temporary_id.reserve(count);
 }
 
 StringRef VariantMetadataBuilder::_temporary_key(uint32_t temporary_id) const noexcept {
     DCHECK_LT(temporary_id, _impl->keys.size());
-    return {_impl->keys[temporary_id].data(), _impl->keys[temporary_id].size()};
+    return metadata_key_ref(_impl->keys[temporary_id]);
 }
 
 size_t VariantMetadataBuilder::_key_capacity() const noexcept {
@@ -255,8 +303,8 @@ void VariantMetadataBuilder::seal() {
         sorted_temporary_ids.push_back(temporary_id);
     }
     std::ranges::sort(sorted_temporary_ids, [this](uint32_t left, uint32_t right) {
-        const StringRef left_key {_impl->keys[left].data(), _impl->keys[left].size()};
-        const StringRef right_key {_impl->keys[right].data(), _impl->keys[right].size()};
+        const StringRef left_key = metadata_key_ref(_impl->keys[left]);
+        const StringRef right_key = metadata_key_ref(_impl->keys[right]);
         return left_key.compare(right_key) < 0;
     });
 
@@ -297,8 +345,12 @@ void VariantMetadataBuilder::seal() {
         append_unsigned(_impl->encoded, offset, offset_width);
     }
     for (uint32_t temporary_id : sorted_temporary_ids) {
-        const auto& key = _impl->keys[temporary_id];
-        _impl->encoded.insert(key.data(), key.data() + key.size());
+        const StringRef key = metadata_key_ref(_impl->keys[temporary_id]);
+        // PaddedPODArray::insert must not receive a null source pointer, and an empty key has no
+        // bytes to append.
+        if (key.size != 0) {
+            _impl->encoded.insert(key.data, key.data + key.size);
+        }
     }
     DCHECK_EQ(_impl->encoded.size(), encoded_size);
     _impl->sealed_key_count = count;
@@ -342,9 +394,10 @@ StringRef VariantMetadataBuilder::encoded_metadata() const {
 PaddedPODArray<char> VariantMetadataBuilder::_take_encoded_metadata() noexcept {
     DCHECK(_impl->sealed);
     PaddedPODArray<char> encoded = std::move(_impl->encoded);
-    release_batch_container(_impl->keys);
-    release_batch_container(_impl->key_to_temporary_id);
-    release_batch_container(_impl->temporary_to_final_id);
+    // Destroy the scratch state directly: the encoded bytes were moved out first, and destroying
+    // the key map and key vectors only releases memory. Rebuilding empty containers here would
+    // allocate inside a noexcept cleanup path.
+    _impl.reset();
     return encoded;
 }
 
@@ -1237,14 +1290,14 @@ class VariantCollectionCore {
     }
 
     void release_scratch() noexcept {
-        release_batch_container(scalar_bytes);
-        release_batch_container(nodes);
-        release_batch_container(containers);
-        release_batch_container(children);
-        release_batch_container(scope_stack);
-        release_batch_container(object_id_scratch);
-        release_batch_container(container_plans);
-        release_batch_container(planned_object_children);
+        release_batch_vector(scalar_bytes);
+        release_batch_vector(nodes);
+        release_batch_vector(containers);
+        release_batch_vector(children);
+        release_batch_vector(scope_stack);
+        release_batch_vector(object_id_scratch);
+        release_batch_vector(container_plans);
+        release_batch_vector(planned_object_children);
         previous_object_tokens = nullptr;
         pending_object_tokens = nullptr;
         active_container = INVALID_INDEX;
@@ -1575,12 +1628,13 @@ VariantBatchBuilder VariantBatchBuilder::finish_batch() {
     _impl->refresh_counters(_impl->metadata._key_capacity(),
                             _impl->metadata._key_capacity_growths(), _impl->metadata.num_keys());
     _impl->counters_finalized = true;
+    DBUG_EXECUTE_IF("VariantBatchBuilder.before_release_scratch", DBUG_RUN_CALLBACK());
 #endif
     PaddedPODArray<char> metadata = _impl->metadata._take_encoded_metadata();
     _impl->collection.release_scratch();
-    release_batch_container(_impl->roots);
-    release_batch_container(_impl->previous_object_tokens);
-    release_batch_container(_impl->pending_object_tokens);
+    release_batch_vector(_impl->roots);
+    release_batch_vector(_impl->previous_object_tokens);
+    release_batch_vector(_impl->pending_object_tokens);
     return {std::move(metadata), std::move(values), std::move(offsets)};
 }
 

be/src/storage/segment/variant/v2/variant_path_builder.cpp

@@ -52,6 +52,9 @@
 #include "core/value/vdatetime_value.h"
 #include "exec/common/variant_util.h"
 #include "exprs/function/parse/variant_jsonb_parse.h"
+#ifdef BE_TEST
+#include "util/debug_points.h"
+#endif
 #include "util/jsonb_document.h"
 #include "util/jsonb_writer.h"
 
@@ -1125,18 +1128,27 @@ Status VariantPathBuilder::append(VariantRef value, size_t row) {
         }
 
         try {
+#ifdef BE_TEST
+            DBUG_EXECUTE_IF("VariantPathBuilder.before_append_value", DBUG_RUN_CALLBACK());
+#endif
             auto& nullable = assert_cast<ColumnNullable&>(*_impl->column);
             append_value(value, _impl->type, &nullable.get_nested_column());
-            nullable.get_null_map_data().push_back(0);
-        } catch (const Exception&) {
-            if (is_array || _impl->type->get_primitive_type() == TYPE_JSONB) {
+        } catch (const Exception& exception) {
+            if (exception.code() == ErrorCode::MEM_ALLOC_FAILED ||
+                exception.code() == ErrorCode::MEM_LIMIT_EXCEEDED || is_array ||
+                _impl->type->get_primitive_type() == TYPE_JSONB) {
                 throw;
             }
             RETURN_IF_ERROR(_impl->promote(jsonb_type(), false));
             auto& nullable = assert_cast<ColumnNullable&>(*_impl->column);
             append_value(value, _impl->type, &nullable.get_nested_column());
-            nullable.get_null_map_data().push_back(0);
         }
+
+#ifdef BE_TEST
+        DBUG_EXECUTE_IF("VariantPathBuilder.before_append_null_map", DBUG_RUN_CALLBACK());
+#endif
+        auto& nullable = assert_cast<ColumnNullable&>(*_impl->column);
+        nullable.get_null_map_data().push_back(0);
         _impl->rowids.push_back(static_cast<uint32_t>(row));
         _impl->logical_rows = row + 1;
 #if defined(BE_TEST) && !defined(BE_BENCHMARK)

be/test/storage/variant/variant_column_writer_reader_test.cpp

@@ -20,6 +20,7 @@
 #include <atomic>
 #include <cstdlib>
 #include <cstring>
+#include <functional>
 #include <mutex>
 #include <optional>
 #include <ranges>
@@ -284,6 +285,36 @@ TEST(VariantPathBuilderTest, PromotesValuesAndMaterializesMissingRows) {
     }
 }
 
+TEST(VariantPathBuilderTest, MemoryFailureBypassesJsonbFallback) {
+    VariantBatchBuilder value_builder;
+    auto row = value_builder.begin_row();
+    row.add_int(42);
+    row.finish();
+    auto values = value_builder.finish_batch();
+    const auto value = values.value_at(0);
+
+    for (const auto point :
+         {"VariantPathBuilder.before_append_value", "VariantPathBuilder.before_append_null_map"}) {
+        SCOPED_TRACE(point);
+        segment_v2::VariantPathBuilder builder(PathInData("metric"));
+        ASSERT_TRUE(builder.append(value, 0).ok());
+        const bool old_enable_debug_points = config::enable_debug_points;
+        config::enable_debug_points = true;
+        DebugPoints::instance()->add_with_callback(
+                point, std::function<void()>([] {
+                    throw Exception(ErrorCode::MEM_ALLOC_FAILED, "Injected allocation failure");
+                }));
+        Defer cleanup {[&] {
+            DebugPoints::instance()->remove(point);
+            config::enable_debug_points = old_enable_debug_points;
+        }};
+        const Status status = builder.append(value, 1);
+        EXPECT_TRUE(status.is<ErrorCode::MEM_ALLOC_FAILED>()) << status;
+        EXPECT_EQ(builder.promotion_count(), 0);
+        EXPECT_EQ(remove_nullable(builder.type())->get_primitive_type(), TYPE_BIGINT);
+    }
+}
+
 TEST(VariantPathBuilderTest, UsesBigintForEncodedIntegersWithoutPromotion) {
     VariantBatchBuilder value_builder;
     auto tiny_row = value_builder.begin_row();