From 3ab54e66b3e06bfc742e45b194224dc087d5fa83 Mon Sep 17 00:00:00 2001
From: Brian Harrington <brharrington@gmail.com>
Date: Fri, 15 Apr 2022 11:32:53 -0500
Subject: [PATCH] allow mixed grouping for custom averages

In the past this wasn't permitted because binary operations
had to have the same grouping on both sides. In #1244 that
was changed so that one side could have a subset of the other.
---
 .../atlas/core/model/CustomVocabulary.scala        | 14 +++++---------
 .../atlas/core/model/CustomVocabularySuite.scala   |  8 ++++++--
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/atlas-core/src/main/scala/com/netflix/atlas/core/model/CustomVocabulary.scala b/atlas-core/src/main/scala/com/netflix/atlas/core/model/CustomVocabulary.scala
index 8ab60070d..34301a819 100644
--- a/atlas-core/src/main/scala/com/netflix/atlas/core/model/CustomVocabulary.scala
+++ b/atlas-core/src/main/scala/com/netflix/atlas/core/model/CustomVocabulary.scala
@@ -165,17 +165,13 @@ object CustomVocabulary {
       val nq = extractCommonQuery(q)
       val numerator = DataExpr.Sum(q)
       val denominator = DataExpr.Sum(baseQuery.and(nq))
-      if (ks.forall(keys.contains)) {
-        // All keys in group by are shared by both sides
-        MathExpr.Divide(DataExpr.GroupBy(numerator, ks), DataExpr.GroupBy(denominator, ks))
-      } else if (ks.forall(k => !keys.contains(k))) {
-        // All keys in group by are only valid for the numerator
+      val denominatorKeys = ks.filter(keys.contains)
+      if (denominatorKeys.isEmpty) {
         MathExpr.Divide(DataExpr.GroupBy(numerator, ks), denominator)
       } else {
-        // Mix of keys that are common to both sides and keys that can only be
-        // applied to the numerator
-        throw new IllegalArgumentException(
-          s"invalid key list for grouping $name, mixes shared and non-shared tag keys"
+        MathExpr.Divide(
+          DataExpr.GroupBy(numerator, ks),
+          DataExpr.GroupBy(denominator, denominatorKeys)
         )
       }
     }
diff --git a/atlas-core/src/test/scala/com/netflix/atlas/core/model/CustomVocabularySuite.scala b/atlas-core/src/test/scala/com/netflix/atlas/core/model/CustomVocabularySuite.scala
index 9f2df7698..6a595f30f 100644
--- a/atlas-core/src/test/scala/com/netflix/atlas/core/model/CustomVocabularySuite.scala
+++ b/atlas-core/src/test/scala/com/netflix/atlas/core/model/CustomVocabularySuite.scala
@@ -139,9 +139,13 @@ class CustomVocabularySuite extends FunSuite {
   }
 
   test("group by mixed keys") {
-    intercept[IllegalArgumentException] {
-      eval("name,(,a,b,c,),:in,app,beacon,:eq,:and,:node-avg,(,name,asg,),:by")
+    val expr = eval("name,(,a,b,c,),:in,app,beacon,:eq,:and,:node-avg,(,name,asg,),:by").rewrite {
+      case MathExpr.NamedRewrite("node-avg", _, e, _, _) => e
     }
+    val expected = eval(
+      s"name,(,a,b,c,),:in,:sum,(,name,asg,),:by,$numInstances,:sum,(,asg,),:by,:div,app,beacon,:eq,:cq"
+    )
+    assertEquals(expr, expected)
   }
 
   test("cg after style") {