fix :cg with non-grouped math aggregate (#1445)

If there was an aggregation of a group by result that
was not grouped, then it was not getting the common
grouping applied.

atlas-core/src/main/scala/com/netflix/atlas/core/model/MathExpr.scala

@@ -661,7 +661,7 @@ object MathExpr {
  }
  }
 
- trait AggrMathExpr extends TimeSeriesExpr {
+ sealed trait AggrMathExpr extends TimeSeriesExpr {
 
  def name: String
 

atlas-core/src/main/scala/com/netflix/atlas/core/model/MathVocabulary.scala

@@ -481,15 +481,27 @@ object MathVocabulary extends Vocabulary {
  nr.groupBy(keys)
  case af: AggregateFunction =>
  DataExpr.GroupBy(af, keys)
+ case af: MathExpr.AggrMathExpr =>
+ val e = addCommonKeys(af.expr, keys)
+ MathExpr.GroupBy(copy(af, e), keys)
  case e: DataExpr.GroupBy =>
  e.copy(keys = mergeKeys(e.keys, keys))
  case e: MathExpr.GroupBy =>
- val af = addCommonKeys(e.expr, keys).asInstanceOf[AggrMathExpr]
+ val af = copy(e.expr, addCommonKeys(e.expr.expr, keys))
  MathExpr.GroupBy(af, mergeKeys(e.keys, keys))
  }
  newExpr.asInstanceOf[TimeSeriesExpr]
  }
 
+ private def copy(aggr: AggrMathExpr, expr: TimeSeriesExpr): AggrMathExpr = {
+ aggr match {
+ case af @ MathExpr.Count(_) => af.copy(expr = expr)
+ case af @ MathExpr.Max(_) => af.copy(expr = expr)
+ case af @ MathExpr.Min(_) => af.copy(expr = expr)
+ case af @ MathExpr.Sum(_) => af.copy(expr = expr)
+ }
+ }
+
  override protected def executor: PartialFunction[List[Any], List[Any]] = {
  case StringListType(keys) :: TimeSeriesType(t) :: stack =>
  addCommonKeys(t, keys) :: stack

atlas-core/src/test/scala/com/netflix/atlas/core/model/MathGroupBySuite.scala

@@ -307,4 +307,13 @@ class MathGroupBySuite extends FunSuite {
  val exprExplicit = eval(inputExplicit)
  assertEquals(expr.toString, exprExplicit.toString)
  }
+
+ test("cg with non-grouped math aggr") {
+ val input = "name,foo,:eq,:sum,(,a,),:by,5,:mul,:sum,(,b,),:cg"
+ val inputExplicit = "name,foo,:eq,:sum,(,a,b,),:by,5,:mul,:sum,(,b,),:by"
+
+ val expr = eval(input)
+ val exprExplicit = eval(inputExplicit)
+ assertEquals(expr.toString, exprExplicit.toString)
+ }
 }