only consider grouping keys for id mapping (#1119)

When using a group by expression, only consider the tags
that are part of the grouping rather than the full set
when mapping the id.

spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/DataExpr.java

@@ -471,15 +471,19 @@ final class Count implements AggregateFunction {
   final class GroupBy implements DataExpr {
 
     private final AggregateFunction af;
-    private final List<String> keys;
+    private final Set<String> keys;
 
     /** Create a new instance. */
-    GroupBy(AggregateFunction af, List<String> keys) {
+    GroupBy(AggregateFunction af, Set<String> keys) {
       Preconditions.checkArg(!keys.isEmpty(), "key list for group by cannot be empty");
       this.af = af;
       this.keys = keys;
     }
 
+    Set<String> keys() {
+      return keys;
+    }
+
     @SuppressWarnings("PMD.ReturnEmptyCollectionRatherThanNull")
     private Map<String, String> keyTags(Map<String, String> tags) {
       Map<String, String> result = new HashMap<>();

spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/Evaluator.java

@@ -34,8 +34,8 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
-import java.util.function.Function;
 
 /**
  * Evaluates all the expressions for a set of subscriptions.
@@ -49,7 +49,7 @@ public class Evaluator {
 
   private final Lock lock = new ReentrantLock();
   private final Map<String, String> commonTags;
-  private final Function<Id, Map<String, String>> idMapper;
+  private final BiFunction<Id, Set<String>, Map<String, String>> idMapper;
   private final long step;
   private final boolean delayGaugeAggregation;
   private final QueryIndex<SubscriptionEntry> index;
@@ -174,11 +174,12 @@ public EvalPayload eval(long timestamp) {
           final double v = consolidator.value(timestamp);
           if (!Double.isNaN(v)) {
             Map<String, String> tags = Collections.emptyMap();
-            if (!(expr instanceof DataExpr.AggregateFunction)) {
+            if (expr instanceof DataExpr.GroupBy) {
               // Aggregation functions only use tags based on the expression. Avoid overhead of
               // considering the tags for the data.
-              tags = idMapper.apply(entry.getKey());
-              tags.putAll(commonTags);
+              Set<String> keys = ((DataExpr.GroupBy) expr).keys();
+              tags = idMapper.apply(entry.getKey(), keys);
+              putCommonTags(tags, keys);
             }
             if (delayGaugeAggr && consolidator.isGauge()) {
               // When performing a group by, datapoints missing tag used for the grouping
@@ -211,6 +212,18 @@ public EvalPayload eval(long timestamp) {
     return new EvalPayload(timestamp, metrics);
   }
 
+  private void putCommonTags(Map<String, String> dst, Set<String> keys) {
+    if (dst.size() < keys.size()) {
+      // Skip this step unless there is something pending
+      for (String key : keys) {
+        String value = commonTags.get(key);
+        if (value != null) {
+          dst.put(key, value);
+        }
+      }
+    }
+  }
+
   private String idHash(Id id) {
     Hash64 hasher = new Hash64();
     hasher.updateString(id.name());

spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/EvaluatorConfig.java

@@ -20,6 +20,8 @@
 import com.netflix.spectator.atlas.AtlasConfig;
 
 import java.util.Map;
+import java.util.Set;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 /**
@@ -45,7 +47,7 @@ static EvaluatorConfig fromAtlasConfig(AtlasConfig config) {
           return config.commonTags();
         }
 
-        @Override public Function<Id, Map<String, String>> idMapper() {
+        @Override public BiFunction<Id, Set<String>, Map<String, String>> idMapper() {
           return new IdMapper(JsonUtils.createReplacementFunction(config.validTagCharacters()));
         }
       };
@@ -59,7 +61,7 @@ static EvaluatorConfig fromAtlasConfig(AtlasConfig config) {
   Map<String, String> commonTags();
 
   /** Function to convert an id to a map of key/value pairs. */
-  default Function<Id, Map<String, String>> idMapper() {
+  default BiFunction<Id, Set<String>, Map<String, String>> idMapper() {
     return new IdMapper(Function.identity());
   }
 

spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/IdMapper.java

@@ -19,6 +19,8 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 /**
@@ -27,7 +29,7 @@
  * <p><b>Classes in this package are only intended for use internally within spectator. They may
  * change at any time and without notice.</b>
  */
-public final class IdMapper implements Function<Id, Map<String, String>> {
+public final class IdMapper implements BiFunction<Id, Set<String>, Map<String, String>> {
 
   private final Function<String, String> fixTagString;
 
@@ -37,21 +39,25 @@ public IdMapper(Function<String, String> fixTagString) {
   }
 
   @Override
-  public Map<String, String> apply(Id id) {
+  public Map<String, String> apply(Id id, Set<String> keys) {
     int size = id.size();
-    Map<String, String> tags = new HashMap<>(size);
+    Map<String, String> tags = new HashMap<>(keys.size());
 
     // Start at 1 as name will be added last
     for (int i = 1; i < size; ++i) {
       String k = fixTagString.apply(id.getKey(i));
-      String v = fixTagString.apply(id.getValue(i));
-      tags.put(k, v);
+      if (keys.contains(k)) {
+        String v = fixTagString.apply(id.getValue(i));
+        tags.put(k, v);
+      }
     }
 
     // Add the name, it is added last so it will have precedence if the user tried to
     // use a tag key of "name".
-    String name = fixTagString.apply(id.name());
-    tags.put("name", name);
+    if (keys.contains("name")) {
+      String name = fixTagString.apply(id.name());
+      tags.put("name", name);
+    }
 
     return tags;
   }

spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/Parser.java

@@ -21,6 +21,7 @@
 import java.util.ArrayList;
 import java.util.Deque;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 
 /**
@@ -196,7 +197,7 @@ private static Object parse(String expr) {
         case ":by":
           tmp = (List<String>) stack.pop();
           af = (DataExpr.AggregateFunction) stack.pop();
-          stack.push(new DataExpr.GroupBy(af, tmp));
+          stack.push(new DataExpr.GroupBy(af, new LinkedHashSet<>(tmp)));
           break;
         case ":rollup-drop":
           tmp = (List<String>) stack.pop();