[IOTDB-876] Add count storage group DDL (#1734)

antlr/src/main/antlr4/org/apache/iotdb/db/qp/strategy/SqlBase.g4

@@ -80,6 +80,7 @@ statement
     | TRACING OFF #tracingOff
     | COUNT TIMESERIES prefixPath? (GROUP BY LEVEL OPERATOR_EQ INT)? #countTimeseries
     | COUNT DEVICES prefixPath? #countDevices
+    | COUNT STORAGE GROUP prefixPath? #countStorageGroup
     | COUNT NODES prefixPath LEVEL OPERATOR_EQ INT #countNodes
     | LOAD CONFIGURATION (MINUS GLOBAL)? #loadConfigurationStatement
     | {hasSingleQuoteString = true;} LOAD stringLiteral autoCreateSchema?#loadFiles

server/src/main/java/org/apache/iotdb/db/metadata/MManager.java

@@ -852,6 +852,18 @@ public class MManager {
     }
   }
 
+  /**
+   * To calculate the count of storage group for given prefix path.
+   */
+  public int getStorageGroupNum(PartialPath prefixPath) throws MetadataException {
+    lock.readLock().lock();
+    try {
+      return mtree.getStorageGroupNum(prefixPath);
+    } finally {
+      lock.readLock().unlock();
+    }
+  }
+
   /**
    * To calculate the count of nodes in the given level for given prefix path.
    *

server/src/main/java/org/apache/iotdb/db/metadata/MTree.java

@@ -662,6 +662,19 @@ public class MTree implements Serializable {
     return getDevicesCount(root, nodes, 1);
   }
 
+  /**
+   * Get the count of storage group under the given prefix path.
+   *
+   * @param prefixPath a prefix path or a full path, may contain '*'.
+   */
+  int getStorageGroupNum(PartialPath prefixPath) throws MetadataException {
+    String[] nodes = prefixPath.getNodes();
+    if (nodes.length == 0 || !nodes[0].equals(root.getName())) {
+      throw new IllegalPathException(prefixPath.getFullPath());
+    }
+    return getStorageGroupCount(root, nodes, 1, "");
+  }
+
   /**
    * Get the count of nodes in the given level under the given prefix path.
    */
@@ -735,6 +748,31 @@ public class MTree implements Serializable {
     return cnt;
   }
 
