Tpcc import with arenas (#17333) (#20756)

Author: eivanov89

.github/last_commit.txt

@@ -1 +1 @@
-b221825ccb0f9e8813d5df19a7167ecdd5b9f098
+42da2fe1b7025f0f6c397b52f085d3ed54e5f690

include/ydb-cpp-sdk/client/value/value.h

@@ -4,6 +4,8 @@
 
 #include <util/datetime/base.h>
 
+#include <google/protobuf/arena.h>
+
 #include <optional>
 #include <memory>
 
@@ -528,7 +530,7 @@ class TValueBuilderBase : public TMoveOnly {
 protected:
     TValueBuilderBase(TValueBuilderBase&&);
 
-    TValueBuilderBase();
+    TValueBuilderBase(google::protobuf::Arena* arena = nullptr);
 
     TValueBuilderBase(const TType& type);
 
@@ -544,7 +546,7 @@ class TValueBuilderBase : public TMoveOnly {
 
 class TValueBuilder : public TValueBuilderBase<TValueBuilder> {
 public:
-    TValueBuilder();
+    TValueBuilder(google::protobuf::Arena* arena = nullptr);
 
     TValueBuilder(const TType& type);
 

src/client/value/value.cpp

@@ -2061,19 +2061,33 @@ class TValueBuilderImpl {
 public:
     TValueBuilderImpl()
         : TypeBuilder_()
+        , Arena(nullptr)
+        , ProtoValue_(ProtoValueHeap)
+    {
+        PushPath(ProtoValue_);
+    }
+
+    TValueBuilderImpl(google::protobuf::Arena* arena)
+        : TypeBuilder_()
+        , Arena(arena)
+        , ProtoValue_(*google::protobuf::Arena::CreateMessage<Ydb::Value>(Arena))
     {
         PushPath(ProtoValue_);
     }
 
     TValueBuilderImpl(const TType& type)
         : TypeBuilder_()
+        , Arena(nullptr)
+        , ProtoValue_(ProtoValueHeap)
     {
         PushPath(ProtoValue_);
         GetType().CopyFrom(type.GetProto());
     }
 
     TValueBuilderImpl(Ydb::Type& type, Ydb::Value& value)
         : TypeBuilder_(type)
+        , Arena(nullptr)
+        , ProtoValue_(ProtoValueHeap)
     {
         PushPath(value);
     }
@@ -2088,10 +2102,15 @@ class TValueBuilderImpl {
     TValue BuildValue() {
         CheckValue();
 
-        Ydb::Value value;
-        value.Swap(&ProtoValue_);
-
-        return TValue(TypeBuilder_.Build(), std::move(value));
+        if (Arena) {
+            auto* value = google::protobuf::Arena::CreateMessage<Ydb::Value>(Arena);
+            value->Swap(&ProtoValue_);
+            return TValue(TypeBuilder_.Build(), value);
+        } else {
+            Ydb::Value value;
+            value.Swap(&ProtoValue_);
+            return TValue(TypeBuilder_.Build(), std::move(value));
+        }
     }
 
     void Bool(bool value) {
@@ -2803,7 +2822,12 @@ class TValueBuilderImpl {
 private:
     //TTypeBuilder TypeBuilder_;
     TTypeBuilder::TImpl TypeBuilder_;
-    Ydb::Value ProtoValue_;
+    google::protobuf::Arena* Arena;
+    Ydb::Value ProtoValueHeap;
+
+    // either ProtoValueHeap or a reference to the arena allocated protobuf
+    Ydb::Value& ProtoValue_;
+
     std::map<const Ydb::StructType*, TMembersMap> StructsMap_;
 
     TStackVec<TProtoPosition, 8> Path_;
@@ -2819,8 +2843,8 @@ template<typename TDerived>
 TValueBuilderBase<TDerived>::~TValueBuilderBase() = default;
 
 template<typename TDerived>
-TValueBuilderBase<TDerived>::TValueBuilderBase()
-    : Impl_(new TValueBuilderImpl()) {}
+TValueBuilderBase<TDerived>::TValueBuilderBase(google::protobuf::Arena* arena)
+    : Impl_(new TValueBuilderImpl(arena)) {}
 
 template<typename TDerived>
 TValueBuilderBase<TDerived>::TValueBuilderBase(const TType& type)
@@ -3382,8 +3406,8 @@ template class TValueBuilderBase<TParamValueBuilder>;
 
 ////////////////////////////////////////////////////////////////////////////////
 
-TValueBuilder::TValueBuilder()
-    : TValueBuilderBase() {}
+TValueBuilder::TValueBuilder(google::protobuf::Arena* arena)
+    : TValueBuilderBase(arena) {}
 
 TValueBuilder::TValueBuilder(const TType& type)
     : TValueBuilderBase(type) {}