Use separate maps for partitioned and not partitioned window functions

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

@@ -15,6 +15,8 @@ import java.util.Map.Entry;
 import org.h2.engine.Session;
 import org.h2.expression.Expression;
 import org.h2.expression.aggregate.DataAnalysisOperation;
+import org.h2.expression.aggregate.PartitionData;
+import org.h2.util.ValueHashMap;
 import org.h2.value.Value;
 import org.h2.value.ValueArray;
 
@@ -209,9 +211,14 @@ public abstract class SelectGroups {
     final HashMap<Expression, Integer> exprToIndexInGroupByData = new HashMap<>();
 
     /**
-     * Maps an expression object to its data.
+     * Maps an window expression object to its data.
      */
-    private final HashMap<DataAnalysisOperation, Object> windowData = new HashMap<>();
+    private final HashMap<DataAnalysisOperation, PartitionData> windowData = new HashMap<>();
+
+    /**
+     * Maps an partitioned window expression object to its data.
+     */
+    private final HashMap<DataAnalysisOperation, ValueHashMap<PartitionData>> windowPartitionData = new HashMap<>();
 
     /**
      * The id of the current group.
@@ -291,10 +298,17 @@ public abstract class SelectGroups {
      *
      * @param expr
      *            expression
+     * @param partitionKey
+     *            a key of partition
      * @return expression data or null
      */
-    public final Object getWindowExprData(DataAnalysisOperation expr) {
+    public final PartitionData getWindowExprData(DataAnalysisOperation expr, ValueArray partitionKey) {
+        if (partitionKey == null) {
             return windowData.get(expr);
+        } else {
+            ValueHashMap<PartitionData> map = windowPartitionData.get(expr);
+            return map != null ? map.get(partitionKey) : null;
+        }
     }
 
     /**
@@ -302,12 +316,23 @@ public abstract class SelectGroups {
      *
      * @param expr
      *            expression
+     * @param partitionKey
+     *            a key of partition
      * @param object
-     *            expression data to set
+     *            window expression data to set
      */
-    public final void setWindowExprData(DataAnalysisOperation expr, Object obj) {
+    public final void setWindowExprData(DataAnalysisOperation expr, ValueArray partitionKey, PartitionData obj) {
+        if (partitionKey == null) {
             Object old = windowData.put(expr, obj);
             assert old == null;
+        } else {
+            ValueHashMap<PartitionData> map = windowPartitionData.get(expr);
+            if (map == null) {
+                map = new ValueHashMap<>();
+                windowPartitionData.put(expr, map);
+            }
+            map.put(partitionKey, obj);
+        }
     }
 
     abstract void updateCurrentGroupExprData();
@@ -329,6 +354,7 @@ public abstract class SelectGroups {
         currentGroupByExprData = null;
         exprToIndexInGroupByData.clear();
         windowData.clear();
+        windowPartitionData.clear();
         currentGroupRowId = 0;
     }
 

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

@@ -20,7 +20,6 @@ import org.h2.message.DbException;
 import org.h2.result.SortOrder;
 import org.h2.table.ColumnResolver;
 import org.h2.table.TableFilter;
-import org.h2.util.ValueHashMap;
 import org.h2.value.Value;
 import org.h2.value.ValueArray;
 import org.h2.value.ValueInt;
@@ -187,38 +186,16 @@ public abstract class DataAnalysisOperation extends Expression {
         Object data;
         if (over != null) {
             ValueArray key = over.getCurrentKey(session);
-            if (key != null) {
-                @SuppressWarnings("unchecked")
-                ValueHashMap<Object> map = (ValueHashMap<Object>) groupData.getWindowExprData(this);
-                if (map == null) {
-                    if (ifExists) {
-                        return null;
-                    }
-                    map = new ValueHashMap<>();
-                    groupData.setWindowExprData(this, map);
-                }
-                PartitionData partition = (PartitionData) map.get(key);
+            PartitionData partition = groupData.getWindowExprData(this, key);
             if (partition == null) {
                 if (ifExists) {
                     return null;
                 }
                 data = forOrderBy ? new ArrayList<>() : createAggregateData();
-                    map.put(key, new PartitionData(data));
+                groupData.setWindowExprData(this, key, new PartitionData(data));
             } else {
                 data = partition.getData();
             }
-            } else {
-                PartitionData partition = (PartitionData) groupData.getWindowExprData(this);
-                if (partition == null) {
-                    if (ifExists) {
-                        return null;
-                    }
-                    data = forOrderBy ? new ArrayList<>() : createAggregateData();
-                    groupData.setWindowExprData(this, new PartitionData(data));
-                } else {
-                    data = partition.getData();
-                }
-            }
         } else {
             data = groupData.getCurrentGroupExprData(this);
             if (data == null) {
@@ -273,31 +250,14 @@ public abstract class DataAnalysisOperation extends Expression {
         Object data;
         boolean forOrderBy = over.getOrderBy() != null;
         ValueArray key = over.getCurrentKey(session);
-        if (key != null) {
-            @SuppressWarnings("unchecked")
-            ValueHashMap<Object> map = (ValueHashMap<Object>) groupData.getWindowExprData(this);
-            if (map == null) {
-                map = new ValueHashMap<>();
-                groupData.setWindowExprData(this, map);
-            }
-            partition = (PartitionData) map.get(key);
-            if (partition == null) {
-                data = forOrderBy ? new ArrayList<>() : createAggregateData();
-                partition = new PartitionData(data);
-                map.put(key, partition);
-            } else {
-                data = partition.getData();
-            }
-        } else {
-            partition = (PartitionData) groupData.getWindowExprData(this);
+        partition = groupData.getWindowExprData(this, key);
         if (partition == null) {
             data = forOrderBy ? new ArrayList<>() : createAggregateData();
             partition = new PartitionData(data);
-                groupData.setWindowExprData(this, partition);
+            groupData.setWindowExprData(this, key, partition);
         } else {
             data = partition.getData();
         }
-        }
         if (over.getOrderBy() != null || !isAggregate()) {
             return getOrderedResult(session, groupData, partition, data);
         }

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

@@ -12,7 +12,7 @@ import org.h2.value.Value;
 /**
  * Partition data of a window aggregate.
  */
-final class PartitionData {
+public final class PartitionData {
 
     /**
      * Aggregate data.