Subscription: add more methods for table session dataset (#15398)

Author: VGalaxies

example/session/src/main/java/org/apache/iotdb/TableModelSubscriptionSessionExample.java

@@ -106,8 +106,8 @@ private static void dataSubscription() throws Exception {
             .password("root")
             .build()) {
       final Properties config = new Properties();
-      config.put(TopicConstant.DATABASE_KEY, "db");
-      config.put(TopicConstant.TABLE_KEY, "test");
+      config.put(TopicConstant.DATABASE_KEY, "db.*");
+      config.put(TopicConstant.TABLE_KEY, "test.*");
       config.put(TopicConstant.START_TIME_KEY, 25);
       config.put(TopicConstant.END_TIME_KEY, 75);
       config.put(TopicConstant.STRICT_KEY, "true");
@@ -133,8 +133,10 @@ private static void dataSubscription() throws Exception {
         for (final SubscriptionMessage message : messages) {
           for (final SubscriptionSessionDataSet dataSet : message.getSessionDataSetsHandler()) {
             System.out.println(dataSet.getDatabaseName());
+            System.out.println(dataSet.getTableName());
             System.out.println(dataSet.getColumnNames());
             System.out.println(dataSet.getColumnTypes());
+            System.out.println(dataSet.getColumnCategories());
             while (dataSet.hasNext()) {
               System.out.println(dataSet.next());
             }
@@ -166,8 +168,14 @@ public static void main(final String[] args) throws Exception {
             .username("root")
             .password("root")
             .build()) {
-      createDataBaseAndTable(session, "db", "test");
-      insertData(session, "db", "test", 0, 100);
+      createDataBaseAndTable(session, "db1", "test1");
+      createDataBaseAndTable(session, "db1", "test2");
+      createDataBaseAndTable(session, "db2", "test1");
+      createDataBaseAndTable(session, "db2", "test2");
+      insertData(session, "db1", "test1", 0, 100);
+      insertData(session, "db1", "test2", 0, 100);
+      insertData(session, "db2", "test1", 0, 100);
+      insertData(session, "db2", "test2", 0, 100);
       dataSubscription();
     }
   }

iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/subscription/annotation/TableModel.java

@@ -28,6 +28,6 @@
  * Indicates that the method is valid only within the subscription module under the table model
  * namespace. Otherwise, the behavior is undefined.
  */
-@Target(ElementType.METHOD)
+@Target({ElementType.METHOD, ElementType.TYPE, ElementType.FIELD})
 @Retention(RetentionPolicy.RUNTIME)
 public @interface TableModel {}

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/payload/SubscriptionSessionDataSet.java

@@ -31,16 +31,19 @@
 import org.apache.tsfile.utils.DateUtils;
 import org.apache.tsfile.write.UnSupportedDataTypeException;
 import org.apache.tsfile.write.record.Tablet;
+import org.apache.tsfile.write.record.Tablet.ColumnCategory;
 import org.apache.tsfile.write.schema.IMeasurementSchema;
 
 import java.time.LocalDate;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public class SubscriptionSessionDataSet implements ISessionDataSet {
 
@@ -52,15 +55,67 @@ public Tablet getTablet() {
     return tablet;
   }
 
+  public SubscriptionSessionDataSet(final Tablet tablet, @Nullable final String databaseName) {
+    this.tablet = tablet;
+    this.databaseName = databaseName;
+    generateRowIterator();
+  }
+
+  /////////////////////////////// table model ///////////////////////////////
+
+  @TableModel private List<ColumnCategory> columnCategoryList;
+
   @TableModel
   public String getDatabaseName() {
     return databaseName;
   }
 
-  public SubscriptionSessionDataSet(final Tablet tablet, @Nullable final String databaseName) {
-    this.tablet = tablet;
-    this.databaseName = databaseName;
-    generateRowIterator();
+  @TableModel
+  public String getTableName() {
+    return tablet.getTableName();
+  }
+
+  @TableModel
+  public List<ColumnCategory> getColumnCategories() {
+    if (Objects.nonNull(columnCategoryList)) {
+      return columnCategoryList;
+    }
+
+    if (!isTableData()) {
+      return Collections.emptyList();
+    }
+
+    return columnCategoryList =
+        Stream.concat(
+                Stream.of(ColumnCategory.TIME),
+                tablet.getColumnTypes().stream()
+                    .map(
+                        columnCategory -> {
+                          switch (columnCategory) {
+                            case FIELD:
+                              return ColumnCategory.FIELD;
+                            case TAG:
+                              return ColumnCategory.TAG;
+                            case ATTRIBUTE:
+                              return ColumnCategory.ATTRIBUTE;
+                            default:
+                              throw new IllegalArgumentException(
+                                  "Unknown column category: " + columnCategory);
+                          }
+                        }))
+            .collect(Collectors.toList());
+  }
+
+  @TableModel
+  public enum ColumnCategory {
+    TIME,
+    TAG,
+    FIELD,
+    ATTRIBUTE
+  }
+
+  private boolean isTableData() {
+    return Objects.nonNull(databaseName);
   }
 
   /////////////////////////////// override ///////////////////////////////
@@ -74,16 +129,17 @@ public List<String> getColumnNames() {
       return columnNameList;
     }
 
-    columnNameList = new ArrayList<>();
-    columnNameList.add("Time");
-
-    String deviceId = tablet.getDeviceId();
     List<IMeasurementSchema> schemas = tablet.getSchemas();
-    columnNameList.addAll(
-        schemas.stream()
-            .map((schema) -> deviceId + "." + schema.getMeasurementName())
-            .collect(Collectors.toList()));
-    return columnNameList;
+    String deviceId = tablet.getDeviceId();
+    return columnNameList =
+        isTableData()
+            ? Stream.concat(
+                    Stream.of("time"), schemas.stream().map(IMeasurementSchema::getMeasurementName))
+                .collect(Collectors.toList())
+            : Stream.concat(
+                    Stream.of("Time"),
+                    schemas.stream().map(schema -> deviceId + "." + schema.getMeasurementName()))
+                .collect(Collectors.toList());
   }
 
   @Override
@@ -92,13 +148,12 @@ public List<String> getColumnTypes() {
       return columnTypeList;
     }
 
-    columnTypeList = new ArrayList<>();
-    columnTypeList.add(TSDataType.INT64.toString());
-
     List<IMeasurementSchema> schemas = tablet.getSchemas();
-    columnTypeList.addAll(
-        schemas.stream().map(schema -> schema.getType().toString()).collect(Collectors.toList()));
-    return columnTypeList;
+    return columnTypeList =
+        Stream.concat(
+                Stream.of(TSDataType.INT64.toString()),
+                schemas.stream().map(schema -> schema.getType().toString()))
+            .collect(Collectors.toList());
   }
 
   public boolean hasNext() {