+  /**
+   * Traverse the MTree to get the count of storage group.
+   */
+  private int getStorageGroupCount(
+      MNode node, String[] nodes, int idx, String parent) throws MetadataException {
+    int cnt = 0;
+    if (node instanceof StorageGroupMNode && idx >= nodes.length) {
+      cnt++;
+      return cnt;
+    }
+    String nodeReg = MetaUtils.getNodeRegByIdx(idx, nodes);
+    if (!(PATH_WILDCARD).equals(nodeReg)) {
+      if (node.hasChild(nodeReg)) {
+        cnt += getStorageGroupCount(node.getChild(nodeReg),
+            nodes, idx + 1, parent + node.getName() + PATH_SEPARATOR);
+      }
+    } else {
+      for (MNode child : node.getChildren().values()) {
+        cnt += getStorageGroupCount(
+            child, nodes, idx + 1, parent + node.getName() + PATH_SEPARATOR);
+      }
+    }
+    return cnt;
+  }
+
   /**
    * Traverse the MTree to get the count of timeseries in the given level.
    *

server/src/main/java/org/apache/iotdb/db/qp/constant/SQLConstant.java

@@ -159,6 +159,7 @@ public class SQLConstant {
   public static final int TOK_TRACING = 91;
 
   public static final int TOK_COUNT_DEVICES = 92;
+  public static final int TOK_COUNT_STORAGE_GROUP = 93;
 
   public static final Map<Integer, String> tokenSymbol = new HashMap<>();
   public static final Map<Integer, String> tokenNames = new HashMap<>();

server/src/main/java/org/apache/iotdb/db/qp/executor/PlanExecutor.java

@@ -386,6 +386,8 @@ public class PlanExecutor implements IPlanExecutor {
         return processCountNodeTimeSeries((CountPlan) showPlan);
       case COUNT_DEVICES:
         return processCountDevices((CountPlan) showPlan);
+      case COUNT_STORAGE_GROUP:
+        return processCountStorageGroup((CountPlan) showPlan);
       case COUNT_NODES:
         return processCountNodes((CountPlan) showPlan);
       case MERGE_STATUS:
@@ -445,10 +447,28 @@ public class PlanExecutor implements IPlanExecutor {
     return singleDataSet;
   }
 
+  private QueryDataSet processCountStorageGroup(CountPlan countPlan) throws MetadataException {
+    int num = getStorageGroupNum(countPlan.getPath());
+    SingleDataSet singleDataSet =
+        new SingleDataSet(
+            Collections.singletonList(new PartialPath(COLUMN_STORAGE_GROUP, false)),
+            Collections.singletonList(TSDataType.INT32));
+    Field field = new Field(TSDataType.INT32);
+    field.setIntV(num);
+    RowRecord record = new RowRecord(0);
+    record.addField(field);
+    singleDataSet.setRecord(record);
+    return singleDataSet;
+  }
+
   private int getDevicesNum(PartialPath path) throws MetadataException {
     return IoTDB.metaManager.getDevicesNum(path);
   }
 
+  private int getStorageGroupNum(PartialPath path) throws MetadataException {
+    return IoTDB.metaManager.getStorageGroupNum(path);
+  }
+
   protected int getPathsNum(PartialPath path) throws MetadataException {
     return IoTDB.metaManager.getAllTimeseriesCount(path);
   }

server/src/main/java/org/apache/iotdb/db/qp/physical/sys/ShowPlan.java

@@ -50,7 +50,7 @@ public class ShowPlan extends PhysicalPlan {
 
   public enum ShowContentType {
     DYNAMIC_PARAMETER, FLUSH_TASK_INFO, TTL, VERSION, TIMESERIES, STORAGE_GROUP, CHILD_PATH, DEVICES,
-    COUNT_TIMESERIES, COUNT_NODE_TIMESERIES, COUNT_NODES, MERGE_STATUS, COUNT_DEVICES
+    COUNT_TIMESERIES, COUNT_NODE_TIMESERIES, COUNT_NODES, MERGE_STATUS, COUNT_DEVICES, COUNT_STORAGE_GROUP
   }
 
 }

server/src/main/java/org/apache/iotdb/db/qp/strategy/LogicalGenerator.java

@@ -86,6 +86,7 @@ import org.apache.iotdb.db.qp.strategy.SqlBaseParser.AttributeClausesContext;
 import org.apache.iotdb.db.qp.strategy.SqlBaseParser.ConstantContext;
 import org.apache.iotdb.db.qp.strategy.SqlBaseParser.CountDevicesContext;
 import org.apache.iotdb.db.qp.strategy.SqlBaseParser.CountNodesContext;
+import org.apache.iotdb.db.qp.strategy.SqlBaseParser.CountStorageGroupContext;
 import org.apache.iotdb.db.qp.strategy.SqlBaseParser.CountTimeseriesContext;
 import org.apache.iotdb.db.qp.strategy.SqlBaseParser.CreateRoleContext;
 import org.apache.iotdb.db.qp.strategy.SqlBaseParser.CreateSnapshotContext;
@@ -230,6 +231,14 @@ public class LogicalGenerator extends SqlBaseBaseListener {
     initializedOperator = new CountOperator(SQLConstant.TOK_COUNT_DEVICES, path);
   }
 
+  @Override
+  public void enterCountStorageGroup(CountStorageGroupContext ctx) {
+    super.enterCountStorageGroup(ctx);
+    PrefixPathContext pathContext = ctx.prefixPath();
+    PartialPath path = (pathContext != null ? parsePrefixPath(pathContext) : new PartialPath(SQLConstant.getSingleRootArray()));
+    initializedOperator = new CountOperator(SQLConstant.TOK_COUNT_STORAGE_GROUP, path);
+  }
+
   @Override
   public void enterFlush(FlushContext ctx) {
     super.enterFlush(ctx);

server/src/main/java/org/apache/iotdb/db/qp/strategy/PhysicalGenerator.java

@@ -242,6 +242,9 @@ public class PhysicalGenerator {
           case SQLConstant.TOK_COUNT_DEVICES:
             return new CountPlan(
                 ShowContentType.COUNT_DEVICES, ((CountOperator) operator).getPath());
+          case SQLConstant.TOK_COUNT_STORAGE_GROUP:
+            return new CountPlan(
+                ShowContentType.COUNT_STORAGE_GROUP, ((CountOperator) operator).getPath());
           case SQLConstant.TOK_COUNT_NODE_TIMESERIES:
             return new CountPlan(
                 ShowContentType.COUNT_NODE_TIMESERIES,

server/src/main/java/org/apache/iotdb/db/service/StaticResps.java

@@ -107,6 +107,10 @@ class StaticResps {
       Collections.singletonList(COLUMN_COUNT),
       Collections.singletonList(TSDataType.INT32.toString()));
 
+  static final TSExecuteStatementResp COUNT_STORAGE_GROUP = getNoTimeExecuteResp(
+      Collections.singletonList(COLUMN_COUNT),
+      Collections.singletonList(TSDataType.INT32.toString()));
+
   static final TSExecuteStatementResp COUNT_NODES = getNoTimeExecuteResp(
       Collections.singletonList(COLUMN_COUNT),
       Collections.singletonList(TSDataType.INT32.toString()));

server/src/main/java/org/apache/iotdb/db/service/TSServiceImpl.java

@@ -732,6 +732,8 @@ public class TSServiceImpl implements TSIService.Iface, ServerContext {
         return StaticResps.COUNT_TIMESERIES;
       case COUNT_DEVICES:
         return StaticResps.COUNT_DEVICES;
+      case COUNT_STORAGE_GROUP:
+        return StaticResps.COUNT_STORAGE_GROUP;
       case MERGE_STATUS:
         return StaticResps.MERGE_STATUS_RESP;
       default:

server/src/test/java/org/apache/iotdb/db/integration/IoTDBMetadataFetchIT.java

@@ -48,6 +48,7 @@ public class IoTDBMetadataFetchIT {
         Statement statement = connection.createStatement()) {
 
       String[] insertSqls = new String[]{"SET STORAGE GROUP TO root.ln.wf01.wt01",
+          "SET STORAGE GROUP TO root.ln1.wf01.wt01", "SET STORAGE GROUP TO root.ln2.wf01.wt01",
           "CREATE TIMESERIES root.ln.wf01.wt01.status WITH DATATYPE = BOOLEAN, ENCODING = PLAIN",
           "CREATE TIMESERIES root.ln.wf01.wt01.status.s1 WITH DATATYPE = BOOLEAN, ENCODING = PLAIN",
           "CREATE TIMESERIES root.ln.wf01.wt01.temperature WITH DATATYPE = FLOAT, ENCODING = RLE, "
@@ -140,7 +141,8 @@ public class IoTDBMetadataFetchIT {
         Statement statement = connection.createStatement()) {
       String[] sqls = new String[]{"show storage group", "show storage group root.ln.wf01",
           "show storage group root.ln.wf01.wt01.status"};
-      String[] standards = new String[]{"root.ln.wf01.wt01,\n", "root.ln.wf01.wt01,\n", ""};
+      String[] standards = new String[]{"root.ln.wf01.wt01,\n" + "root.ln1.wf01.wt01,\n" + "root.ln2.wf01.wt01,\n", 
+                                        "root.ln.wf01.wt01,\n", ""};
       for (int n = 0; n < sqls.length; n++) {
         String sql = sqls[n];
         String standard = standards[n];
@@ -158,7 +160,7 @@ public class IoTDBMetadataFetchIT {
               }
             }
           }
-          Assert.assertEquals(builder.toString(), standard);
+          Assert.assertEquals(standard, builder.toString());
         } catch (SQLException e) {
           logger.error("showStorageGroupTest() failed", e);
           fail(e.getMessage());
@@ -199,7 +201,7 @@ public class IoTDBMetadataFetchIT {
         if (hasResultSet) {
           try (ResultSet resultSet = statement.getResultSet()) {
             resultSet.next();
-            Assert.assertEquals(resultSet.getString(1), IoTDBConstant.VERSION);
+            Assert.assertEquals(IoTDBConstant.VERSION, resultSet.getString(1));
           }
         }
       } catch (Exception e) {
@@ -233,7 +235,7 @@ public class IoTDBMetadataFetchIT {
               }
             }
           }
-          Assert.assertEquals(builder.toString(), standard);
+          Assert.assertEquals(standard, builder.toString());
         } catch (SQLException e) {
           logger.error("showDevicesTest() failed", e);
           fail(e.getMessage());
@@ -267,7 +269,7 @@ public class IoTDBMetadataFetchIT {
               }
             }
           }
-          Assert.assertEquals(builder.toString(), standard);
+          Assert.assertEquals(standard, builder.toString());
         } catch (SQLException e) {
           logger.error("showChildPaths() failed", e);
           fail(e.getMessage());
@@ -344,6 +346,41 @@ public class IoTDBMetadataFetchIT {
     }
   }
 
+  @Test
+  public void showCountStorageGroup() throws SQLException, ClassNotFoundException {
+    Class.forName(Config.JDBC_DRIVER_NAME);
+    try (Connection connection = DriverManager
+        .getConnection(Config.IOTDB_URL_PREFIX + "127.0.0.1:6667/", "root", "root");
+        Statement statement = connection.createStatement()) {
+      String[] sqls = new String[]{"count storage group root.ln", "count storage group",
+          "count storage group root.ln.wf01.wt01.status"};
+      String[] standards = new String[]{"1,\n", "3,\n", "0,\n"};
+      for (int n = 0; n < sqls.length; n++) {
+        String sql = sqls[n];
+        String standard = standards[n];
+        StringBuilder builder = new StringBuilder();
+        try {
+          boolean hasResultSet = statement.execute(sql);
+          if (hasResultSet) {
+            try (ResultSet resultSet = statement.getResultSet()) {
+              ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+              while (resultSet.next()) {
+                for (int i = 1; i <= resultSetMetaData.getColumnCount(); i++) {
+                  builder.append(resultSet.getString(i)).append(",");
+                }
+                builder.append("\n");
+              }
+            }
+          }
+          Assert.assertEquals(standard, builder.toString());
+        } catch (SQLException e) {
+          logger.error("showCountStorageGroup() failed", e);
+          fail(e.getMessage());
+        }
+      }
+    }
+  }
+
   @Test
   public void showCountTimeSeriesGroupBy() throws SQLException, ClassNotFoundException {
     Class.forName(Config.JDBC_DRIVER_NAME);
@@ -351,7 +388,7 @@ public class IoTDBMetadataFetchIT {
         .getConnection(Config.IOTDB_URL_PREFIX + "127.0.0.1:6667/", "root", "root");
         Statement statement = connection.createStatement()) {
       String[] sqls = new String[]{"COUNT TIMESERIES root group by level=1"};
-      String[] standards = new String[]{"root.ln,3,\n"};
+      String[] standards = new String[]{"root.ln,3,\n" + "root.ln1,0,\n" + "root.ln2,0,\n"};
       for (int n = 0; n < sqls.length; n++) {
         String sql = sqls[n];
         String standard = standards[n];
@@ -369,7 +406,7 @@ public class IoTDBMetadataFetchIT {
               }
             }
           }
-          Assert.assertEquals(builder.toString(), standard);
+          Assert.assertEquals(standard, builder.toString());
         } catch (SQLException e) {
           logger.error("showCountTimeSeriesGroupBy() failed", e);
           fail(e.getMessage());
@@ -385,7 +422,7 @@ public class IoTDBMetadataFetchIT {
         .getConnection(Config.IOTDB_URL_PREFIX + "127.0.0.1:6667/", "root", "root");
         Statement statement = connection.createStatement()) {
       String[] sqls = new String[]{"COUNT NODES root level=1"};
-      String[] standards = new String[]{"1,\n"};
+      String[] standards = new String[]{"3,\n"};
       for (int n = 0; n < sqls.length; n++) {
         String sql = sqls[n];
         String standard = standards[n];
@@ -403,7 +440,7 @@ public class IoTDBMetadataFetchIT {
               }
             }
           }
-          Assert.assertEquals(builder.toString(), standard);
+          Assert.assertEquals(standard, builder.toString());
         } catch (SQLException e) {
           logger.error("showCountNodes() failed", e);
           fail(e.getMessage());
@@ -423,6 +460,11 @@ public class IoTDBMetadataFetchIT {
             + "\n"
             + "{\n"
             + "\t\"root\":{\n"
+            + "\t\t\"ln2\":{\n"
+            + "\t\t\t\"wf01\":{\n"
+            + "\t\t\t\t\"wt01\":{}\n"
+            + "\t\t\t}\n"
+            + "\t\t},\n"
             + "\t\t\"ln\":{\n"
             + "\t\t\t\"wf01\":{\n"
             + "\t\t\t\t\"wt01\":{\n"
@@ -443,6 +485,11 @@ public class IoTDBMetadataFetchIT {
             + "\t\t\t\t\t}\n"
             + "\t\t\t\t}\n"
             + "\t\t\t}\n"
+            + "\t\t},\n"
+            + "\t\t\"ln1\":{\n"
+            + "\t\t\t\"wf01\":{\n"
+            + "\t\t\t\t\"wt01\":{}\n"
+            + "\t\t\t}\n"
             + "\t\t}\n"
             + "\t}\n"
             + "}";