Cli parquet

ydb/apps/ydb/CHANGELOG.md

@@ -1,3 +1,4 @@
+* Use parquet format instead of CSV to fill tables in `ydb workload` benchmarks
 * Made `--consumer` flag in `ydb topic read` command optional. Now if this flag is not specified, reading is performed in no-consumer mode. In this mode partition IDs should be specified with `--partition-ids` option. 
 * Fixed a bug in `ydb import file csv` where multiple columns with escaped quotes in the same row were parsed incorrectly
 * Truncate query results output in benchmarks

ydb/core/grpc_services/rpc_load_rows.cpp

@@ -460,46 +460,13 @@ class TUploadColumnsRPCPublic : public NTxProxy::TUploadRowsBase<NKikimrServices
             case EUploadSource::CSV:
             {
                 auto& data = GetSourceData();
-                auto& cvsSettings = GetCsvSettings();
-                ui32 skipRows = cvsSettings.skip_rows();
-                auto& delimiter = cvsSettings.delimiter();
-                auto& nullValue = cvsSettings.null_value();
-                bool withHeader = cvsSettings.header();
-
-                auto reader = NFormats::TArrowCSV::Create(SrcColumns, withHeader, NotNullColumns);
+                auto& csvSettings = GetCsvSettings();
+                auto reader = NFormats::TArrowCSVScheme::Create(SrcColumns, csvSettings.header(), NotNullColumns);
                 if (!reader.ok()) {
                     errorMessage = reader.status().ToString();
                     return false;
                 }
-                const auto& quoting = cvsSettings.quoting();
-                if (quoting.quote_char().length() > 1) {
-                    errorMessage = TStringBuilder() << "Wrong quote char '" << quoting.quote_char() << "'";
-                    return false;
-                }
-                const char qchar = quoting.quote_char().empty() ? '"' : quoting.quote_char().front();
-                reader->SetQuoting(!quoting.disabled(), qchar, !quoting.double_quote_disabled());
-                reader->SetSkipRows(skipRows);
-
-                if (!delimiter.empty()) {
-                    if (delimiter.size() != 1) {
-                        errorMessage = TStringBuilder() << "Wrong delimiter '" << delimiter << "'";
-                        return false;
-                    }
-
-                    reader->SetDelimiter(delimiter[0]);
-                }
-
-                if (!nullValue.empty()) {
-                    reader->SetNullValue(nullValue);
-                }
-
-                if (data.size() > NFormats::TArrowCSV::DEFAULT_BLOCK_SIZE) {
-                    ui32 blockSize = NFormats::TArrowCSV::DEFAULT_BLOCK_SIZE;
-                    blockSize *= data.size() / blockSize + 1;
-                    reader->SetBlockSize(blockSize);
-                }
-
-                Batch = reader->ReadSingleBatch(data, errorMessage);
+                Batch = reader->ReadSingleBatch(data, csvSettings, errorMessage);
                 if (!Batch) {
                     return false;
                 }

ydb/core/io_formats/arrow/csv_arrow.cpp → ydb/core/io_formats/arrow/csv_arrow/csv_arrow.cpp

@@ -1,8 +1,7 @@
 #include "csv_arrow.h"
 
-#include <ydb/core/formats/arrow/arrow_helpers.h>
-#include <ydb/core/formats/arrow/serializer/stream.h>
-
+#include <contrib/libs/apache/arrow/cpp/src/arrow/array.h>
+#include <contrib/libs/apache/arrow/cpp/src/arrow/array/builder_primitive.h>
 #include <contrib/libs/apache/arrow/cpp/src/arrow/record_batch.h>
 #include <contrib/libs/apache/arrow/cpp/src/arrow/util/value_parsing.h>
 #include <util/string/join.h>
@@ -43,29 +42,6 @@ class TimestampIntParser: public arrow::TimestampParser {
 
 }
 
-arrow::Result<TArrowCSV> TArrowCSV::Create(const TVector<std::pair<TString, NScheme::TTypeInfo>>& columns, bool header, const std::set<std::string>& notNullColumns) {
-        TVector<TString> errors;
-        TColummns convertedColumns;
-        convertedColumns.reserve(columns.size());
-        for (auto& [name, type] : columns) {
-            const auto arrowType = NArrow::GetArrowType(type);
-            if (!arrowType.ok()) {
-                errors.emplace_back("column " + name + ": " + arrowType.status().ToString());
-                continue;
-            }
-            const auto csvArrowType = NArrow::GetCSVArrowType(type);
-            if (!csvArrowType.ok()) {
-                errors.emplace_back("column " + name + ": " + csvArrowType.status().ToString());
-                continue;
-            }
-            convertedColumns.emplace_back(TColumnInfo{name, *arrowType, *csvArrowType});
-        }
-        if (!errors.empty()) {
-            return arrow::Status::TypeError(ErrorPrefix() + "columns errors: " + JoinSeq("; ", errors));
-        }
-        return TArrowCSV(convertedColumns, header, notNullColumns);
-}
-
 TArrowCSV::TArrowCSV(const TColummns& columns, bool header, const std::set<std::string>& notNullColumns)
     : ReadOptions(arrow::csv::ReadOptions::Defaults())
     , ParseOptions(arrow::csv::ParseOptions::Defaults())
@@ -107,6 +83,27 @@ TArrowCSV::TArrowCSV(const TColummns& columns, bool header, const std::set<std::
     SetNullValue(); // set default null value
 }
 
+namespace {
+
+    template<class TBuilder, class TOriginalArray>
+    std::shared_ptr<arrow::Array> ConvertArray(std::shared_ptr<arrow::ArrayData> data, ui64 dev) {
+        auto originalArr = std::make_shared<TOriginalArray>(data);
+        TBuilder aBuilder;
+        Y_ABORT_UNLESS(aBuilder.Reserve(originalArr->length()).ok());
+        for (long i = 0; i < originalArr->length(); ++i) {
+            if (originalArr->IsNull(i)) {
+                Y_ABORT_UNLESS(aBuilder.AppendNull().ok());
+            } else {
+                aBuilder.UnsafeAppend(originalArr->Value(i) / dev);
+            }
+        }
+        auto res = aBuilder.Finish();
+        Y_ABORT_UNLESS(res.ok());
+        return *res;
+    }
+
+}
+
 std::shared_ptr<arrow::RecordBatch> TArrowCSV::ConvertColumnTypes(std::shared_ptr<arrow::RecordBatch> parsedBatch) const {
     if (!parsedBatch) {
         return nullptr;
@@ -134,59 +131,20 @@ std::shared_ptr<arrow::RecordBatch> TArrowCSV::ConvertColumnTypes(std::shared_pt
         if (fArr->type()->Equals(originalType)) {
             resultColumns.emplace_back(fArr);
         } else if (fArr->type()->id() == arrow::TimestampType::type_id) {
-            arrow::Result<std::shared_ptr<arrow::Array>> arrResult;
-            {
-                std::shared_ptr<arrow::TimestampArray> i64Arr = std::make_shared<arrow::TimestampArray>(fArr->data());
-                if (originalType->id() == arrow::UInt16Type::type_id) {
-                    arrow::UInt16Builder aBuilder;
-                    Y_ABORT_UNLESS(aBuilder.Reserve(parsedBatch->num_rows()).ok());
-                    for (long i = 0; i < parsedBatch->num_rows(); ++i) {
-                        if (i64Arr->IsNull(i)) {
-                            Y_ABORT_UNLESS(aBuilder.AppendNull().ok());
-                        } else {
-                            aBuilder.UnsafeAppend(i64Arr->Value(i) / 86400ull);
-                        }
-                    }
-                    arrResult = aBuilder.Finish();
-                } else if (originalType->id() == arrow::UInt32Type::type_id) {
-                    arrow::UInt32Builder aBuilder;
-                    Y_ABORT_UNLESS(aBuilder.Reserve(parsedBatch->num_rows()).ok());
-                    for (long i = 0; i < parsedBatch->num_rows(); ++i) {
-                        if (i64Arr->IsNull(i)) {
-                            Y_ABORT_UNLESS(aBuilder.AppendNull().ok());
-                        } else {
-                            aBuilder.UnsafeAppend(i64Arr->Value(i));
-                        }
-                    }
-                    arrResult = aBuilder.Finish();
-                } else if (originalType->id() == arrow::Int32Type::type_id) {
-                    arrow::Int32Builder aBuilder;
-                    Y_ABORT_UNLESS(aBuilder.Reserve(parsedBatch->num_rows()).ok());
-                    for (long i = 0; i < parsedBatch->num_rows(); ++i) {
-                        if (i64Arr->IsNull(i)) {
-                            Y_ABORT_UNLESS(aBuilder.AppendNull().ok());
-                        } else {
-                            aBuilder.UnsafeAppend(i64Arr->Value(i) / 86400);
-                        }
-                    }
-                    arrResult = aBuilder.Finish();
-                } else if (originalType->id() == arrow::Int64Type::type_id) {
-                    arrow::Int64Builder aBuilder;
-                    Y_ABORT_UNLESS(aBuilder.Reserve(parsedBatch->num_rows()).ok());
-                    for (long i = 0; i < parsedBatch->num_rows(); ++i) {
-                        if (i64Arr->IsNull(i)) {
-                            Y_ABORT_UNLESS(aBuilder.AppendNull().ok());
-                        } else {
-                            aBuilder.UnsafeAppend(i64Arr->Value(i));
-                        }
-                    }
-                    arrResult = aBuilder.Finish();
-                } else {
+            resultColumns.emplace_back([originalType, fArr]() {
+                switch (originalType->id()) {
+                case arrow::UInt16Type::type_id: // Date
+                    return ConvertArray<arrow::UInt16Builder, arrow::TimestampArray>(fArr->data(), 86400);
+                case arrow::UInt32Type::type_id: // Datetime
+                    return ConvertArray<arrow::UInt32Builder, arrow::TimestampArray>(fArr->data(), 1);
+                case arrow::Int32Type::type_id: // Date32
+                    return ConvertArray<arrow::Int32Builder, arrow::TimestampArray>(fArr->data(), 86400);
+                case arrow::Int64Type::type_id:// Datetime64, Timestamp64
+                    return ConvertArray<arrow::Int64Builder, arrow::TimestampArray>(fArr->data(), 1);
+                default:
                     Y_ABORT_UNLESS(false);
                 }
-            }
-            Y_ABORT_UNLESS(arrResult.ok());
-            resultColumns.emplace_back(*arrResult);
+            }());
         } else {
             Y_ABORT_UNLESS(false);
         }
@@ -204,7 +162,7 @@ std::shared_ptr<arrow::RecordBatch> TArrowCSV::ReadNext(const TString& csv, TStr
             return {};
         }
 
-        auto buffer = std::make_shared<NArrow::NSerialization::TBufferOverString>(csv);
+        auto buffer = std::make_shared<arrow::Buffer>(arrow::util::string_view(csv.c_str(), csv.length()));
         auto input = std::make_shared<arrow::io::BufferReader>(buffer);
         auto res = arrow::csv::StreamingReader::Make(arrow::io::default_io_context(), input,
                                                      ReadOptions, ParseOptions, ConvertOptions);
@@ -249,11 +207,9 @@ std::shared_ptr<arrow::RecordBatch> TArrowCSV::ReadNext(const TString& csv, TStr
         return {};
     }
 
-    if (batch && ResultColumns.size()) {
-        batch = NArrow::TColumnOperator().ErrorIfAbsent().Extract(batch, ResultColumns);
-        if (!batch) {
-            errString = ErrorPrefix() + "not all result columns present";
-        }
+    if (batch && ResultColumns.size() && batch->schema()->fields().size() != ResultColumns.size()) {
+        errString = ErrorPrefix() + "not all result columns present";
+        batch.reset();
     }
     return batch;
 }
@@ -279,5 +235,34 @@ std::shared_ptr<arrow::RecordBatch> TArrowCSV::ReadSingleBatch(const TString& cs
     }
     return batch;
 }
+std::shared_ptr<arrow::RecordBatch> TArrowCSV::ReadSingleBatch(const TString& csv, const Ydb::Formats::CsvSettings& csvSettings, TString& errString) {
+    const auto& quoting = csvSettings.quoting();
+    if (quoting.quote_char().length() > 1) {
+        errString = ErrorPrefix() + "Wrong quote char '" + quoting.quote_char() + "'";
+        return {};
+    }
+
+    const char qchar = quoting.quote_char().empty() ? '"' : quoting.quote_char().front();
+    SetQuoting(!quoting.disabled(), qchar, !quoting.double_quote_disabled());
+    if (csvSettings.delimiter()) {
+        if (csvSettings.delimiter().size() != 1) {
+            errString = ErrorPrefix() + "Invalid delimitr in csv: " + csvSettings.delimiter();
+            return {};
+        }
+        SetDelimiter(csvSettings.delimiter().front());
+    }
+    SetSkipRows(csvSettings.skip_rows());
+
+    if (csvSettings.null_value()) {
+        SetNullValue(csvSettings.null_value());
+    }
+
+    if (csv.size() > NKikimr::NFormats::TArrowCSV::DEFAULT_BLOCK_SIZE) {
+        ui32 blockSize = NKikimr::NFormats::TArrowCSV::DEFAULT_BLOCK_SIZE;
+        blockSize *= csv.size() / blockSize + 1;
+        SetBlockSize(blockSize);
+    }
+    return ReadSingleBatch(csv, errString);
+}
 
 }

ydb/core/io_formats/arrow/csv_arrow.h → ydb/core/io_formats/arrow/csv_arrow/csv_arrow.h

@@ -1,22 +1,23 @@
 #pragma once
 
-#include <ydb/core/scheme_types/scheme_type_info.h>
-
+#include <ydb/public/api/protos/ydb_formats.pb.h>
 #include <contrib/libs/apache/arrow/cpp/src/arrow/csv/api.h>
 #include <contrib/libs/apache/arrow/cpp/src/arrow/io/api.h>
+#include <util/generic/string.h>
+#include <util/generic/vector.h>
+#include <set>
+#include <vector>
+#include <unordered_map>
 
 namespace NKikimr::NFormats {
 
 class TArrowCSV {
 public:
     static constexpr ui32 DEFAULT_BLOCK_SIZE = 1024 * 1024;
 
-    /// If header is true read column names from first line after skipRows. Parse columns as strings in this case.
-    /// @note It's possible to skip header with skipRows and use typed columns instead.
-    static arrow::Result<TArrowCSV> Create(const TVector<std::pair<TString, NScheme::TTypeInfo>>& columns, bool header = false, const std::set<std::string>& notNullColumns = {});
-
     std::shared_ptr<arrow::RecordBatch> ReadNext(const TString& csv, TString& errString);
     std::shared_ptr<arrow::RecordBatch> ReadSingleBatch(const TString& csv, TString& errString);
+    std::shared_ptr<arrow::RecordBatch> ReadSingleBatch(const TString& csv, const Ydb::Formats::CsvSettings& csvSettings, TString& errString);
 
     void Reset() {
         Reader = {};
@@ -49,14 +50,20 @@ class TArrowCSV {
 
     void SetNullValue(const TString& null = "");
 
-private:
+protected:
     struct TColumnInfo {
         TString Name;
         std::shared_ptr<arrow::DataType> ArrowType;
         std::shared_ptr<arrow::DataType>CsvArrowType;
     };
     using TColummns = TVector<TColumnInfo>;
     TArrowCSV(const TColummns& columns, bool header, const std::set<std::string>& notNullColumns);
+
+    static TString ErrorPrefix() {
+        return "Cannot read CSV: ";
+    }
+
+private:
     arrow::csv::ReadOptions ReadOptions;
     arrow::csv::ParseOptions ParseOptions;
     arrow::csv::ConvertOptions ConvertOptions;
@@ -66,10 +73,6 @@ class TArrowCSV {
     std::set<std::string> NotNullColumns;
 
     std::shared_ptr<arrow::RecordBatch> ConvertColumnTypes(std::shared_ptr<arrow::RecordBatch> parsedBatch) const;
-
-    static TString ErrorPrefix() {
-        return "Cannot read CSV: ";
-    }
 };
 
 }

ydb/core/io_formats/arrow/csv_arrow/ya.make

@@ -0,0 +1,12 @@
+LIBRARY()
+
+SRCS(
+    csv_arrow.cpp
+)
+
+PEERDIR(
+    contrib/libs/apache/arrow
+    ydb/public/api/protos
+)
+
+END()

ydb/core/io_formats/arrow/csv_arrow_ut.cpp

@@ -1,4 +1,4 @@
-#include "csv_arrow.h"
+#include <ydb/core/io_formats/arrow/scheme/scheme.h>
 
 #include <ydb/core/formats/arrow/arrow_helpers.h>
 #include <library/cpp/testing/unittest/registar.h>
@@ -64,7 +64,7 @@ TestReadSingleBatch(TArrowCSV& reader,
 std::shared_ptr<arrow::RecordBatch>
 TestReadSingleBatch(const TVector<std::pair<TString, NScheme::TTypeInfo>>& columns, const TString& data,
                     char delimiter, bool header, ui32 numRows, ui32 skipRows = 0, std::optional<char> escape = {}) {
-    auto reader = TArrowCSV::Create(columns, header);
+    auto reader = TArrowCSVScheme::Create(columns, header);
     UNIT_ASSERT_C(reader.ok(), reader.status().ToString());
     reader->SetDelimiter(delimiter);
     if (skipRows) {
@@ -98,7 +98,7 @@ Y_UNIT_TEST_SUITE(FormatCSV) {
             };
             TInstant dtInstant;
             Y_ABORT_UNLESS(TInstant::TryParseIso8601(dateTimeString, dtInstant));
-            auto reader = TArrowCSV::Create(columns, false);
+            auto reader = TArrowCSVScheme::Create(columns, false);
             UNIT_ASSERT_C(reader.ok(), reader.status().ToString());
 
             TString errorMessage;
@@ -159,7 +159,7 @@ Y_UNIT_TEST_SUITE(FormatCSV) {
         TVector<std::pair<TString, NScheme::TTypeInfo>> columns;
 
         {
-            auto reader = TArrowCSV::Create(columns, false);
+            auto reader = TArrowCSVScheme::Create(columns, false);
             UNIT_ASSERT_C(reader.ok(), reader.status().ToString());
 
             TString errorMessage;
@@ -175,7 +175,7 @@ Y_UNIT_TEST_SUITE(FormatCSV) {
                 {"i64", NScheme::TTypeInfo(NScheme::NTypeIds::Int64)}
             };
 
-            auto reader = TArrowCSV::Create(columns, false);
+            auto reader = TArrowCSVScheme::Create(columns, false);
             UNIT_ASSERT_C(reader.ok(), reader.status().ToString());
 
             TString errorMessage;
@@ -297,7 +297,7 @@ Y_UNIT_TEST_SUITE(FormatCSV) {
             csv += TString() + null + delimiter + q + null + q + delimiter + q + null + q + endLine;
             csv += TString() + null + delimiter + null + delimiter + null + endLine;
 
-            auto reader = TArrowCSV::Create(columns, false);
+            auto reader = TArrowCSVScheme::Create(columns, false);
             UNIT_ASSERT_C(reader.ok(), reader.status().ToString());
             if (!nulls.empty() || !defaultNull) {
                 reader->SetNullValue(null);