Specify scheme for CSV files

Author: stebloev

ydb/library/workload/vector/vector_data_generator.cpp

@@ -18,7 +18,7 @@ void TWorkloadVectorFilesDataInitializer::ConfigureOpts(NLastGetopt::TOpts& opts
 
 TBulkDataGeneratorList TWorkloadVectorFilesDataInitializer::DoGetBulkInitialData() {
     return {
-        std::make_shared<TDataGenerator>(*this, Params.TableName, 0, Params.TableName, DataFiles, Default<TVector<TString>>(), TDataGenerator::EPortionSizeUnit::Line)
+        std::make_shared<TDataGenerator>(*this, Params.TableName, 0, Params.TableName, DataFiles, Params.GetColumns(), TDataGenerator::EPortionSizeUnit::Line)
     };
 }
 

ydb/library/workload/vector/vector_workload_params.cpp

@@ -92,6 +92,15 @@ void TVectorWorkloadParams::ConfigureIndexOpts(NLastGetopt::TOpts& opts) {
         .Required().StoreResult(&KmeansTreeClusters);
 }
 
+TVector<TString> TVectorWorkloadParams::GetColumns() const {
+    TVector<TString> result(KeyColumns.begin(), KeyColumns.end());
+    result.emplace_back(EmbeddingColumn);
+    if (PrefixColumn.has_value()) {
+        result.emplace_back(PrefixColumn.value());
+    }
+    return result;
+}
+
 void TVectorWorkloadParams::Init() {
     const TString tablePath = GetFullTableName(TableName.c_str());
 

ydb/library/workload/vector/vector_workload_params.h

@@ -27,6 +27,8 @@ class TVectorWorkloadParams final: public TWorkloadBaseParams {
     void ConfigureCommonOpts(NLastGetopt::TOpts& opts);
     void ConfigureIndexOpts(NLastGetopt::TOpts& opts);
 
+    TVector<TString> GetColumns() const;
+
     TString TableName;
     TString QueryTableName;
     TString IndexName;