Add Arena option to BulkUpsert, add hidden --send-format option (#20061)

Co-authored-by: Vladislav Artiukhov <vladislav0art.work@gmail.com>

Author: pnv1

ydb/public/lib/ydb_cli/commands/ydb_service_import.cpp

@@ -6,7 +6,6 @@
 #include <ydb/public/lib/ydb_cli/common/print_operation.h>
 #include <ydb/public/lib/ydb_cli/common/interactive.h>
 #include <ydb/public/lib/ydb_cli/dump/files/files.h>
-#include <ydb/public/lib/ydb_cli/import/import.h>
 #include <ydb/library/backup/util.h>
 
 #include <util/string/builder.h>
@@ -302,6 +301,17 @@ void TCommandImportFromCsv::Config(TConfig& config) {
     config.Opts->AddLongOption("newline-delimited",
             "No newline characters inside records, enables some import optimizations (see docs)")
         .StoreTrue(&NewlineDelimited);
+    TStringStream description;
+    description << "Format that data will be serialized to before sending to YDB. Available options: ";
+    NColorizer::TColors colors = NColorizer::AutoColors(Cout);
+    description << "\n  " << colors.BoldColor() << "tvalue" << colors.OldColor()
+        << "\n    " << "A default YDB protobuf format";
+    description << "\n  " << colors.BoldColor() << "arrow" << colors.OldColor()
+        << "\n    " << "Apache Arrow format";
+    description << "\nDefault: " << colors.CyanColor() << "\"" << "tvalue" << "\"" << colors.OldColor() << ".";
+    config.Opts->AddLongOption("send-format", description.Str())
+        .RequiredArgument("STRING").StoreResult(&SendFormat)
+        .Hidden();
     if (InputFormat == EDataFormat::Csv) {
         config.Opts->AddLongOption("delimiter", "Field delimiter in rows")
             .RequiredArgument("STRING").StoreResult(&Delimiter).DefaultValue(Delimiter);
@@ -325,6 +335,7 @@ int TCommandImportFromCsv::Run(TConfig& config) {
     settings.HeaderRow(HeaderRow);
     settings.NullValue(NullValue);
     settings.Verbose(config.IsVerbose());
+    settings.SendFormat(SendFormat);
 
     if (Delimiter.size() != 1) {
         throw TMisuseException()

ydb/public/lib/ydb_cli/commands/ydb_service_import.h

@@ -7,6 +7,7 @@
 #include <ydb/public/lib/ydb_cli/common/aws.h>
 #include <ydb/public/lib/ydb_cli/common/format.h>
 #include <ydb/public/lib/ydb_cli/common/parseable_struct.h>
+#include <ydb/public/lib/ydb_cli/import/import.h>
 
 namespace NYdb::NConsoleClient {
 
@@ -86,6 +87,7 @@ class TCommandImportFromCsv : public TCommandImportFileBase {
     ui32 SkipRows = 0;
     bool Header = false;
     bool NewlineDelimited = true;
+    NConsoleClient::ESendFormat SendFormat = NConsoleClient::ESendFormat::Default;
 };
 
 class TCommandImportFromTsv : public TCommandImportFromCsv {

ydb/public/lib/ydb_cli/common/csv_parser.cpp

@@ -15,9 +15,7 @@ class TCsvToYdbConverter {
 public:
     explicit TCsvToYdbConverter(TTypeParser& parser, const std::optional<TString>& nullValue)
         : Parser(parser)
-        , NullValue(nullValue)
-    {
-    }
+        , NullValue(nullValue) {}
 
     template <class T, std::enable_if_t<std::is_integral_v<T> && std::is_signed_v<T>, std::nullptr_t> = nullptr>
     static i64 StringToArithmetic(const TString& token, size_t& cnt) {
@@ -165,7 +163,7 @@ class TCsvToYdbConverter {
         }
         case EPrimitiveType::Interval64:
             Builder.Interval64(GetArithmetic<i64>(token));
-            break;            
+            break;
         case EPrimitiveType::TzDate:
             Builder.TzDate(token);
             break;
@@ -441,7 +439,7 @@ TStringBuf Consume(NCsvFormat::CsvSplitter& splitter,
 
 TCsvParser::TCsvParser(TString&& headerRow, const char delimeter, const std::optional<TString>& nullValue,
                        const std::map<std::string, TType>* destinationTypes,
-                       const std::map<TString, TString>* paramSources) 
+                       const std::map<TString, TString>* paramSources)
     : HeaderRow(std::move(headerRow))
     , Delimeter(delimeter)
     , NullValue(nullValue)
@@ -454,7 +452,7 @@ TCsvParser::TCsvParser(TString&& headerRow, const char delimeter, const std::opt
 
 TCsvParser::TCsvParser(TVector<TString>&& header, const char delimeter, const std::optional<TString>& nullValue,
                        const std::map<std::string, TType>* destinationTypes,
-                       const std::map<TString, TString>* paramSources) 
+                       const std::map<TString, TString>* paramSources)
     : Header(std::move(header))
     , Delimeter(delimeter)
     , NullValue(nullValue)
@@ -529,41 +527,91 @@ void TCsvParser::BuildValue(TString& data, TValueBuilder& builder, const TType&
     builder.EndStruct();
 }
 
-TValue TCsvParser::BuildList(std::vector<TString>& lines, const TString& filename, std::optional<ui64> row) const {
+TValue TCsvParser::BuildList(const std::vector<TString>& lines, const TString& filename, std::optional<ui64> row) const {
     std::vector<std::unique_ptr<TTypeParser>> columnTypeParsers;
     columnTypeParsers.reserve(ResultColumnCount);
     for (const TType* type : ResultLineTypesSorted) {
         columnTypeParsers.push_back(std::make_unique<TTypeParser>(*type));
     }
-    
-    Ydb::Value listValue;
-    auto* listItems = listValue.mutable_items();
+
+    // Original path with local value object
+    Ydb::Value listProtoValue;
+    auto* listItems = listProtoValue.mutable_items();
     listItems->Reserve(lines.size());
-    for (auto& line : lines) {
-        NCsvFormat::CsvSplitter splitter(line, Delimeter);
-        TParseMetadata meta {row, filename};
-        auto* structItems = listItems->Add()->mutable_items();
-        structItems->Reserve(ResultColumnCount);
-        auto headerIt = Header.cbegin();
-        auto skipIt = SkipBitMap.begin();
-        auto typeParserIt = columnTypeParsers.begin();
-        do {
-            if (headerIt == Header.cend()) { // SkipBitMap has same size as Header
-                throw FormatError(yexception() << "Header contains less fields than data. Header: \"" << HeaderRow << "\", data: \"" << line << "\"", meta);
-            }
-            TStringBuf nextField = Consume(splitter, meta, *headerIt);
-            if (!*skipIt) {
-                *structItems->Add() = FieldToValue(*typeParserIt->get(), nextField, NullValue, meta, *headerIt).GetProto();
-                ++typeParserIt;
-            }
-            ++headerIt;
-            ++skipIt;
-        } while (splitter.Step());
+
+    for (const auto& line : lines) {
+        ProcessCsvLine(line, listItems, columnTypeParsers, row, filename);
         if (row.has_value()) {
             ++row.value();
         }
     }
-    return TValue(ResultListType.value(), std::move(listValue));
+
+    // Return a TValue that takes ownership via move
+    return TValue(ResultListType.value(), std::move(listProtoValue));
+}
+
+TValue TCsvParser::BuildListOnArena(
+    const std::vector<TString>& lines,
+    const TString& filename,
+    google::protobuf::Arena* arena,
+    std::optional<ui64> row
+) const {
+    Y_ASSERT(arena != nullptr);
+
+    std::vector<std::unique_ptr<TTypeParser>> columnTypeParsers;
+    columnTypeParsers.reserve(ResultColumnCount);
+    for (const TType* type : ResultLineTypesSorted) {
+        columnTypeParsers.push_back(std::make_unique<TTypeParser>(*type));
+    }
+
+    // allocate Ydb::Value on arena
+    Ydb::Value* listProtoValue = google::protobuf::Arena::CreateMessage<Ydb::Value>(arena);
+    auto* listItems = listProtoValue->mutable_items();
+    listItems->Reserve(lines.size());
+
+    for (const auto& line : lines) {
+        ProcessCsvLine(line, listItems, columnTypeParsers, row, filename);
+        if (row.has_value()) {
+            ++row.value();
+        }
+    }
+
+    // Return a TValue that references the arena-allocated message
+    return TValue(ResultListType.value(), listProtoValue);
+}
+
+// Helper method to process a single CSV line
+void TCsvParser::ProcessCsvLine(
+    const TString& line,
+    google::protobuf::RepeatedPtrField<Ydb::Value>* listItems,
+    const std::vector<std::unique_ptr<TTypeParser>>& columnTypeParsers,
+    std::optional<ui64> currentRow,
+    const TString& filename
+) const {
+    NCsvFormat::CsvSplitter splitter(line, Delimeter);
+    auto* structItems = listItems->Add()->mutable_items();
+    structItems->Reserve(ResultColumnCount);
+
+    const TParseMetadata meta {currentRow, filename};
+
+    auto headerIt = Header.cbegin();
+    auto skipIt = SkipBitMap.begin();
+    auto typeParserIt = columnTypeParsers.begin();
+
+    do {
+        if (headerIt == Header.cend()) { // SkipBitMap has same size as Header
+            throw FormatError(yexception() << "Header contains less fields than data. Header: \"" << HeaderRow << "\", data: \"" << line << "\"", meta);
+        }
+        TStringBuf nextField = Consume(splitter, meta, *headerIt);
+        if (!*skipIt) {
+            TValue builtValue = FieldToValue(*typeParserIt->get(), nextField, NullValue, meta, *headerIt);
+            *structItems->Add() = std::move(builtValue.GetProto());
+
+            ++typeParserIt;
+        }
+        ++headerIt;
+        ++skipIt;
+    } while (splitter.Step());
 }
 
 void TCsvParser::BuildLineType() {
@@ -607,5 +655,10 @@ const TVector<TString>& TCsvParser::GetHeader() {
     return Header;
 }
 
+const TString& TCsvParser::GetHeaderRow() const {
+    return HeaderRow;
+}
+
 }
 }
+

ydb/public/lib/ydb_cli/common/csv_parser.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/params/params.h>
+#include <google/protobuf/arena.h>
 
 #include <library/cpp/string_utils/csv/csv.h>
 
@@ -35,10 +36,20 @@ class TCsvParser {
 
     void BuildParams(TString& data, TParamsBuilder& builder, const TParseMetadata& meta) const;
     void BuildValue(TString& data, TValueBuilder& builder, const TType& type, const TParseMetadata& meta) const;
-    TValue BuildList(std::vector<TString>& lines, const TString& filename,
+
+    TValue BuildList(const std::vector<TString>& lines, const TString& filename,
                      std::optional<ui64> row = std::nullopt) const;
+
+    TValue BuildListOnArena(
+        const std::vector<TString>& lines,
+        const TString& filename,
+        google::protobuf::Arena* arena,
+        std::optional<ui64> row = std::nullopt
+    ) const;
+
     void BuildLineType();
     const TVector<TString>& GetHeader();
+    const TString& GetHeaderRow() const;
 
 private:
     TVector<TString> Header;
@@ -60,6 +71,15 @@ class TCsvParser {
     // Types of columns in a single row in resulting TValue.
     // Column order according to the header, though can have less elements than the Header
     std::vector<const TType*> ResultLineTypesSorted;
+
+    // Helper method to process a single line of CSV data
+    void ProcessCsvLine(
+        const TString& line,
+        google::protobuf::RepeatedPtrField<Ydb::Value>* listItems,
+        const std::vector<std::unique_ptr<TTypeParser>>& columnTypeParsers,
+        std::optional<ui64> currentRow,
+        const TString& filename
+    ) const;
 };
 
 }

ydb/public/lib/ydb_cli/common/csv_parser_ut.cpp

@@ -317,7 +317,7 @@ Y_UNIT_TEST_SUITE(YdbCliCsvParserTests) {
             {"col2", TTypeBuilder().BeginOptional().Primitive(EPrimitiveType::Int64).EndOptional().Build()},
             {"col3", TTypeBuilder().Primitive(EPrimitiveType::Bool).Build()},
         };
-        
+
         TString csvHeader = "col4,col3,col5,col1,col6";
         std::vector<TString> data = {
             "col4 unused value,true,col5 unused value,col1 value,col6 unused value"