Multi source aggregate

cqframework · Nov 2, 2023 · 308b143 · 308b143
1 parent 2ba9d37
commit 308b143
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 22 deletions.
diff --git a/...gine/src/main/java/org/opencds/cqf/cql/engine/elm/executing/AggregateClauseEvaluator.java b/...gine/src/main/java/org/opencds/cqf/cql/engine/elm/executing/AggregateClauseEvaluator.java
@@ -2,8 +2,10 @@
 
 import org.cqframework.cql.elm.visiting.ElmLibraryVisitor;
 import org.hl7.elm.r1.AggregateClause;
+import org.opencds.cqf.cql.engine.exception.CqlException;
 import org.opencds.cqf.cql.engine.execution.State;
 import org.opencds.cqf.cql.engine.execution.Variable;
+import org.opencds.cqf.cql.engine.runtime.Tuple;
 
 import java.util.List;
 import java.util.Objects;
@@ -20,7 +22,7 @@
 
 public class AggregateClauseEvaluator {
 
-    public static Object aggregate(AggregateClause elm, String alias, State state, ElmLibraryVisitor<Object, State> visitor, List<Object> elements) {
+    public static Object aggregate(AggregateClause elm, State state, ElmLibraryVisitor<Object, State> visitor, List<Object> elements) {
         Objects.requireNonNull(elm, "elm can not be null");
         Objects.requireNonNull(elements, "elements can not be null");
         Objects.requireNonNull(state, "state can not be null");
@@ -36,15 +38,25 @@ public static Object aggregate(AggregateClause elm, String alias, State state, E
 
         for(var e : elements) {
             state.push(new Variable().withName(elm.getIdentifier()).withValue(aggregatedValue));
-            state.push(new Variable().withName(alias).withValue(e));
+            int pushes = 1;
+            if (!(e instanceof Tuple)) {
+                throw new CqlException("expected aggregation source to be a Tuple");
+            }
+            var tuple = (Tuple)e;
+            for (var p : tuple.getElements().entrySet()) {
+                state.push(new Variable().withName(p.getKey()).withValue(p.getValue()));
+                pushes++;
+            }
+
             try {
                 aggregatedValue= visitor.visitExpression(elm.getExpression(), state);
             }
             finally {
-                state.pop();
-                state.pop();
+                while(pushes > 0) {
+                    state.pop();
+                    pushes--;
+                }
             }
-
         }
 
         return aggregatedValue;

diff --git a/Src/java/engine/src/main/java/org/opencds/cqf/cql/engine/elm/executing/QueryEvaluator.java b/Src/java/engine/src/main/java/org/opencds/cqf/cql/engine/elm/executing/QueryEvaluator.java
@@ -87,18 +87,22 @@ private static Object evaluateReturn(Query elm, State state, List<Variable> vari
                 : constructResult(state, variables, elements);
     }
 
-    private static List<Object> evaluateAggregate(AggregateClause elm, String alias, State state, ElmLibraryVisitor<Object, State> visitor, List<Object> elements) {
-        return Collections.singletonList(AggregateClauseEvaluator.aggregate(elm, alias, state, visitor, elements));
+    private static List<Object> evaluateAggregate(AggregateClause elm, State state, ElmLibraryVisitor<Object, State> visitor, List<Object> elements) {
+        return Collections.singletonList(AggregateClauseEvaluator.aggregate(elm, state, visitor, elements));
+    }
+
+    private static Object constructTuple(State state, List<Variable> variables, List<Object> elements) {
+        LinkedHashMap<String, Object> elementMap = new LinkedHashMap<>();
+        for (int i = 0; i < variables.size(); i++) {
+            elementMap.put(variables.get(i).getName(), variables.get(i).getValue());
+        }
+
+        return new Tuple(state).withElements(elementMap);
     }
 
     private static Object constructResult(State state, List<Variable> variables, List<Object> elements) {
         if (variables.size() > 1) {
-            LinkedHashMap<String, Object> elementMap = new LinkedHashMap<>();
-            for (int i = 0; i < variables.size(); i++) {
-                elementMap.put(variables.get(i).getName(), variables.get(i).getValue());
-            }
-
-            return new Tuple(state).withElements(elementMap);
+            return constructTuple(state, variables, elements);
         }
 
         return elements.get(0);
@@ -156,6 +160,9 @@ public Iterable<Object> getData() {
 
     @SuppressWarnings("unchecked")
     public static Object internalEvaluate(Query elm, State state, ElmLibraryVisitor<Object, State> visitor) {
+        if (elm.getAggregate() != null && elm.getReturn() != null) {
+            throw new CqlException("aggregate and return are mutually exclusive");
+        }
 
         var sources = new ArrayList<Iterator<Object>>();
         var variables = new ArrayList<Variable>();
@@ -203,7 +210,13 @@ public static Object internalEvaluate(Query elm, State state, ElmLibraryVisitor<
                     continue;
                 }
 
-                result.add(evaluateReturn(elm, state, variables, elements, visitor));
+                if (elm.getAggregate() != null) {
+                    result.add(constructTuple(state, variables, elements));
+                }
+                else {
+                    result.add(evaluateReturn(elm, state, variables, elements, visitor));
+                }
+
             }
         } finally {
             while (pushCount > 0) {
@@ -212,18 +225,12 @@ public static Object internalEvaluate(Query elm, State state, ElmLibraryVisitor<
             }
         }
 
-        if (elm.getAggregate() != null && elm.getReturn() != null) {
-            throw new CqlException("aggregate and return are mutually exclusive");
-        }
-
         if (elm.getReturn() != null && elm.getReturn().isDistinct()) {
             result = DistinctEvaluator.distinct(result, state);
         }
 
         if (elm.getAggregate() != null) {
-            // TODO: JP - Hmm... this ain't right...
-            var alias = elm.getSource().get(0).getAlias();
-            result = evaluateAggregate(elm.getAggregate(), alias, state, visitor, result);
+            result = evaluateAggregate(elm.getAggregate(), state, visitor, result);
         }
 
         sortResult(elm, result, state, null, visitor);

diff --git a/.../engine/src/test/resources/org/opencds/cqf/cql/engine/execution/CqlAggregateQueryTest.cql b/.../engine/src/test/resources/org/opencds/cqf/cql/engine/execution/CqlAggregateQueryTest.cql
@@ -31,7 +31,7 @@ define "B": {1, 2}
 define "C": {1, 2}
 
 define MegaMulti:
-   from "1" X, "2" Y, "3" Z
+   from "A" X, "B" Y, "C" Z
     aggregate Agg starting 0: Agg + X + Y + Z // 36 -- (1+1+1)+(1+1+2)+(1+2+1)+(1+2+2)+(2+1+1)+(2+1+2)+(2+2+1)+(2+2+2)