improve legend for custom avg with nested grouping (#1475)

Updates the simple legends generation to handle the case
of custom averages used with nested grouping. Before the
logic to remove the named rewrite would also remove the
first grouping and break the rewrite of any later groupings.

atlas-eval/src/main/scala/com/netflix/atlas/eval/graph/SimpleLegends.scala

@@ -118,8 +118,9 @@ private[graph] object SimpleLegends extends StrictLogging {
     // a simple aggregate like sum based on the display expression.
     expr
       .rewrite {
-        case MathExpr.NamedRewrite(n, q: Query, _, _, _) if n.endsWith("-avg") =>
-          DataExpr.Sum(q)
+        case MathExpr.NamedRewrite(n, q: Query, evalExpr, _, _) if n.endsWith("-avg") =>
+          val aggr = DataExpr.Sum(q)
+          if (evalExpr.isGrouped) DataExpr.GroupBy(aggr, evalExpr.finalGrouping) else aggr
       }
       .asInstanceOf[StyleExpr]
   }

atlas-eval/src/test/scala/com/netflix/atlas/eval/graph/SimpleLegendsSuite.scala

@@ -102,6 +102,15 @@ class SimpleLegendsSuite extends FunSuite {
     assertEquals(legends("name,cpu,:eq,:node-avg"), List("cpu"))
   }
 
+  test("name with node avg and grouping") {
+    assertEquals(legends("name,cpu,:eq,:node-avg,(,app,),:by"), List("$app"))
+  }
+
+  test("name with node avg and nested grouping") {
+    val expr = "name,cpu,:eq,:node-avg,(,app,region,),:by,:max,(,region,),:by"
+    assertEquals(legends(expr), List("$region"))
+  }
+
   test("group by with offsets") {
     val expr = "name,cpu,:eq,:sum,(,id,),:by,(,0h,1w,),:offset"
     assertEquals(legends(expr), List("$id", "$id (offset=$atlas.offset)"))