feat: TsFile 2.x table-model real dataset (read + generate)

configuration/conf/config.properties

@@ -428,6 +428,9 @@
 #   + schema.txt 每一行解释每个Sensor的Type,如"d_0 s_0 3\n d_0 s_1 4"
 # IS_COPY_MODE=false
 
+# Format of the real dataset for generateDataMode export and verification replay: CSV or TSFILE (TsFile 2.x table model)
+REAL_DATASET_FORMAT=CSV
+
 # IS_ADD_ANOMALY=false
 
 # ANOMALY_RATE=0.2

core/pom.xml

@@ -92,6 +92,11 @@
       <artifactId>opencsv</artifactId>
       <version>5.5.2</version>
     </dependency>
+    <dependency>
+      <groupId>org.apache.tsfile</groupId>
+      <artifactId>tsfile</artifactId>
+      <version>2.2.1</version>
+    </dependency>
   </dependencies>
 
 </project>

core/src/main/java/cn/edu/tsinghua/iot/benchmark/client/generate/GenerateDataWriteClient.java

@@ -28,7 +28,7 @@
 
 /** This client is using by GenerateMode */
 public class GenerateDataWriteClient extends GenerateBaseClient {
-  private DataWriter dataWriter = DataWriter.getDataWriter();
+  private DataWriter dataWriter = DataWriter.getDataWriter(clientThreadId);
 
   public GenerateDataWriteClient(
       int id, CountDownLatch countDownLatch, CyclicBarrier barrier, TaskProgress taskProgress) {
@@ -38,14 +38,18 @@ public GenerateDataWriteClient(
   /** Do Operations */
   @Override
   protected void doTest() {
-    taskProgress.resetLoopIndex();
-    for (; taskProgress.getLoopIndex() < config.getLOOP(); taskProgress.incrementLoopIndex()) {
-      if (!doGenerate()) {
-        break;
-      }
-      if (isStop.get()) {
-        break;
+    try {
+      taskProgress.resetLoopIndex();
+      for (; taskProgress.getLoopIndex() < config.getLOOP(); taskProgress.incrementLoopIndex()) {
+        if (!doGenerate()) {
+          break;
+        }
+        if (isStop.get()) {
+          break;
+        }
       }
+    } finally {
+      dataWriter.close();
     }
   }
 

core/src/main/java/cn/edu/tsinghua/iot/benchmark/conf/Config.java

@@ -194,6 +194,12 @@ public class Config {
   /** The path of file */
   private String FILE_PATH = "data/test";
 
+  /**
+   * Format of the real dataset on disk: CSV or TSFILE (table model). Used by generate +
+   * verification.
+   */
+  private RealDatasetFormat REAL_DATASET_FORMAT = RealDatasetFormat.CSV;
+
   /** The size of Big Batch */
   private int BIG_BATCH_SIZE = 100;
 
@@ -1003,6 +1009,14 @@ public void setFILE_PATH(String FILE_PATH) {
     this.FILE_PATH = FILE_PATH;
   }
 
+  public RealDatasetFormat getREAL_DATASET_FORMAT() {
+    return REAL_DATASET_FORMAT;
+  }
+
+  public void setREAL_DATASET_FORMAT(RealDatasetFormat REAL_DATASET_FORMAT) {
+    this.REAL_DATASET_FORMAT = REAL_DATASET_FORMAT;
+  }
+
   public int getDEVICE_NUMBER() {
     return DEVICE_NUMBER;
   }
@@ -2044,7 +2058,10 @@ public String toInfoText() {
         + "\nQUERY_SEED="
         + QUERY_SEED
         + "\nWORKLOAD_BUFFER_SIZE="
-        + WORKLOAD_BUFFER_SIZE;
+        + WORKLOAD_BUFFER_SIZE
+        + "\nREAL_DATASET_FORMAT='"
+        + REAL_DATASET_FORMAT
+        + '\'';
   }
 
   /** get properties from config, one property in one line. */
@@ -2169,6 +2186,7 @@ public ConfigProperties getAllConfigProperties() {
     configProperties.addProperty("Extern Param", "ANOMALY_RATE", this.ANOMALY_RATE);
     configProperties.addProperty("Extern Param", "ANOMALY_TIMES", this.ANOMALY_TIMES);
     configProperties.addProperty("Extern Param", "IS_COPY_MODE", this.IS_COPY_MODE);
+    configProperties.addProperty("Extern Param", "REAL_DATASET_FORMAT", this.REAL_DATASET_FORMAT);
 
     /* The config of schema */
     configProperties.addProperty("Extern Param", "IS_CLIENT_BIND", this.IS_CLIENT_BIND);

core/src/main/java/cn/edu/tsinghua/iot/benchmark/conf/ConfigDescriptor.java

@@ -458,6 +458,10 @@ private void loadProps() {
         config.setIS_COPY_MODE(
             Boolean.parseBoolean(
                 properties.getProperty("IS_COPY_MODE", config.isIS_COPY_MODE() + "")));
+        config.setREAL_DATASET_FORMAT(
+            RealDatasetFormat.getRealDatasetFormat(
+                properties.getProperty(
+                    "REAL_DATASET_FORMAT", config.getREAL_DATASET_FORMAT().name())));
         config.setIS_ADD_ANOMALY(
             Boolean.parseBoolean(
                 properties.getProperty("IS_ADD_ANOMALY", config.isIS_ADD_ANOMALY() + "")));

core/src/main/java/cn/edu/tsinghua/iot/benchmark/conf/RealDatasetFormat.java

@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package cn.edu.tsinghua.iot.benchmark.conf;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public enum RealDatasetFormat {
+  CSV,
+  TSFILE;
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(RealDatasetFormat.class);
+
+  public static RealDatasetFormat getRealDatasetFormat(String value) {
+    if (value != null) {
+      for (RealDatasetFormat format : values()) {
+        if (format.name().equalsIgnoreCase(value.trim())) {
+          return format;
+        }
+      }
+    }
+    LOGGER.warn("Unknown REAL_DATASET_FORMAT '{}', falling back to CSV", value);
+    return CSV;
+  }
+}

core/src/main/java/cn/edu/tsinghua/iot/benchmark/extern/DataWriter.java

@@ -27,7 +27,11 @@ public abstract class DataWriter {
 
   protected static final Config config = ConfigDescriptor.getInstance().getConfig();
 
-  public static DataWriter getDataWriter() {
+  public static DataWriter getDataWriter(int clientId) {
+    if (config.getREAL_DATASET_FORMAT()
+        == cn.edu.tsinghua.iot.benchmark.conf.RealDatasetFormat.TSFILE) {
+      return new TsFileDataWriter(clientId);
+    }
     return new CSVDataWriter();
   }
 
@@ -37,4 +41,7 @@ public static DataWriter getDataWriter() {
    * @param insertLoopIndex loop index of batch
    */
   public abstract boolean writeBatch(IBatch batch, long insertLoopIndex) throws Exception;
+
+  /** Flush and release any open resources. No-op for append-style writers (CSV). */
+  public void close() {}
 }

core/src/main/java/cn/edu/tsinghua/iot/benchmark/extern/SchemaWriter.java

@@ -32,6 +32,10 @@ public abstract class SchemaWriter {
 
   /** Get Basic Writer */
   public static SchemaWriter getBasicWriter() {
+    if (config.getREAL_DATASET_FORMAT()
+        == cn.edu.tsinghua.iot.benchmark.conf.RealDatasetFormat.TSFILE) {
+      return new TsFileSchemaWriter();
+    }
     return new CSVSchemaWriter();
   }
 

core/src/main/java/cn/edu/tsinghua/iot/benchmark/extern/TsFileDataWriter.java

@@ -0,0 +1,187 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package cn.edu.tsinghua.iot.benchmark.extern;
+
+import cn.edu.tsinghua.iot.benchmark.entity.Batch.IBatch;
+import cn.edu.tsinghua.iot.benchmark.entity.Record;
+import cn.edu.tsinghua.iot.benchmark.entity.Sensor;
+import cn.edu.tsinghua.iot.benchmark.schema.schemaImpl.DeviceSchema;
+import cn.edu.tsinghua.iot.benchmark.source.TsFileTableModelMapping;
+import cn.edu.tsinghua.iot.benchmark.utils.FileUtils;
+import org.apache.tsfile.enums.ColumnCategory;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.metadata.TableSchema;
+import org.apache.tsfile.write.record.Tablet;
+import org.apache.tsfile.write.schema.IMeasurementSchema;
+import org.apache.tsfile.write.schema.MeasurementSchema;
+import org.apache.tsfile.write.v4.DeviceTableModelWriter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class TsFileDataWriter extends DataWriter {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(TsFileDataWriter.class);
+  private static final long MEM_THRESHOLD = 32 * 1024 * 1024L;
+
+  private final int clientId;
+  private final Map<String, DeviceTableModelWriter> writers = new HashMap<>();
+
+  public TsFileDataWriter(int clientId) {
+    this.clientId = clientId;
+  }
+
+  @Override
+  public boolean writeBatch(IBatch batch, long insertLoopIndex) throws Exception {
+    while (true) {
+      DeviceSchema deviceSchema = batch.getDeviceSchema();
+      try {
+        DeviceTableModelWriter writer = writerFor(deviceSchema);
+        writer.write(buildTablet(deviceSchema, batch.getRecords()));
+      } catch (Exception e) {
+        LOGGER.error("Failed to write batch for device {}", batch.getDeviceSchema().getDevice(), e);
+        return false;
+      }
+      if (!batch.hasNext()) {
+        break;
+      }
+      batch.next();
+    }
+    batch.finishCheck();
+    return true;
+  }
+
+  private DeviceTableModelWriter writerFor(DeviceSchema deviceSchema) throws IOException {
+    String table = deviceSchema.getTable();
+    DeviceTableModelWriter writer = writers.get(table);
+    if (writer == null) {
+      File dir = Paths.get(config.getFILE_PATH(), table).toFile();
+      Files.createDirectories(dir.toPath());
+      File file = new File(FileUtils.union(dir.getAbsolutePath(), "data_" + clientId + ".tsfile"));
+      Files.deleteIfExists(file.toPath());
+      writer = new DeviceTableModelWriter(file, tableSchemaOf(deviceSchema), MEM_THRESHOLD);
+      writers.put(table, writer);
+    }
+    return writer;
+  }
+
+  /** Column order: FIELD sensors, then device_id TAG, then each tag key as TAG. */
+  private TableSchema tableSchemaOf(DeviceSchema deviceSchema) {
+    List<IMeasurementSchema> cols = new ArrayList<>();
+    List<ColumnCategory> cats = new ArrayList<>();
+    for (Sensor s : deviceSchema.getSensors()) {
+      cols.add(
+          new MeasurementSchema(
+              s.getName(), TsFileTableModelMapping.toTsDataType(s.getSensorType())));
+      cats.add(ColumnCategory.FIELD);
+    }
+    cols.add(new MeasurementSchema(TsFileTableModelMapping.DEVICE_ID_COLUMN, TSDataType.STRING));
+    cats.add(ColumnCategory.TAG);
+    for (String tagKey : deviceSchema.getTags().keySet()) {
+      cols.add(new MeasurementSchema(tagKey, TSDataType.STRING));
+      cats.add(ColumnCategory.TAG);
+    }
+    return new TableSchema(deviceSchema.getTable(), cols, cats);
+  }
+
+  private Tablet buildTablet(DeviceSchema deviceSchema, List<Record> records) {
+    List<Sensor> sensors = deviceSchema.getSensors();
+    List<String> names = new ArrayList<>();
+    List<TSDataType> types = new ArrayList<>();
+    List<ColumnCategory> cats = new ArrayList<>();
+    for (Sensor s : sensors) {
+      names.add(s.getName());
+      types.add(TsFileTableModelMapping.toTsDataType(s.getSensorType()));
+      cats.add(ColumnCategory.FIELD);
+    }
+    names.add(TsFileTableModelMapping.DEVICE_ID_COLUMN);
+    types.add(TSDataType.STRING);
+    cats.add(ColumnCategory.TAG);
+    List<String> tagKeys = new ArrayList<>(deviceSchema.getTags().keySet());
+    for (String tagKey : tagKeys) {
+      names.add(tagKey);
+      types.add(TSDataType.STRING);
+      cats.add(ColumnCategory.TAG);
+    }
+
+    Tablet tablet = new Tablet(deviceSchema.getTable(), names, types, cats, records.size());
+    for (int row = 0; row < records.size(); row++) {
+      Record record = records.get(row);
+      tablet.addTimestamp(row, record.getTimestamp());
+      tablet.addValue(row, TsFileTableModelMapping.DEVICE_ID_COLUMN, deviceSchema.getDevice());
+      for (String tagKey : tagKeys) {
+        tablet.addValue(row, tagKey, deviceSchema.getTags().get(tagKey));
+      }
+      for (int i = 0; i < sensors.size(); i++) {
+        addFieldValue(tablet, row, sensors.get(i), record.getRecordDataValue().get(i));
+      }
+    }
+    tablet.setRowSize(records.size());
+    return tablet;
+  }
+
+  private void addFieldValue(Tablet tablet, int row, Sensor sensor, Object value) {
+    if (value == null) {
+      return; // leave null
+    }
+    String col = sensor.getName();
+    // value is Object; cast to the wrapper (not the primitive) before unboxing.
+    switch (sensor.getSensorType()) {
+      case BOOLEAN:
+        tablet.addValue(row, col, (Boolean) value);
+        break;
+      case INT32:
+        tablet.addValue(row, col, ((Number) value).intValue());
+        break;
+      case INT64:
+      case TIMESTAMP:
+        tablet.addValue(row, col, ((Number) value).longValue());
+        break;
+      case FLOAT:
+        tablet.addValue(row, col, ((Number) value).floatValue());
+        break;
+      case DOUBLE:
+        tablet.addValue(row, col, ((Number) value).doubleValue());
+        break;
+      default:
+        tablet.addValue(row, col, String.valueOf(value));
+    }
+  }
+
+  @Override
+  public void close() {
+    for (DeviceTableModelWriter writer : writers.values()) {
+      try {
+        writer.close();
+      } catch (Exception e) {
+        LOGGER.error("Failed to close TsFile writer", e);
+      }
+    }
+    writers.clear();
+  }
+}