BanyanDB: Support @ShardingKey for Measure tags and set to TopNAggregation group tag by default (#13180)

Author: wankai123

docs/en/changes/changes.md

@@ -12,6 +12,7 @@
 * PromQL Service: traffic query support `limit` and regex match.
 * Fix an edge case of HashCodeSelector(Integer#MIN_VALUE causes ArrayIndexOutOfBoundsException).
 * Support Flink monitoring.
+* BanyanDB: Support `@ShardingKey` for Measure tags and set to TopNAggregation group tag by default.
 
 #### UI
 

oap-server/oal-rt/src/main/java/org/apache/skywalking/oal/rt/util/OALClassGenerator.java

@@ -203,9 +203,9 @@ private Class generateMetricsClass(AnalysisResult metricsStmt) throws OALCompile
                 annotationsAttribute.addAnnotation(columnAnnotation);
                 if (field.isID()) {
                     // Add SeriesID = 0 annotation to ID field.
-                    Annotation banyanShardingKeyAnnotation = new Annotation(BanyanDB.SeriesID.class.getName(), constPool);
-                    banyanShardingKeyAnnotation.addMemberValue("index", new IntegerMemberValue(constPool, 0));
-                    annotationsAttribute.addAnnotation(banyanShardingKeyAnnotation);
+                    Annotation banyanSeriesIDAnnotation = new Annotation(BanyanDB.SeriesID.class.getName(), constPool);
+                    banyanSeriesIDAnnotation.addMemberValue("index", new IntegerMemberValue(constPool, 0));
+                    annotationsAttribute.addAnnotation(banyanSeriesIDAnnotation);
 
                     // Entity id field should enable doc values.
                     final var enableDocValuesAnnotation = new Annotation(ElasticSearch.EnableDocValues.class.getName(), constPool);
@@ -215,6 +215,10 @@ private Class generateMetricsClass(AnalysisResult metricsStmt) throws OALCompile
                 if (field.isGroupByCondInTopN()) {
                     Annotation banyanTopNAggregationAnnotation = new Annotation(BanyanDB.TopNAggregation.class.getName(), constPool);
                     annotationsAttribute.addAnnotation(banyanTopNAggregationAnnotation);
+                    // If TopN, add ShardingKey to group field.
+                    Annotation banyanShardingKeyAnnotation = new Annotation(BanyanDB.ShardingKey.class.getName(), constPool);
+                    banyanShardingKeyAnnotation.addMemberValue("index", new IntegerMemberValue(constPool, 0));
+                    annotationsAttribute.addAnnotation(banyanShardingKeyAnnotation);
                 }
 
                 newField.getFieldInfo().addAttribute(annotationsAttribute);

oap-server/server-core/src/main/java/org/apache/skywalking/oap/server/core/storage/annotation/BanyanDB.java

@@ -77,6 +77,30 @@
         int index() default -1;
     }
 
+    /**
+     * ShardingKey is used to group time series data per metric in one place. Optional. Only support Measure Tag.
+     * If ShardingKey is not set, the default ShardingKey is based on the combination of 'name' and 'entity' according to the {@link SeriesID}.
+     * <p>
+     * The typical scenario to specify the ShardingKey to the Group tag when the metric generate a TopNAggregation:
+     * If not set, the default data distribution based on the combination of 'name' and 'entity', can lead to performance issues when calculating the 'TopNAggregation'.
+     * This is because each shard only has a subset of the top-n list, and the query process has to be responsible for aggregating those lists to obtain the final result.
+     * This introduces overhead in terms of querying performance and disk usage.
+     *
+     * @since 10.3.0
+     */
+    @Target({ElementType.FIELD})
+    @Retention(RetentionPolicy.RUNTIME)
+    @interface ShardingKey {
+        /**
+         * Relative sharding tag
+         * <p>
+         * The index number determines the order of the column placed in the ShardingKey.
+         *
+         * @return non-negative if this column be used for sharding. -1 means not as a sharding key
+         */
+        int index() default -1;
+    }
+
     /**
      * Force disabling indexing declare through {@link Column}.
      * In BanyanDB, some additional conditions could be done in server memory, no indexing required in this case.

oap-server/server-core/src/main/java/org/apache/skywalking/oap/server/core/storage/model/BanyanDBExtension.java

@@ -39,6 +39,9 @@ public class BanyanDBExtension {
     @Getter
     private final int seriesIDIdx;
 
+    @Getter
+    private final int shardingKeyIdx;
+
     /**
      * {@link BanyanDB.NoIndexing} exists to override {@link ModelColumn#shouldIndex()}, or be the same as {@link
      * ModelColumn#shouldIndex()}.
@@ -82,6 +85,13 @@ public boolean isSeriesID() {
         return this.seriesIDIdx > -1;
     }
 
+    /**
+     * @return true if this column is a part of sharding key
+     */
+    public boolean isShardingKey() {
+        return this.shardingKeyIdx > -1;
+    }
+
     /**
      * @return true if this column should be indexing in the BanyanDB
      */

oap-server/server-core/src/main/java/org/apache/skywalking/oap/server/core/storage/model/StorageModels.java

@@ -58,11 +58,15 @@ public Model add(Class<?> aClass, int scopeId, Storage storage) throws StorageEx
         DefaultScopeDefine.nameOf(scopeId);
 
         List<ModelColumn> modelColumns = new ArrayList<>();
-        SeriesIDChecker checker = new SeriesIDChecker();
+        SeriesIDChecker seriesIDChecker = new SeriesIDChecker();
+        ShardingKeyChecker shardingKeyChecker = new ShardingKeyChecker();
         SQLDatabaseModelExtension sqlDBModelExtension = new SQLDatabaseModelExtension();
         BanyanDBModelExtension banyanDBModelExtension = new BanyanDBModelExtension();
         ElasticSearchModelExtension elasticSearchModelExtension = new ElasticSearchModelExtension();
-        retrieval(aClass, storage.getModelName(), modelColumns, scopeId, checker, sqlDBModelExtension, banyanDBModelExtension);
+        retrieval(
+            aClass, storage.getModelName(), modelColumns, scopeId, seriesIDChecker, shardingKeyChecker, sqlDBModelExtension,
+            banyanDBModelExtension
+        );
         // Add extra column for additional entities
         if (aClass.isAnnotationPresent(SQLDatabase.ExtraColumn4AdditionalEntity.class)
             || aClass.isAnnotationPresent(SQLDatabase.MultipleExtraColumn4AdditionalEntity.class)) {
@@ -106,7 +110,8 @@ public Model add(Class<?> aClass, int scopeId, Storage storage) throws StorageEx
         // Set routing rules for ElasticSearch
         elasticSearchModelExtension.setRouting(storage.getModelName(), modelColumns);
 
-        checker.check(storage.getModelName());
+        seriesIDChecker.check(storage.getModelName());
+        shardingKeyChecker.check(storage.getModelName());
 
         Model model = new Model(
             storage.getModelName(),
@@ -153,7 +158,8 @@ private void retrieval(final Class<?> clazz,
                            final String modelName,
                            final List<ModelColumn> modelColumns,
                            final int scopeId,
-                           SeriesIDChecker checker,
+                           SeriesIDChecker seriesIDChecker,
+                           ShardingKeyChecker shardingKeyChecker,
                            final SQLDatabaseModelExtension sqlDBModelExtension,
                            final BanyanDBModelExtension banyanDBModelExtension) {
         if (log.isDebugEnabled()) {
@@ -212,6 +218,8 @@ private void retrieval(final Class<?> clazz,
                 // BanyanDB extension
                 final BanyanDB.SeriesID banyanDBSeriesID = field.getAnnotation(
                     BanyanDB.SeriesID.class);
+                final BanyanDB.ShardingKey banyanDBShardingKey = field.getAnnotation(
+                    BanyanDB.ShardingKey.class);
                 final BanyanDB.NoIndexing banyanDBNoIndex = field.getAnnotation(
                     BanyanDB.NoIndexing.class);
                 final BanyanDB.IndexRule banyanDBIndexRule = field.getAnnotation(
@@ -227,6 +235,7 @@ private void retrieval(final Class<?> clazz,
                 final boolean shouldIndex = (banyanDBNoIndex == null) && !column.storageOnly();
                 BanyanDBExtension banyanDBExtension = new BanyanDBExtension(
                     banyanDBSeriesID == null ? -1 : banyanDBSeriesID.index(),
+                    banyanDBShardingKey == null ? -1 : banyanDBShardingKey.index(),
                     shouldIndex,
                     banyanDBIndexRule == null ? BanyanDB.IndexRule.IndexType.INVERTED : banyanDBIndexRule.indexType(),
                     banyanDBMeasureField != null,
@@ -255,7 +264,10 @@ private void retrieval(final Class<?> clazz,
                     banyanDBExtension
                 );
                 if (banyanDBExtension.isSeriesID()) {
-                    checker.accept(modelName, modelColumn);
+                    seriesIDChecker.accept(modelName, modelColumn);
+                }
+                if (banyanDBExtension.isShardingKey()) {
+                    shardingKeyChecker.accept(modelName, modelColumn);
                 }
 
                 if (field.isAnnotationPresent(SQLDatabase.AdditionalEntity.class)) {
@@ -282,7 +294,10 @@ private void retrieval(final Class<?> clazz,
         }
 
         if (Objects.nonNull(clazz.getSuperclass())) {
-            retrieval(clazz.getSuperclass(), modelName, modelColumns, scopeId, checker, sqlDBModelExtension, banyanDBModelExtension);
+            retrieval(
+                clazz.getSuperclass(), modelName, modelColumns, scopeId, seriesIDChecker, shardingKeyChecker,
+                sqlDBModelExtension, banyanDBModelExtension
+            );
         }
     }
 
@@ -359,4 +374,40 @@ private void check(String modelName) throws IllegalStateException {
             }
         }
     }
+
+    private static class ShardingKeyChecker {
+        private final ArrayList<ModelColumn> keys = new ArrayList<>();
+
+        /**
+         * @throws IllegalStateException if sharding key indices are conflicting.
+         */
+        private void accept(String modelName, ModelColumn modelColumn) throws IllegalStateException {
+            final int idx = modelColumn.getBanyanDBExtension().getShardingKeyIdx();
+            while (idx + 1 > keys.size()) {
+                keys.add(null);
+            }
+            ModelColumn exist = keys.get(idx);
+            if (exist != null) {
+                throw new IllegalStateException(
+                    modelName + "'s "
+                        + "Column [" + exist.getColumnName() + "] and column [" + modelColumn.getColumnName()
+                        + " are conflicting with sharding key index=" + modelColumn.getBanyanDBExtension()
+                                                                                   .getShardingKeyIdx());
+            }
+            keys.set(idx, modelColumn);
+        }
+
+        /**
+         * @param modelName model name of the entity
+         * @throws IllegalStateException if sharding key indices are not continuous
+         */
+        private void check(String modelName) throws IllegalStateException {
+            for (int i = 0; i < keys.size(); i++) {
+                final ModelColumn modelColumn = keys.get(i);
+                if (modelColumn == null) {
+                    throw new IllegalStateException("sharding key index=" + i + " is missing in " + modelName);
+                }
+            }
+        }
+    }
 }

oap-server/server-core/src/test/java/org/apache/skywalking/oap/server/core/storage/model/ModelColumnTest.java

@@ -50,7 +50,7 @@ public void testColumnDefine() {
                                              new ElasticSearchExtension(
                                                  ElasticSearch.MatchQuery.AnalyzerType.OAP_ANALYZER, null, false, false, true),
                                              new BanyanDBExtension(
-                                                 -1, true, BanyanDB.IndexRule.IndexType.INVERTED, false,
+                                                 -1, -1, true, BanyanDB.IndexRule.IndexType.INVERTED, false,
                                                  BanyanDB.MatchQuery.AnalyzerType.SIMPLE, true
                                              )
         );
@@ -62,7 +62,7 @@ public void testColumnDefine() {
                                  new SQLDatabaseExtension(),
                                  new ElasticSearchExtension(ElasticSearch.MatchQuery.AnalyzerType.OAP_ANALYZER, null, false, false, true),
                                  new BanyanDBExtension(
-                                     -1, true, BanyanDB.IndexRule.IndexType.INVERTED, false,
+                                     -1, -1, true, BanyanDB.IndexRule.IndexType.INVERTED, false,
                                      BanyanDB.MatchQuery.AnalyzerType.SIMPLE, true
                                  )
         );
@@ -74,7 +74,7 @@ public void testColumnDefine() {
                                  false, false, true, 200,
                                  new SQLDatabaseExtension(),
                                  new ElasticSearchExtension(ElasticSearch.MatchQuery.AnalyzerType.OAP_ANALYZER, null, false, false, true),
-                                 new BanyanDBExtension(-1, true, BanyanDB.IndexRule.IndexType.INVERTED, false, BanyanDB.MatchQuery.AnalyzerType.SIMPLE, true)
+                                 new BanyanDBExtension(-1, -1, true, BanyanDB.IndexRule.IndexType.INVERTED, false, BanyanDB.MatchQuery.AnalyzerType.SIMPLE, true)
         );
         Assertions.assertFalse(column.isStorageOnly());
         Assertions.assertEquals("abc", column.getColumnName().getName());
@@ -89,7 +89,7 @@ public void testConflictDefinition() {
                     new ElasticSearchExtension(
                             ElasticSearch.MatchQuery.AnalyzerType.OAP_ANALYZER, "abc", false, false, true),
                             new BanyanDBExtension(
-                                -1, true, BanyanDB.IndexRule.IndexType.INVERTED, false,
+                                -1, -1, true, BanyanDB.IndexRule.IndexType.INVERTED, false,
                                 BanyanDB.MatchQuery.AnalyzerType.SIMPLE, true
                             )
             );
@@ -104,7 +104,7 @@ public void testConflictDefinitionIndexOnly() {
                     new SQLDatabaseExtension(),
                     new ElasticSearchExtension(
                             ElasticSearch.MatchQuery.AnalyzerType.OAP_ANALYZER, "abc", false, false, true),
-                            new BanyanDBExtension(-1, true, BanyanDB.IndexRule.IndexType.INVERTED, false,
+                            new BanyanDBExtension(-1, -1, true, BanyanDB.IndexRule.IndexType.INVERTED, false,
                                                   BanyanDB.MatchQuery.AnalyzerType.SIMPLE, true
                             )
             );
@@ -120,7 +120,7 @@ public void testConflictDefinitionStorageOnly() {
                             new ElasticSearchExtension(
                                 ElasticSearch.MatchQuery.AnalyzerType.OAP_ANALYZER, "abc", false, false, true),
                             new BanyanDBExtension(
-                                -1, false, BanyanDB.IndexRule.IndexType.INVERTED, false,
+                                -1, -1, false, BanyanDB.IndexRule.IndexType.INVERTED, false,
                                 BanyanDB.MatchQuery.AnalyzerType.SIMPLE, true
                             )
             );
@@ -136,7 +136,7 @@ public void testConflictDefinitionEnableSort() {
                             new ElasticSearchExtension(
                                 ElasticSearch.MatchQuery.AnalyzerType.OAP_ANALYZER, "abc", false, false, true),
                             new BanyanDBExtension(
-                                -1, false, BanyanDB.IndexRule.IndexType.INVERTED, true,
+                                -1, -1, false, BanyanDB.IndexRule.IndexType.INVERTED, true,
                                 null, true
                             )
             );

oap-server/server-storage-plugin/storage-banyandb-plugin/src/main/java/org/apache/skywalking/oap/server/storage/plugin/banyandb/MetadataRegistry.java

@@ -131,6 +131,7 @@ public MeasureModel registerMeasureModel(Model model, BanyanDBStorageConfig conf
                 .collect(Collectors.toMap(modelColumn -> modelColumn.getColumnName().getStorageName(), Function.identity()));
         // parse and set seriesIDs
         List<String> seriesIDColumns = parseEntityNames(modelColumnMap);
+        List<String> shardingKeyColumns = parseShardingKeyNames(modelColumnMap);
         if (seriesIDColumns.isEmpty()) {
             throw new StorageException("model " + model.getName() + " doesn't contain series id");
         }
@@ -160,6 +161,9 @@ public MeasureModel registerMeasureModel(Model model, BanyanDBStorageConfig conf
                 .setName(schemaMetadata.name()));
         builder.setInterval(downSamplingDuration(model.getDownsampling()).format());
         builder.setEntity(BanyandbDatabase.Entity.newBuilder().addAllTagNames(seriesIDColumns));
+        if (CollectionUtils.isNotEmpty(shardingKeyColumns)) {
+            builder.setShardingKey(BanyandbDatabase.ShardingKey.newBuilder().addAllTagNames(shardingKeyColumns));
+        }
         builder.addAllTagFamilies(tagFamilySpecs);
         if (model.getBanyanDBModelExtension().isIndexMode()) {
             builder.setIndexMode(true);
@@ -351,6 +355,19 @@ List<String> parseEntityNames(Map<String, ModelColumn> modelColumnMap) {
                 .collect(Collectors.toList());
     }
 
+    List<String> parseShardingKeyNames(Map<String, ModelColumn> modelColumnMap) {
+        List<ModelColumn> shardingKeyColumns = new ArrayList<>();
+        for (final ModelColumn col : modelColumnMap.values()) {
+            if (col.getBanyanDBExtension().isShardingKey()) {
+                shardingKeyColumns.add(col);
+            }
+        }
+        return shardingKeyColumns.stream()
+                .sorted(Comparator.comparingInt(col -> col.getBanyanDBExtension().getShardingKeyIdx()))
+                .map(col -> col.getColumnName().getName())
+                .collect(Collectors.toList());
+    }
+
     /**
      * Parse tags' metadata for {@link Stream}
      * Every field of a class is registered as a {@link org.apache.skywalking.banyandb.model.v1.BanyandbModel.Tag}