fix: csv data source supports data caching (#30)

* feat: csv data source supports data caching

* fix: adjust the precision of the timestamp

* feat: optimize CSV data generation

Author: YamingPei

conf/tdengine-csv.yaml

@@ -22,9 +22,7 @@ schema:
   columns:
     - name: ts
       type: timestamp
-      start: now
       precision : us
-      step: 1
     - name: current
       type: float
     - name: voltage
@@ -40,6 +38,7 @@ schema:
     interlace: 1
     rows_per_table: 100
     rows_per_batch: 10000
+    tables_reuse_data: false
 
 jobs:
   # TDengine insert job

src/actions/components/memory_pool/test/TestMemoryPool.cpp

@@ -233,6 +233,8 @@ void test_memory_pool_cache_mode_basic() {
     std::vector<RowData> cache_data_0, cache_data_1;
     for (size_t i = 0; i < num_rows; ++i) {
         RowData row0, row1;
+        row0.timestamp = 0;
+        row1.timestamp = 0;
         row0.columns = {int32_t(i * 10), std::string("cache0_t0_" + std::to_string(i))};
         row1.columns = {int32_t(i * 100), std::string("cache0_t1_" + std::to_string(i))};
         cache_data_0.push_back(row0);

src/actions/components/reader/csv/inc/ColumnsCSVReader.hpp

@@ -23,7 +23,7 @@ class ColumnsCSVReader {
 
     ~ColumnsCSVReader() = default;
 
-    std::vector<TableData> generate() const;
+    std::unordered_map<std::string, TableData> generate() const;
 
 private:
     ColumnsCSV config_;

src/actions/components/reader/csv/src/ColumnsCSVReader.cpp

@@ -90,7 +90,7 @@ ColumnType ColumnsCSVReader::convert_to_type(const std::string& value, ColumnTyp
     return CSVUtils::convert_to_type(value, target_type);
 }
 
-std::vector<TableData> ColumnsCSVReader::generate() const {
+std::unordered_map<std::string, TableData> ColumnsCSVReader::generate() const {
     try {
         // Create CSV reader
         CSVReader reader(
@@ -242,12 +242,7 @@ std::vector<TableData> ColumnsCSVReader::generate() const {
             data.rows.push_back(std::move(data_row));
         }
 
-        // Convert to std::vector
-        table_data.reserve(table_map.size());
-        for (auto& [_, data] : table_map) {
-            table_data.push_back(std::move(data));
-        }
-        return table_data;
+        return table_map;
 
     } catch (const std::exception& e) {
         std::stringstream ss;

src/actions/components/reader/csv/test/TestColumnsCSVReader.cpp

@@ -71,12 +71,15 @@ void test_generate_table_data_with_default_timestamp() {
     auto table_data = columns_csv.generate();
 
     assert(table_data.size() == 1 && "Expected 1 table");
-    assert(table_data[0].timestamps.size() == 2 && "Expected 2 timestamps");
-    assert(table_data[0].timestamps[0] == 1622505600000 && "Expected first timestamp to match");
-    assert(table_data[0].timestamps[1] == 1622592000000 && "Expected second timestamp to match");
-    assert(table_data[0].rows.size() == 2 && "Expected 2 rows of data");
-    assert(std::get<std::string>(table_data[0].rows[0][0]) == "Alice" && "Expected first column to be 'Alice'");
-    assert(std::get<std::string>(table_data[0].rows[0][1]) == "New York" && "Expected second column to be 'New York'");
+    const auto& table_pair = *table_data.begin();
+    const auto& table = table_pair.second;
+
+    assert(table.timestamps.size() == 2 && "Expected 2 timestamps");
+    assert(table.timestamps[0] == 1622505600000 && "Expected first timestamp to match");
+    assert(table.timestamps[1] == 1622592000000 && "Expected second timestamp to match");
+    assert(table.rows.size() == 2 && "Expected 2 rows of data");
+    assert(std::get<std::string>(table.rows[0][0]) == "Alice" && "Expected first column to be 'Alice'");
+    assert(std::get<std::string>(table.rows[0][1]) == "New York" && "Expected second column to be 'New York'");
     std::cout << "test_generate_table_data_with_default_timestamp passed\n";
 }
 
@@ -106,12 +109,15 @@ void test_generate_table_data_with_timestamp() {
     auto table_data = columns_csv.generate();
 
     assert(table_data.size() == 1 && "Expected 1 table");
-    assert(table_data[0].timestamps.size() == 2 && "Expected 2 timestamps");
-    assert(table_data[0].timestamps[0] == 1622505600000 && "Expected first timestamp to match");
-    assert(table_data[0].timestamps[1] == 1622592000000 && "Expected second timestamp to match");
-    assert(table_data[0].rows.size() == 2 && "Expected 2 rows of data");
-    assert(std::get<std::string>(table_data[0].rows[0][0]) == "Alice" && "Expected first column to be 'Alice'");
-    assert(std::get<std::string>(table_data[0].rows[0][1]) == "New York" && "Expected second column to be 'New York'");
+    const auto& table_pair = *table_data.begin();
+    const auto& table = table_pair.second;
+
+    assert(table.timestamps.size() == 2 && "Expected 2 timestamps");
+    assert(table.timestamps[0] == 1622505600000 && "Expected first timestamp to match");
+    assert(table.timestamps[1] == 1622592000000 && "Expected second timestamp to match");
+    assert(table.rows.size() == 2 && "Expected 2 rows of data");
+    assert(std::get<std::string>(table.rows[0][0]) == "Alice" && "Expected first column to be 'Alice'");
+    assert(std::get<std::string>(table.rows[0][1]) == "New York" && "Expected second column to be 'New York'");
     std::cout << "test_generate_table_data_with_timestamp passed\n";
 }
 
@@ -140,15 +146,17 @@ void test_generate_table_data_with_generated_timestamp() {
     auto table_data = columns_csv.generate();
 
     assert(table_data.size() == 1 && "Expected 1 table");
-    assert(table_data[0].timestamps.size() == 2 && "Expected 2 timestamps");
-    assert(table_data[0].rows.size() == 2 && "Expected 2 rows of data");
-    assert(std::get<std::string>(table_data[0].rows[0][0]) == "Alice" && "Expected first column to be 'Alice'");
-    assert(std::get<int32_t>(table_data[0].rows[0][1]) == 30 && "Expected second column to be 30");
-    assert(std::get<std::string>(table_data[0].rows[0][2]) == "New York" && "Expected third column to be 'New York'");
+    const auto& table_pair = *table_data.begin();
+    const auto& table = table_pair.second;
+
+    assert(table.timestamps.size() == 2 && "Expected 2 timestamps");
+    assert(table.rows.size() == 2 && "Expected 2 rows of data");
+    assert(std::get<std::string>(table.rows[0][0]) == "Alice" && "Expected first column to be 'Alice'");
+    assert(std::get<int32_t>(table.rows[0][1]) == 30 && "Expected second column to be 30");
+    assert(std::get<std::string>(table.rows[0][2]) == "New York" && "Expected third column to be 'New York'");
     std::cout << "test_generate_table_data_with_generated_timestamp passed\n";
 }
 
-
 void test_generate_table_data_include_tbname() {
     ColumnsCSV config;
     config.file_path = "include_tbname.csv";
@@ -175,30 +183,20 @@ void test_generate_table_data_include_tbname() {
     assert(table_data.size() == 2 && "Expected 2 tables");
 
     // Check table names and data
-    bool table1_found = false;
-    bool table2_found = false;
-
-    for (const auto& table : table_data) {
-        if (table.table_name == "table1") {
-            table1_found = true;
-            assert(table.timestamps.size() == 1 && "Expected 1 timestamp for table1");
-            assert(table.rows.size() == 1 && "Expected 1 row of data for table1");
-            assert(std::get<int32_t>(table.rows[0][0]) == 30 && "Expected first column to be 30 for table1");
-            assert(std::get<std::string>(table.rows[0][1]) == "New York" && "Expected second column to be 'New York' for table1");
-        } else if (table.table_name == "table2") {
-            table2_found = true;
-            assert(table.timestamps.size() == 1 && "Expected 1 timestamp for table2");
-            assert(table.rows.size() == 1 && "Expected 1 row of data for table2");
-            assert(std::get<int32_t>(table.rows[0][0]) == 25 && "Expected first column to be 25 for table2");
-            assert(std::get<std::string>(table.rows[0][1]) == "Los Angeles" && "Expected second column to be 'Los Angeles' for table2");
-        }
-    }
+    assert(table_data.find("table1") != table_data.end() && "Expected table1 to be found");
+    assert(table_data.find("table2") != table_data.end() && "Expected table2 to be found");
+
+    const auto& table1 = table_data.at("table1");
+    assert(table1.timestamps.size() == 1 && "Expected 1 timestamp for table1");
+    assert(table1.rows.size() == 1 && "Expected 1 row of data for table1");
+    assert(std::get<int32_t>(table1.rows[0][0]) == 30 && "Expected first column to be 30 for table1");
+    assert(std::get<std::string>(table1.rows[0][1]) == "New York" && "Expected second column to be 'New York' for table1");
 
-    // Verify both tables are found
-    (void)table1_found;
-    (void)table2_found;
-    assert(table1_found && "Expected table1 to be found");
-    assert(table2_found && "Expected table2 to be found");
+    const auto& table2 = table_data.at("table2");
+    assert(table2.timestamps.size() == 1 && "Expected 1 timestamp for table2");
+    assert(table2.rows.size() == 1 && "Expected 1 row of data for table2");
+    assert(std::get<int32_t>(table2.rows[0][0]) == 25 && "Expected first column to be 25 for table2");
+    assert(std::get<std::string>(table2.rows[0][1]) == "Los Angeles" && "Expected second column to be 'Los Angeles' for table2");
 
     std::cout << "test_generate_table_data_include_tbname passed\n";
 }
@@ -223,11 +221,15 @@ void test_generate_table_data_default_column_types() {
     auto table_data = columns_csv.generate();
 
     assert(table_data.size() == 1 && "Expected 1 table");
-    assert(table_data[0].rows.size() == 2 && "Expected 2 rows of data");
-    assert(std::get<std::string>(table_data[0].rows[0][0]) == "Alice" && "Expected first column to be 'Alice'");
-    assert(std::get<std::string>(table_data[0].rows[0][1]) == "30" && "Expected second column to be '30'");
-    assert(std::get<std::string>(table_data[0].rows[0][2]) == "New York" && "Expected third column to be 'New York'");
+    const auto& table_pair = *table_data.begin();
+    const auto& table = table_pair.second;
+
+    assert(table.rows.size() == 2 && "Expected 2 rows of data");
+    assert(std::get<std::string>(table.rows[0][0]) == "Alice" && "Expected first column to be 'Alice'");
+    assert(std::get<std::string>(table.rows[0][1]) == "30" && "Expected second column to be '30'");
+    assert(std::get<std::string>(table.rows[0][2]) == "New York" && "Expected third column to be 'New York'");
     std::cout << "test_generate_table_data_default_column_types passed\n";
+
 }
 
 int main() {

src/actions/config/inc/SchemaConfig.hpp

@@ -61,5 +61,19 @@ struct SchemaConfig {
             columns_cfg.generator.schema = ColumnConfigVector(columns.begin() + 1, columns.end());
             columns_cfg.generator.timestamp_strategy.timestamp_config = columns[0].ts.generator;
         }
+
+        if (generation.data_cache.enabled) {
+            if (from_csv.columns.enabled) {
+                if (from_csv.columns.tbname_index >= 0) {
+                    std::cerr << "[Config Warning] data_cache.enabled is set to false because from_csv.columns.tbname_index in effect." << std::endl;
+                    generation.data_cache.enabled = false;
+                }
+
+                if (from_csv.columns.enabled && !generation.tables_reuse_data) {
+                    std::cerr << "[Config Warning] data_cache.enabled is set to false because tables_reuse_data is false." << std::endl;
+                    generation.data_cache.enabled = false;
+                }
+            }
+        }
     }
 };

src/actions/core/insert/src/InsertDataAction.cpp

@@ -457,7 +457,7 @@ void InsertDataAction::init_cache_units_data(
     size_t max_rows_per_table
 ) {
     for (size_t table_idx = 0; table_idx < max_tables_per_block; ++table_idx) {
-        std::string table_name = "cache_table_" + std::to_string(table_idx);
+        std::string table_name = "default_table";
         RowDataGenerator generator(table_name, config_, col_instances_);
 
         for (size_t cache_idx = 0; cache_idx < num_cached_batches; ++cache_idx) {

src/actions/core/insert/src/generator/src/RowDataGenerator.cpp

@@ -137,37 +137,38 @@ void RowDataGenerator::init_generator() {
 void RowDataGenerator::init_csv_reader() {
     use_generator_ = false;
 
+    if (use_cache_ && timestamp_generator_) {
+        return;
+    }
+
     csv_precision_ = columns_config_.csv.timestamp_strategy.get_precision();
 
-    // Create ColumnsCSV Reader
-    columns_csv_ = std::make_unique<ColumnsCSVReader>(columns_config_.csv, instances_);
+    static std::once_flag csv_once_flag;
+    static std::unordered_map<std::string, TableData> all_tables;
 
-    // TODO: ColumnsCSV Reader needs to support table name index interface
-    // Get all table data
-    std::vector<TableData> all_tables = columns_csv_->generate();
+    std::call_once(csv_once_flag, [this]() {
+        // Create ColumnsCSV Reader
+        auto columns_csv = std::make_unique<ColumnsCSVReader>(columns_config_.csv, instances_);
+        all_tables = columns_csv->generate();
+    });
 
     // Find current table data
-    bool found = false;
-    for (const auto& table_data : all_tables) {
-        if (table_data.table_name == table_name_ || table_data.table_name == "default_table") {
-            found = true;
-
-            for (size_t i = 0; i < table_data.rows.size(); i++) {
-                RowData row;
-                // row.table_name = table_name_;
-                row.timestamp = TimestampUtils::convert_timestamp_precision(table_data.timestamps[i], csv_precision_, target_precision_);
-                row.columns = table_data.rows[i];
-                csv_rows_.push_back(row);
-            }
-            break;
-        }
+    auto it = all_tables.find(table_name_);
+    if (it == all_tables.end()) {
+        it = all_tables.find("default_table");
     }
-
-    if (!found) {
+    if (it == all_tables.end()) {
         throw std::runtime_error("Table '" + table_name_ + "' not found in CSV file");
     }
-}
 
+    const auto& table_data = it->second;
+    for (size_t i = 0; i < table_data.rows.size(); i++) {
+        RowData row;
+        row.timestamp = TimestampUtils::convert_timestamp_precision(table_data.timestamps[i], csv_precision_, target_precision_);
+        row.columns = table_data.rows[i];
+        csv_rows_.push_back(row);
+    }
+}
 
 std::optional<RowData> RowDataGenerator::next_row() {
     if (generated_rows_ >= total_rows_) {
@@ -318,10 +319,19 @@ void RowDataGenerator::generate_from_generator() {
 }
 
 bool RowDataGenerator::generate_from_csv() {
-    cached_row_.timestamp =  csv_rows_[csv_row_index_].timestamp;
+    if (timestamp_generator_) {
+        cached_row_.timestamp = TimestampUtils::convert_timestamp_precision(timestamp_generator_->generate(),
+            timestamp_generator_->timestamp_precision(), target_precision_);
+    } else {
+        cached_row_.timestamp =  csv_rows_[csv_row_index_].timestamp;
+    }
+
     if (!use_cache_) {
         cached_row_.columns = csv_rows_[csv_row_index_].columns;
     }
-    csv_row_index_ = (csv_row_index_ + 1) % csv_rows_.size();
+
+    if (!use_cache_ || !timestamp_generator_) {
+        csv_row_index_ = (csv_row_index_ + 1) % csv_rows_.size();
+    }
     return true;
 }

src/parameter/src/ParameterContext.cpp

@@ -732,6 +732,9 @@ void ParameterContext::parse_commandline(int argc, char* argv[]) {
 
             cli_params[key] = value;
         }
+        else {
+            throw std::runtime_error("Unknown argument: " + arg);
+        }
     }
 }