Merge pull request #1452 from katzyn/aggregate

Reset aggregates before reuse

h2/src/main/org/h2/command/dml/Query.java

@@ -276,9 +276,9 @@ public abstract class Query extends Prepared {
      * Update all aggregate function values.
      *
      * @param s the session
-     * @param window true for window processing stage, false for group stage
+     * @param stage select stage
      */
-    public abstract void updateAggregate(Session s, boolean window);
+    public abstract void updateAggregate(Session s, int stage);
 
     /**
      * Call the before triggers on all tables.

h2/src/main/org/h2/command/dml/Select.java

@@ -24,6 +24,7 @@ import org.h2.expression.ExpressionColumn;
 import org.h2.expression.ExpressionVisitor;
 import org.h2.expression.Parameter;
 import org.h2.expression.Wildcard;
+import org.h2.expression.aggregate.Aggregate;
 import org.h2.index.Cursor;
 import org.h2.index.Index;
 import org.h2.index.IndexType;
@@ -358,18 +359,9 @@ public class Select extends Query {
     }
 
     private void queryWindow(int columnCount, LocalResult result, long offset, boolean quickOffset) {
-        if (isGroupQuery) {
-            queryGroupWindow(columnCount, result, offset, quickOffset);
-            return;
-        }
-        if (groupData == null) {
-            groupData = SelectGroups.getInstance(session, expressions, isGroupQuery, groupIndex);
-        }
-        groupData.reset();
-        groupData.resetCounter();
+        initGroupData(columnCount);
         try {
-            gatherGroup(columnCount, true);
-            groupData.resetCounter();
+            gatherGroup(columnCount, Aggregate.STAGE_WINDOW);
             processGroupResult(columnCount, result, offset, quickOffset);
         } finally {
             groupData.reset();
@@ -377,19 +369,13 @@ public class Select extends Query {
     }
 
     private void queryGroupWindow(int columnCount, LocalResult result, long offset, boolean quickOffset) {
-        if (groupData == null) {
-            groupData = SelectGroups.getInstance(session, expressions, isGroupQuery, groupIndex);
-        }
-        groupData.reset();
-        groupData.resetCounter();
+        initGroupData(columnCount);
         try {
-            gatherGroup(columnCount, false);
-            groupData.resetCounter();
+            gatherGroup(columnCount, Aggregate.STAGE_GROUP);
             while (groupData.next() != null) {
-                updateAgg(columnCount, true);
+                updateAgg(columnCount, Aggregate.STAGE_WINDOW);
             }
             groupData.done();
-            groupData.resetCounter();
             try {
                 isGroupWindowStage2 = true;
                 processGroupResult(columnCount, result, offset, quickOffset);
@@ -402,19 +388,25 @@ public class Select extends Query {
     }
 
     private void queryGroup(int columnCount, LocalResult result, long offset, boolean quickOffset) {
-        if (groupData == null) {
-            groupData = SelectGroups.getInstance(session, expressions, isGroupQuery, groupIndex);
-        }
-        groupData.reset();
+        initGroupData(columnCount);
         try {
-            gatherGroup(columnCount, false);
+            gatherGroup(columnCount, Aggregate.STAGE_GROUP);
             processGroupResult(columnCount, result, offset, quickOffset);
         } finally {
             groupData.reset();
         }
     }
 
-    private void gatherGroup(int columnCount, boolean window) {
+    private void initGroupData(int columnCount) {
+        if (groupData == null) {
+            groupData = SelectGroups.getInstance(session, expressions, isGroupQuery, groupIndex);
+        } else {
+            updateAgg(columnCount, Aggregate.STAGE_RESET);
+        }
+        groupData.reset();
+    }
+
+    private void gatherGroup(int columnCount, int stage) {
         int rowNumber = 0;
         setCurrentRowNumber(0);
         int sampleSize = getSampleSizeValue(session);
@@ -423,7 +415,7 @@ public class Select extends Query {
             if (isConditionMet()) {
                 rowNumber++;
                 groupData.nextSource();
-                updateAgg(columnCount, window);
+                updateAgg(columnCount, stage);
                 if (sampleSize > 0 && rowNumber >= sampleSize) {
                     break;
                 }
@@ -432,11 +424,11 @@ public class Select extends Query {
         groupData.done();
     }
 
-    private void updateAgg(int columnCount, boolean window) {
+    void updateAgg(int columnCount, int stage) {
         for (int i = 0; i < columnCount; i++) {
             if (groupByExpression == null || !groupByExpression[i]) {
                 Expression expr = expressions.get(i);
-                expr.updateAggregate(session, window);
+                expr.updateAggregate(session, stage);
             }
         }
     }
@@ -1529,15 +1521,15 @@ public class Select extends Query {
     }
 
     @Override
-    public void updateAggregate(Session s, boolean window) {
+    public void updateAggregate(Session s, int stage) {
         for (Expression e : expressions) {
-            e.updateAggregate(s, window);
+            e.updateAggregate(s, stage);
         }
         if (condition != null) {
-            condition.updateAggregate(s, window);
+            condition.updateAggregate(s, stage);
         }
         if (having != null) {
-            having.updateAggregate(s, window);
+            having.updateAggregate(s, stage);
         }
     }
 
@@ -1694,6 +1686,7 @@ public class Select extends Query {
                 groupData = SelectGroups.getInstance(getSession(), Select.this.expressions, isGroupQuery, groupIndex);
             } else {
                 // TODO is this branch possible?
+                updateAgg(columnCount, Aggregate.STAGE_RESET);
                 groupData.resetLazy();
             }
         }
@@ -1729,13 +1722,7 @@ public class Select extends Query {
                         groupData.nextLazyGroup();
                     }
                     groupData.nextLazyRow();
-
-                    for (int i = 0; i < columnCount; i++) {
-                        if (groupByExpression == null || !groupByExpression[i]) {
-                            Expression expr = expressions.get(i);
-                            expr.updateAggregate(getSession(), false);
-                        }
-                    }
+                    updateAgg(columnCount, Aggregate.STAGE_GROUP);
                     if (row != null) {
                         return row;
                     }

h2/src/main/org/h2/command/dml/SelectGroups.java

@@ -115,6 +115,7 @@ public abstract class SelectGroups {
 
         @Override
         public void done() {
+            super.done();
             if (groupIndex == null && groupByData.size() == 0) {
                 groupByData.put(defaultGroup,
                         new Object[Math.max(exprToIndexInGroupByData.size(), expressions.size())]);
@@ -175,6 +176,7 @@ public abstract class SelectGroups {
 
         @Override
         public void done() {
+            super.done();
             cursor = rows.iterator();
         }
 
@@ -314,13 +316,6 @@ public abstract class SelectGroups {
         currentGroupByExprData = null;
         exprToIndexInGroupByData.clear();
         windowData.clear();
-    }
-
-    /**
-     * Reset the row id counter.
-     */
-    public void resetCounter() {
-        // TODO merge into reset() and done()
         currentGroupRowId = 0;
     }
 
@@ -333,7 +328,9 @@ public abstract class SelectGroups {
     /**
      * Invoked after all source rows are evaluated.
      */
-    public abstract void done();
+    public void done() {
+        currentGroupRowId = 0;
+    }
 
     /**
      * Returns the key of the next group.
@@ -347,6 +344,7 @@ public abstract class SelectGroups {
      */
     public void resetLazy() {
         currentGroupByExprData = null;
+        currentGroupRowId = 0;
     }
 
     /**

h2/src/main/org/h2/command/dml/SelectUnion.java

@@ -449,9 +449,9 @@ public class SelectUnion extends Query {
     }
 
     @Override
-    public void updateAggregate(Session s, boolean window) {
-        left.updateAggregate(s, window);
-        right.updateAggregate(s, window);
+    public void updateAggregate(Session s, int stage) {
+        left.updateAggregate(s, stage);
+        right.updateAggregate(s, stage);
     }
 
     @Override

h2/src/main/org/h2/expression/Alias.java

@@ -83,8 +83,8 @@ public class Alias extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        expr.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        expr.updateAggregate(session, stage);
     }
 
     @Override

h2/src/main/org/h2/expression/BinaryOperation.java

@@ -436,9 +436,9 @@ public class BinaryOperation extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        left.updateAggregate(session, window);
-        right.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        left.updateAggregate(session, stage);
+        right.updateAggregate(session, stage);
     }
 
     @Override

h2/src/main/org/h2/expression/CompareLike.java

@@ -495,11 +495,11 @@ public class CompareLike extends Condition {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        left.updateAggregate(session, window);
-        right.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        left.updateAggregate(session, stage);
+        right.updateAggregate(session, stage);
         if (escape != null) {
-            escape.updateAggregate(session, window);
+            escape.updateAggregate(session, stage);
         }
     }
 

h2/src/main/org/h2/expression/Comparison.java

@@ -475,10 +475,10 @@ public class Comparison extends Condition {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        left.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        left.updateAggregate(session, stage);
         if (right != null) {
-            right.updateAggregate(session, window);
+            right.updateAggregate(session, stage);
         }
     }
 

h2/src/main/org/h2/expression/ConditionAndOr.java

@@ -268,9 +268,9 @@ public class ConditionAndOr extends Condition {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        left.updateAggregate(session, window);
-        right.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        left.updateAggregate(session, stage);
+        right.updateAggregate(session, stage);
     }
 
     @Override

h2/src/main/org/h2/expression/ConditionExists.java

@@ -46,7 +46,7 @@ public class ConditionExists extends Condition {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         // TODO exists: is it allowed that the subquery contains aggregates?
         // probably not
         // select id from test group by id having exists (select * from test2

h2/src/main/org/h2/expression/ConditionIn.java

@@ -163,10 +163,10 @@ public class ConditionIn extends Condition {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        left.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        left.updateAggregate(session, stage);
         for (Expression e : valueList) {
-            e.updateAggregate(session, window);
+            e.updateAggregate(session, stage);
         }
     }
 

h2/src/main/org/h2/expression/ConditionInConstantSet.java

@@ -125,8 +125,8 @@ public class ConditionInConstantSet extends Condition {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        left.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        left.updateAggregate(session, stage);
     }
 
     @Override

h2/src/main/org/h2/expression/ConditionInParameter.java

@@ -145,8 +145,8 @@ public class ConditionInParameter extends Condition {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        left.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        left.updateAggregate(session, stage);
     }
 
     @Override

h2/src/main/org/h2/expression/ConditionInSelect.java

@@ -146,9 +146,9 @@ public class ConditionInSelect extends Condition {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        left.updateAggregate(session, window);
-        query.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        left.updateAggregate(session, stage);
+        query.updateAggregate(session, stage);
     }
 
     @Override

h2/src/main/org/h2/expression/ConditionNot.java

@@ -70,8 +70,8 @@ public class ConditionNot extends Condition {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        condition.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        condition.updateAggregate(session, stage);
     }
 
     @Override

h2/src/main/org/h2/expression/Expression.java

@@ -105,9 +105,9 @@ public abstract class Expression {
      * be used to make sure the internal state is only updated once.
      *
      * @param session the session
-     * @param window true for window processing stage, false for group stage
+     * @param stage select stage
      */
-    public abstract void updateAggregate(Session session, boolean window);
+    public abstract void updateAggregate(Session session, int stage);
 
     /**
      * Check if this expression and all sub-expressions can fulfill a criteria.

h2/src/main/org/h2/expression/ExpressionColumn.java

@@ -153,7 +153,7 @@ public class ExpressionColumn extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         Value now = columnResolver.getValue(column);
         Select select = columnResolver.getSelect();
         if (select == null) {

h2/src/main/org/h2/expression/ExpressionList.java

@@ -98,9 +98,9 @@ public class ExpressionList extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         for (Expression e : list) {
-            e.updateAggregate(session, window);
+            e.updateAggregate(session, stage);
         }
     }
 

h2/src/main/org/h2/expression/Function.java

@@ -2633,10 +2633,10 @@ public class Function extends Expression implements FunctionCall {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         for (Expression e : args) {
             if (e != null) {
-                e.updateAggregate(session, window);
+                e.updateAggregate(session, stage);
             }
         }
     }

h2/src/main/org/h2/expression/IntervalOperation.java

@@ -291,9 +291,9 @@ public class IntervalOperation extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        left.updateAggregate(session, window);
-        right.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        left.updateAggregate(session, stage);
+        right.updateAggregate(session, stage);
     }
 
     @Override

h2/src/main/org/h2/expression/JavaFunction.java

@@ -107,10 +107,10 @@ public class JavaFunction extends Expression implements FunctionCall {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         for (Expression e : args) {
             if (e != null) {
-                e.updateAggregate(session, window);
+                e.updateAggregate(session, stage);
             }
         }
     }

h2/src/main/org/h2/expression/Parameter.java

@@ -141,7 +141,7 @@ public class Parameter extends Expression implements ParameterInterface {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         // nothing to do
     }
 

h2/src/main/org/h2/expression/Rownum.java

@@ -73,7 +73,7 @@ public class Rownum extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         // nothing to do
     }
 

h2/src/main/org/h2/expression/SequenceValue.java

@@ -72,7 +72,7 @@ public class SequenceValue extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         // nothing to do
     }
 

h2/src/main/org/h2/expression/Subquery.java

@@ -94,8 +94,8 @@ public class Subquery extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        query.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        query.updateAggregate(session, stage);
     }
 
     private Expression getExpression() {

h2/src/main/org/h2/expression/UnaryOperation.java

@@ -82,8 +82,8 @@ public class UnaryOperation extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
-        arg.updateAggregate(session, window);
+    public void updateAggregate(Session session, int stage) {
+        arg.updateAggregate(session, stage);
     }
 
     @Override

h2/src/main/org/h2/expression/ValueExpression.java

@@ -143,7 +143,7 @@ public class ValueExpression extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         // nothing to do
     }
 

h2/src/main/org/h2/expression/Variable.java

@@ -101,7 +101,7 @@ public class Variable extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         // nothing to do
     }
 

h2/src/main/org/h2/expression/Wildcard.java

@@ -91,7 +91,7 @@ public class Wildcard extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         DbException.throwInternalError(toString());
     }
 

h2/src/main/org/h2/expression/aggregate/AbstractAggregate.java

@@ -110,14 +110,16 @@ public abstract class AbstractAggregate extends Expression {
     }
 
     @Override
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
+        if (stage == Aggregate.STAGE_RESET) {
+            updateSubAggregates(session, Aggregate.STAGE_RESET);
+            lastGroupRowId = 0;
+            return;
+        }
+        boolean window = stage == Aggregate.STAGE_WINDOW;
         if (window != (over != null)) {
             if (!window && select.isWindowQuery()) {
-                updateGroupAggregates(session);
-                if (filterCondition != null) {
-                    filterCondition.updateAggregate(session, false);
-                }
-                over.updateAggregate(session, false);
+                updateSubAggregates(session, stage);
             }
             return;
         }
@@ -140,7 +142,7 @@ public abstract class AbstractAggregate extends Expression {
 
         if (over != null) {
             if (!select.isGroupQuery()) {
-                over.updateAggregate(session, true);
+                over.updateAggregate(session, stage);
             }
         }
         if (filterCondition != null) {
@@ -158,6 +160,16 @@ public abstract class AbstractAggregate extends Expression {
         updateAggregate(session, getData(session, groupData, false, false));
     }
 
+    private void updateSubAggregates(Session session, int stage) {
+        updateGroupAggregates(session, stage);
+        if (filterCondition != null) {
+            filterCondition.updateAggregate(session, stage);
+        }
+        if (over != null) {
+            over.updateAggregate(session, stage);
+        }
+    }
+
     /**
      * Updates an aggregate value.
      *
@@ -174,8 +186,10 @@ public abstract class AbstractAggregate extends Expression {
      *
      * @param session
      *            the session
+     * @param stage
+     *            select stage
      */
-    protected abstract void updateGroupAggregates(Session session);
+    protected abstract void updateGroupAggregates(Session session, int stage);
 
     /**
      * Returns the number of expressions, excluding FILTER and OVER clauses.

h2/src/main/org/h2/expression/aggregate/Aggregate.java

@@ -150,6 +150,21 @@ public class Aggregate extends AbstractAggregate {
         ENVELOPE,
     }
 
+    /**
+     * Reset stage. Used to reset internal data to its initial state.
+     */
+    public static final int STAGE_RESET = 0;
+
+    /**
+     * Group stage, used for explicit or implicit GROUP BY operation.
+     */
+    public static final int STAGE_GROUP = 1;
+
+    /**
+     * Window processing stage.
+     */
+    public static final int STAGE_WINDOW = 2;
+
     private static final HashMap<String, AggregateType> AGGREGATES = new HashMap<>(64);
 
     private final AggregateType type;
@@ -291,13 +306,13 @@ public class Aggregate extends AbstractAggregate {
     }
 
     @Override
-    protected void updateGroupAggregates(Session session) {
+    protected void updateGroupAggregates(Session session, int stage) {
         if (on != null) {
-            on.updateAggregate(session, false);
+            on.updateAggregate(session, stage);
         }
         if (orderByList != null) {
             for (SelectOrderBy orderBy : orderByList) {
-                orderBy.expression.updateAggregate(session, false);
+                orderBy.expression.updateAggregate(session, stage);
             }
         }
     }

h2/src/main/org/h2/expression/aggregate/JavaAggregate.java

@@ -228,9 +228,9 @@ public class JavaAggregate extends AbstractAggregate {
     }
 
     @Override
-    protected void updateGroupAggregates(Session session) {
+    protected void updateGroupAggregates(Session session, int stage) {
         for (Expression expr : args) {
-            expr.updateAggregate(session, false);
+            expr.updateAggregate(session, stage);
         }
     }
 

h2/src/main/org/h2/expression/aggregate/Window.java

@@ -163,19 +163,19 @@ public final class Window {
      *
      * @param session
      *            the session
-     * @param window
-     *            true for window processing stage, false for group stage
-     * @see Expression#updateAggregate(Session, boolean)
+     * @param stage
+     *            select stage
+     * @see Expression#updateAggregate(Session, int)
      */
-    public void updateAggregate(Session session, boolean window) {
+    public void updateAggregate(Session session, int stage) {
         if (partitionBy != null) {
             for (Expression expr : partitionBy) {
-                expr.updateAggregate(session, window);
+                expr.updateAggregate(session, stage);
             }
         }
         if (orderBy != null) {
             for (SelectOrderBy o : orderBy) {
-                o.expression.updateAggregate(session, false);
+                o.expression.updateAggregate(session, stage);
             }
         }
     }