Add basic import from files

Author: stebloev

ydb/library/workload/vector/vector_data_generator.cpp

@@ -0,0 +1,25 @@
+#include "vector_data_generator.h"
+
+namespace NYdbWorkload {
+
+TWorkloadVectorFilesDataInitializer::TWorkloadVectorFilesDataInitializer(const TVectorWorkloadParams& params)
+    : TWorkloadDataInitializerBase("files", "Import vectors from files", params)
+    , Params(params)
+{ }
+
+void TWorkloadVectorFilesDataInitializer::ConfigureOpts(NLastGetopt::TOpts& opts) {
+    opts.AddLongOption('i', "input",
+            "File or Directory with dataset. If directory is set, all its available files will be used. "
+            "Supports zipped and unzipped csv, tsv files and parquet ones that may be downloaded here: "
+            "https://huggingface.co/datasets/Cohere/wikipedia-22-12-simple-embeddings. "
+            "For better perfomanse you may split it to some parts for parrallel upload."
+        ).Required().StoreResult(&DataFiles);
+}
+
+TBulkDataGeneratorList TWorkloadVectorFilesDataInitializer::DoGetBulkInitialData() {
+    return {
+        std::make_shared<TDataGenerator>(*this, Params.TableName, 0, Params.TableName, DataFiles, Default<TVector<TString>>(), TDataGenerator::EPortionSizeUnit::Line)
+    };
+}
+
+} // namespace NYdbWorkload

ydb/library/workload/vector/vector_data_generator.h

@@ -0,0 +1,22 @@
+#pragma once
+
+#include "vector_workload_params.h"
+
+#include <ydb/library/workload/benchmark_base/workload.h>
+#include <ydb/library/workload/benchmark_base/data_generator.h>
+
+namespace NYdbWorkload {
+
+class TWorkloadVectorFilesDataInitializer : public TWorkloadDataInitializerBase {
+private:
+    const TVectorWorkloadParams& Params;
+    TString DataFiles;
+
+public:
+    TWorkloadVectorFilesDataInitializer(const TVectorWorkloadParams& params);
+
+    virtual void ConfigureOpts(NLastGetopt::TOpts& opts) override;
+    virtual TBulkDataGeneratorList DoGetBulkInitialData() override;
+};
+
+} // namespace NYdbWorkload

ydb/library/workload/vector/vector_workload_params.cpp

@@ -1,3 +1,4 @@
+#include "vector_data_generator.h"
 #include "vector_enums.h"
 #include "vector_workload_params.h"
 #include "vector_workload_generator.h"
@@ -193,6 +194,12 @@ THolder<IWorkloadQueryGenerator> TVectorWorkloadParams::CreateGenerator() const
     return MakeHolder<TVectorWorkloadGenerator>(this);
 }
 
+TWorkloadDataInitializer::TList TVectorWorkloadParams::CreateDataInitializers() const {
+    return {
+        std::make_shared<TWorkloadVectorFilesDataInitializer>(*this)
+    };
+}
+
 TString TVectorWorkloadParams::GetWorkloadName() const {
     return "vector";
 }

ydb/library/workload/vector/vector_workload_params.h

@@ -18,6 +18,7 @@ class TVectorWorkloadParams final: public TWorkloadBaseParams {
 public:
     void ConfigureOpts(NLastGetopt::TOpts& opts, const ECommandType commandType, int workloadType) override;
     THolder<IWorkloadQueryGenerator> CreateGenerator() const override;
+    TWorkloadDataInitializer::TList CreateDataInitializers() const override;
     TString GetWorkloadName() const override;
     void Validate(const ECommandType commandType, int workloadType) override;
 

ydb/library/workload/vector/ya.make

@@ -2,6 +2,7 @@ LIBRARY()
 
 SRCS(
     vector_command_index.cpp
+    vector_data_generator.cpp
     vector_recall_evaluator.cpp
     vector_sampler.cpp
     vector_sql.cpp