Extract SelectGroups from Select

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

@@ -8,9 +8,7 @@ package org.h2.command.dml;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.BitSet;
-import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Map;
 import org.h2.api.ErrorCode;
 import org.h2.api.Trigger;
 import org.h2.command.Parser;
@@ -103,23 +101,7 @@ public class Select extends Query {
      */
     boolean[] groupByExpression;
 
-    /**
-     * The array of current group-by expression data e.g. AggregateData.
-     */
-    Object[] currentGroupByExprData;
-    /**
-     * Maps an expression object to an index, to use in accessing the Object[]
-     * pointed to by groupByData.
-     */
-    final HashMap<Expression,Integer> exprToIndexInGroupByData = new HashMap<>();
-    /**
-     * Map of group-by key to group-by expression data e.g. AggregateData
-     */
-    private HashMap<Value, Object[]> groupByData;
-    /**
-     * Key into groupByData that produces currentGroupByExprData. Not used in lazy mode.
-     */
-    ValueArray currentGroupsKey;
+    SelectGroups groupData;
 
     private int havingIndex;
     private boolean isGroupQuery, isGroupSortedQuery;
@@ -129,11 +111,6 @@ public class Select extends Query {
     private boolean isPrepared, checkInit;
     private boolean sortUsingIndex;
 
-    /**
-     * The id of the current group.
-     */
-    int currentGroupRowId;
-
     public Select(Session session) {
         super(session);
     }
@@ -191,49 +168,8 @@ public class Select extends Query {
         return group;
     }
 
-    /**
-     * Is there currently a group-by active
-     */
-    public boolean isCurrentGroup() {
-        return currentGroupByExprData != null;
-    }
-
-    /**
-     * Get the group-by data for the current group and the passed in expression.
-     */
-    public Object getCurrentGroupExprData(Expression expr) {
-        Integer index = exprToIndexInGroupByData.get(expr);
-        if (index == null) {
-            return null;
-        }
-        return currentGroupByExprData[index];
-    }
-
-    /**
-     * Set the group-by data for the current group and the passed in expression.
-     */
-    public void setCurrentGroupExprData(Expression expr, Object obj) {
-        Integer index = exprToIndexInGroupByData.get(expr);
-        if (index != null) {
-            assert currentGroupByExprData[index] == null;
-            currentGroupByExprData[index] = obj;
-            return;
-        }
-        index = exprToIndexInGroupByData.size();
-        exprToIndexInGroupByData.put(expr, index);
-        if (index >= currentGroupByExprData.length) {
-            currentGroupByExprData = Arrays.copyOf(currentGroupByExprData, currentGroupByExprData.length * 2);
-            // this can be null in lazy mode
-            if (currentGroupsKey != null) {
-                // since we changed the size of the array, update the object in the groups map
-                groupByData.put(currentGroupsKey, currentGroupByExprData);
-            }
-        }
-        currentGroupByExprData[index] = obj;
-    }
-
-    public int getCurrentGroupRowId() {
-        return currentGroupRowId;
+    public SelectGroups getGroupDataIfCurrent() {
+        return groupData != null && groupData.isCurrentGroup() ? groupData : null;
     }
 
     @Override
@@ -411,38 +347,19 @@ public class Select extends Query {
     }
 
     private void queryGroup(int columnCount, LocalResult result, long offset, boolean quickOffset) {
-        groupByData = new HashMap<>();
-        currentGroupByExprData = null;
-        currentGroupsKey = null;
-        exprToIndexInGroupByData.clear();
+        if (groupData == null) {
+            groupData = new SelectGroups(session, expressions, groupIndex);
+        }
+        groupData.reset();
         try {
             int rowNumber = 0;
             setCurrentRowNumber(0);
-            ValueArray defaultGroup = ValueArray.get(new Value[0]);
             int sampleSize = getSampleSizeValue(session);
             while (topTableFilter.next()) {
                 setCurrentRowNumber(rowNumber + 1);
                 if (isConditionMet()) {
                     rowNumber++;
-                    if (groupIndex == null) {
-                        currentGroupsKey = defaultGroup;
-                    } else {
-                        Value[] keyValues = new Value[groupIndex.length];
-                        // update group
-                        for (int i = 0; i < groupIndex.length; i++) {
-                            int idx = groupIndex[i];
-                            Expression expr = expressions.get(idx);
-                            keyValues[i] = expr.getValue(session);
-                        }
-                        currentGroupsKey = ValueArray.get(keyValues);
-                    }
-                    Object[] values = groupByData.get(currentGroupsKey);
-                    if (values == null) {
-                        values = new Object[Math.max(exprToIndexInGroupByData.size(), expressions.size())];
-                        groupByData.put(currentGroupsKey, values);
-                    }
-                    currentGroupByExprData = values;
-                    currentGroupRowId++;
+                    groupData.nextSource();
                     for (int i = 0; i < columnCount; i++) {
                         if (groupByExpression == null || !groupByExpression[i]) {
                             Expression expr = expressions.get(i);
@@ -454,13 +371,8 @@ public class Select extends Query {
                     }
                 }
             }
-            if (groupIndex == null && groupByData.size() == 0) {
-                groupByData.put(defaultGroup,
-                        new Object[Math.max(exprToIndexInGroupByData.size(), expressions.size())]);
-            }
-            for (Map.Entry<Value, Object[]> entry : groupByData.entrySet()) {
-                currentGroupsKey = (ValueArray) entry.getKey();
-                currentGroupByExprData = entry.getValue();
+            groupData.done();
+            for (ValueArray currentGroupsKey; (currentGroupsKey = groupData.next()) != null;) {
                 Value[] keyValues = currentGroupsKey.getList();
                 Value[] row = new Value[columnCount];
                 for (int j = 0; groupIndex != null && j < groupIndex.length; j++) {
@@ -484,10 +396,7 @@ public class Select extends Query {
                 result.addRow(row);
             }
         } finally {
-            groupByData = null;
-            currentGroupsKey = null;
-            currentGroupByExprData = null;
-            exprToIndexInGroupByData.clear();
+            groupData.reset();
         }
     }
 
@@ -1676,19 +1585,23 @@ public class Select extends Query {
      */
     private final class LazyResultGroupSorted extends LazyResultSelect {
 
-        Value[] previousKeyValues;
+        private Value[] previousKeyValues;
 
         LazyResultGroupSorted(Expression[] expressions, int columnCount) {
             super(expressions, columnCount);
-            currentGroupByExprData = null;
-            currentGroupsKey = null;
+            if (groupData == null) {
+                groupData = new SelectGroups(getSession(), Select.this.expressions, groupIndex);
+            } else {
+                // TODO is this branch possible?
+                groupData.resetLazy();
+            }
         }
 
         @Override
         public void reset() {
             super.reset();
-            currentGroupByExprData = null;
-            currentGroupsKey = null;
+            groupData.resetLazy();
+            previousKeyValues = null;
         }
 
         @Override
@@ -1708,15 +1621,13 @@ public class Select extends Query {
                     Value[] row = null;
                     if (previousKeyValues == null) {
                         previousKeyValues = keyValues;
-                        currentGroupByExprData =new Object[Math.max(exprToIndexInGroupByData.size(),
-                                expressions.size())];
+                        groupData.nextLazyGroup();
                     } else if (!Arrays.equals(previousKeyValues, keyValues)) {
                         row = createGroupSortedRow(previousKeyValues, columnCount);
                         previousKeyValues = keyValues;
-                        currentGroupByExprData = new Object[Math.max(exprToIndexInGroupByData.size(),
-                                expressions.size())];
+                        groupData.nextLazyGroup();
                     }
-                    currentGroupRowId++;
+                    groupData.nextLazyRow();
 
                     for (int i = 0; i < columnCount; i++) {
                         if (groupByExpression == null || !groupByExpression[i]) {

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

+/*
+ * Copyright 2004-2018 H2 Group. Multiple-Licensed under the MPL 2.0,
+ * and the EPL 1.0 (http://h2database.com/html/license.html).
+ * Initial Developer: H2 Group
+ */
+package org.h2.command.dml;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Map.Entry;
+
+import org.h2.engine.Session;
+import org.h2.expression.Expression;
+import org.h2.value.Value;
+import org.h2.value.ValueArray;
+
+/**
+ * Grouped data for aggregates.
+ *
+ * <p>
+ * Call sequence:
+ * </p>
+ * <ul>
+ * <li>{@link #reset()} (not required before the first execution).</li>
+ * <li>For each source row {@link #nextSource()} should be invoked.</li>
+ * <li>{@link #done()}.</li>
+ * <li>{@link #next()} is invoked inside a loop until it returns null.</li>
+ * </ul>
+ * <p>
+ * Call sequence for lazy group sorted result:
+ * </p>
+ * <ul>
+ * <li>{@link #resetLazy()} (not required before the first execution).</li>
+ * <li>For each source group {@link #nextLazyGroup()} should be invoked.</li>
+ * <li>For each source row {@link #nextLazyRow()} should be invoked. Each group
+ * can have one or more rows.</li>
+ * </ul>
+ */
+public final class SelectGroups {
+
+    private final Session session;
+
+    private final ArrayList<Expression> expressions;
+
+    private final int[] groupIndex;
+
+    /**
+     * The array of current group-by expression data e.g. AggregateData.
+     */
+    private Object[] currentGroupByExprData;
+
+    /**
+     * Maps an expression object to an index, to use in accessing the Object[]
+     * pointed to by groupByData.
+     */
+    private final HashMap<Expression, Integer> exprToIndexInGroupByData = new HashMap<>();
+
+    /**
+     * Map of group-by key to group-by expression data e.g. AggregateData
+     */
+    private HashMap<ValueArray, Object[]> groupByData;
+
+    /**
+     * Key into groupByData that produces currentGroupByExprData. Not used in
+     * lazy mode.
+     */
+    private ValueArray currentGroupsKey;
+
+    /**
+     * The id of the current group.
+     */
+    private int currentGroupRowId;
+
+    /**
+     * The key for the default group.
+     */
+    // Can be static, but TestClearReferences complains about it
+    private ValueArray defaultGroup = ValueArray.get(new Value[0]);
+
+    /**
+     * Cursor for {@link #next()} method.
+     */
+    private Iterator<Entry<ValueArray, Object[]>> cursor;
+
+    /**
+     * Creates new instance of grouped data.
+     *
+     * @param session
+     *            the session
+     * @param expressions
+     *            the expressions
+     * @param groupIndex
+     *            the indexes of group expressions, or null
+     */
+    public SelectGroups(Session session, ArrayList<Expression> expressions, int[] groupIndex) {
+        this.session = session;
+        this.expressions = expressions;
+        this.groupIndex = groupIndex;
+    }
+
+    /**
+     * Is there currently a group-by active
+     */
+    public boolean isCurrentGroup() {
+        return currentGroupByExprData != null;
+    }
+
+    /**
+     * Get the group-by data for the current group and the passed in expression.
+     */
+    public Object getCurrentGroupExprData(Expression expr) {
+        Integer index = exprToIndexInGroupByData.get(expr);
+        if (index == null) {
+            return null;
+        }
+        return currentGroupByExprData[index];
+    }
+
+    /**
+     * Set the group-by data for the current group and the passed in expression.
+     */
+    public void setCurrentGroupExprData(Expression expr, Object obj) {
+        Integer index = exprToIndexInGroupByData.get(expr);
+        if (index != null) {
+            assert currentGroupByExprData[index] == null;
+            currentGroupByExprData[index] = obj;
+            return;
+        }
+        index = exprToIndexInGroupByData.size();
+        exprToIndexInGroupByData.put(expr, index);
+        if (index >= currentGroupByExprData.length) {
+            currentGroupByExprData = Arrays.copyOf(currentGroupByExprData, currentGroupByExprData.length * 2);
+            // this can be null in lazy mode
+            if (currentGroupsKey != null) {
+                // since we changed the size of the array, update the object in
+                // the groups map
+                groupByData.put(currentGroupsKey, currentGroupByExprData);
+            }
+        }
+        currentGroupByExprData[index] = obj;
+    }
+
+    /**
+     * Returns identity of the current row. Used by aggregates to check whether
+     * they already processed this row or not.
+     *
+     * @return identity of the current row
+     */
+    public int getCurrentGroupRowId() {
+        return currentGroupRowId;
+    }
+
+    /**
+     * Resets this group data for reuse.
+     */
+    public void reset() {
+        groupByData = new HashMap<>();
+        currentGroupByExprData = null;
+        currentGroupsKey = null;
+        exprToIndexInGroupByData.clear();
+        cursor = null;
+    }
+
+    /**
+     * Invoked for each source row to evaluate group key and setup all necessary
+     * data for aggregates.
+     */
+    public void nextSource() {
+        if (groupIndex == null) {
+            currentGroupsKey = defaultGroup;
+        } else {
+            Value[] keyValues = new Value[groupIndex.length];
+            // update group
+            for (int i = 0; i < groupIndex.length; i++) {
+                int idx = groupIndex[i];
+                Expression expr = expressions.get(idx);
+                keyValues[i] = expr.getValue(session);
+            }
+            currentGroupsKey = ValueArray.get(keyValues);
+        }
+        Object[] values = groupByData.get(currentGroupsKey);
+        if (values == null) {
+            values = new Object[Math.max(exprToIndexInGroupByData.size(), expressions.size())];
+            groupByData.put(currentGroupsKey, values);
+        }
+        currentGroupByExprData = values;
+        currentGroupRowId++;
+    }
+
+    /**
+     * Invoked after all source rows are evaluated.
+     */
+    public void done() {
+        if (groupIndex == null && groupByData.size() == 0) {
+            groupByData.put(defaultGroup, new Object[Math.max(exprToIndexInGroupByData.size(), expressions.size())]);
+        }
+        cursor = groupByData.entrySet().iterator();
+    }
+
+    /**
+     * Returns the key of the next group.
+     *
+     * @return the key of the next group, or null
+     */
+    public ValueArray next() {
+        if (cursor.hasNext()) {
+            Map.Entry<ValueArray, Object[]> entry = cursor.next();
+            currentGroupByExprData = entry.getValue();
+            return entry.getKey();
+        }
+        return null;
+    }
+
+    /**
+     * Resets this group data for reuse in lazy mode.
+     */
+    public void resetLazy() {
+        currentGroupByExprData = null;
+        currentGroupsKey = null;
+    }
+
+    /**
+     * Moves group data to the next group in lazy mode.
+     */
+    public void nextLazyGroup() {
+        currentGroupByExprData = new Object[Math.max(exprToIndexInGroupByData.size(), expressions.size())];
+    }
+
+    /**
+     * Moves group data to the next row in lazy mode.
+     */
+    public void nextLazyRow() {
+        currentGroupRowId++;
+    }
+
+}

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

@@ -8,6 +8,7 @@ package org.h2.expression;
 import org.h2.api.ErrorCode;
 import org.h2.command.Parser;
 import org.h2.command.dml.Select;
+import org.h2.command.dml.SelectGroups;
 import org.h2.command.dml.SelectListColumnResolver;
 import org.h2.engine.Database;
 import org.h2.engine.Session;
@@ -158,13 +159,14 @@ public class ExpressionColumn extends Expression {
         if (select == null) {
             throw DbException.get(ErrorCode.MUST_GROUP_BY_COLUMN_1, getSQL());
         }
-        if (!select.isCurrentGroup()) {
+        SelectGroups groupData = select.getGroupDataIfCurrent();
+        if (groupData == null) {
             // this is a different level (the enclosing query)
             return;
         }
-        Value v = (Value) select.getCurrentGroupExprData(this);
+        Value v = (Value) groupData.getCurrentGroupExprData(this);
         if (v == null) {
-            select.setCurrentGroupExprData(this, now);
+            groupData.setCurrentGroupExprData(this, now);
         } else {
             if (!database.areEqual(now, v)) {
                 throw DbException.get(ErrorCode.MUST_GROUP_BY_COLUMN_1, getSQL());
@@ -176,8 +178,9 @@ public class ExpressionColumn extends Expression {
     public Value getValue(Session session) {
         Select select = columnResolver.getSelect();
         if (select != null) {
-            if (select.isCurrentGroup()) {
-                Value v = (Value) select.getCurrentGroupExprData(this);
+            SelectGroups groupData = select.getGroupDataIfCurrent();
+            if (groupData != null) {
+                Value v = (Value) groupData.getCurrentGroupExprData(this);
                 if (v != null) {
                     return v;
                 }

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

@@ -11,6 +11,7 @@ import java.util.Comparator;
 import java.util.HashMap;
 import org.h2.api.ErrorCode;
 import org.h2.command.dml.Select;
+import org.h2.command.dml.SelectGroups;
 import org.h2.command.dml.SelectOrderBy;
 import org.h2.engine.Session;
 import org.h2.expression.Expression;
@@ -295,12 +296,13 @@ public class Aggregate extends Expression {
         // if (on != null) {
         // on.updateAggregate();
         // }
-        if (!select.isCurrentGroup()) {
+        SelectGroups groupData = select.getGroupDataIfCurrent();
+        if (groupData == null) {
             // this is a different level (the enclosing query)
             return;
         }
 
-        int groupRowId = select.getCurrentGroupRowId();
+        int groupRowId = groupData.getCurrentGroupRowId();
         if (lastGroupRowId == groupRowId) {
             // already visited
             return;
@@ -312,10 +314,10 @@ public class Aggregate extends Expression {
                 return;
             }
         }
-        AggregateData data = (AggregateData) select.getCurrentGroupExprData(this);
+        AggregateData data = (AggregateData) groupData.getCurrentGroupExprData(this);
         if (data == null) {
             data = AggregateData.create(type);
-            select.setCurrentGroupExprData(this, data);
+            groupData.setCurrentGroupExprData(this, data);
         }
         Value v = on == null ? null : on.getValue(session);
         if (type == AggregateType.GROUP_CONCAT) {
@@ -378,13 +380,14 @@ public class Aggregate extends Expression {
                 DbException.throwInternalError("type=" + type);
             }
         }
-        if (!select.isCurrentGroup()) {
+        SelectGroups groupData = select.getGroupDataIfCurrent();
+        if (groupData == null) {
             throw DbException.get(ErrorCode.INVALID_USE_OF_AGGREGATE_FUNCTION_1, getSQL());
         }
-        AggregateData data = (AggregateData)select.getCurrentGroupExprData(this);
+        AggregateData data = (AggregateData) groupData.getCurrentGroupExprData(this);
         if (data == null) {
             data = AggregateData.create(type);
-            select.setCurrentGroupExprData(this, data);
+            groupData.setCurrentGroupExprData(this, data);
         }
         switch (type) {
         case GROUP_CONCAT: {

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

@@ -11,6 +11,7 @@ import org.h2.api.Aggregate;
 import org.h2.api.ErrorCode;
 import org.h2.command.Parser;
 import org.h2.command.dml.Select;
+import org.h2.command.dml.SelectGroups;
 import org.h2.engine.Session;
 import org.h2.engine.UserAggregate;
 import org.h2.expression.Expression;
@@ -168,14 +169,15 @@ public class JavaAggregate extends Expression {
 
     @Override
     public Value getValue(Session session) {
-        if (!select.isCurrentGroup()) {
+        SelectGroups groupData = select.getGroupDataIfCurrent();
+        if (groupData == null) {
             throw DbException.get(ErrorCode.INVALID_USE_OF_AGGREGATE_FUNCTION_1, getSQL());
         }
         try {
             Aggregate agg;
             if (distinct) {
                 agg = getInstance();
-                AggregateDataCollecting data = (AggregateDataCollecting) select.getCurrentGroupExprData(this);
+                AggregateDataCollecting data = (AggregateDataCollecting) groupData.getCurrentGroupExprData(this);
                 if (data != null) {
                     for (Value value : data.values) {
                         if (args.length == 1) {
@@ -191,7 +193,7 @@ public class JavaAggregate extends Expression {
                     }
                 }
             } else {
-                agg = (Aggregate) select.getCurrentGroupExprData(this);
+                agg = (Aggregate) groupData.getCurrentGroupExprData(this);
                 if (agg == null) {
                     agg = getInstance();
                 }
@@ -208,12 +210,13 @@ public class JavaAggregate extends Expression {
 
     @Override
     public void updateAggregate(Session session) {
-        if (!select.isCurrentGroup()) {
+        SelectGroups groupData = select.getGroupDataIfCurrent();
+        if (groupData == null) {
             // this is a different level (the enclosing query)
             return;
         }
 
-        int groupRowId = select.getCurrentGroupRowId();
+        int groupRowId = groupData.getCurrentGroupRowId();
         if (lastGroupRowId == groupRowId) {
             // already visited
             return;
@@ -228,10 +231,10 @@ public class JavaAggregate extends Expression {
 
         try {
             if (distinct) {
-                AggregateDataCollecting data = (AggregateDataCollecting) select.getCurrentGroupExprData(this);
+                AggregateDataCollecting data = (AggregateDataCollecting) groupData.getCurrentGroupExprData(this);
                 if (data == null) {
                     data = new AggregateDataCollecting();
-                    select.setCurrentGroupExprData(this, data);
+                    groupData.setCurrentGroupExprData(this, data);
                 }
                 Value[] argValues = new Value[args.length];
                 Value arg = null;
@@ -242,10 +245,10 @@ public class JavaAggregate extends Expression {
                 }
                 data.add(session.getDatabase(), dataType, true, args.length == 1 ? arg : ValueArray.get(argValues));
             } else {
-                Aggregate agg = (Aggregate) select.getCurrentGroupExprData(this);
+                Aggregate agg = (Aggregate) groupData.getCurrentGroupExprData(this);
                 if (agg == null) {
                     agg = getInstance();
-                    select.setCurrentGroupExprData(this, agg);
+                    groupData.setCurrentGroupExprData(this, agg);
                 }
                 Object[] argValues = new Object[args.length];
                 Object arg = null;