[YDB CLI][CSV] Optimization: Use Apache Arrow and Protobuf Arena API as fallback in CSV upload module

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;
@@ -470,7 +468,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)
@@ -483,7 +481,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)
@@ -558,41 +556,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 valueProto;
+    auto* listItems = valueProto.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 a TValue that takes ownership via move
+    return TValue(ResultListType.value(), std::move(valueProto));
+}
+
+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* value = google::protobuf::Arena::CreateMessage<Ydb::Value>(arena);
+    auto* listItems = value->mutable_items();
+    listItems->Reserve(lines.size());
+
+    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 references the arena-allocated message
+    return TValue(ResultListType.value(), value);
+}
+
+// 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).ExtractProto();
+
+            ++typeParserIt;
+        }
+        ++headerIt;
+        ++skipIt;
+    } while (splitter.Step());
 }
 
 void TCsvParser::BuildLineType() {
@@ -761,5 +809,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>
 
@@ -66,10 +67,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;
     void ParseLineTypes(TString& line, TPossibleTypes& possibleTypes, const TParseMetadata& meta);
 
 private:
@@ -92,6 +103,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"