Implement Spilling for TopNRowNumber Operator

presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java

@@ -131,6 +131,7 @@ public final class SystemSessionProperties
     public static final String SPILL_ENABLED = "spill_enabled";
     public static final String JOIN_SPILL_ENABLED = "join_spill_enabled";
     public static final String AGGREGATION_SPILL_ENABLED = "aggregation_spill_enabled";
+    public static final String TOPN_SPILL_ENABLED = "topn_spill_enabled";
     public static final String DISTINCT_AGGREGATION_SPILL_ENABLED = "distinct_aggregation_spill_enabled";
     public static final String DEDUP_BASED_DISTINCT_AGGREGATION_SPILL_ENABLED = "dedup_based_distinct_aggregation_spill_enabled";
     public static final String DISTINCT_AGGREGATION_LARGE_BLOCK_SPILL_ENABLED = "distinct_aggregation_large_block_spill_enabled";
@@ -139,6 +140,7 @@ public final class SystemSessionProperties
     public static final String WINDOW_SPILL_ENABLED = "window_spill_enabled";
     public static final String ORDER_BY_SPILL_ENABLED = "order_by_spill_enabled";
     public static final String AGGREGATION_OPERATOR_UNSPILL_MEMORY_LIMIT = "aggregation_operator_unspill_memory_limit";
+    public static final String TOPN_OPERATOR_UNSPILL_MEMORY_LIMIT = "topn_operator_unspill_memory_limit";
     public static final String QUERY_MAX_REVOCABLE_MEMORY_PER_NODE = "query_max_revocable_memory_per_node";
     public static final String TEMP_STORAGE_SPILLER_BUFFER_SIZE = "temp_storage_spiller_buffer_size";
     public static final String OPTIMIZE_DISTINCT_AGGREGATIONS = "optimize_mixed_distinct_aggregations";
@@ -689,6 +691,11 @@ public SystemSessionProperties(
                         "Enable aggregate spilling if spill_enabled",
                         featuresConfig.isAggregationSpillEnabled(),
                         false),
+                booleanProperty(
+                        TOPN_SPILL_ENABLED,
+                        "Enable topN spilling if spill_enabled",
+                        featuresConfig.isTopNSpillEnabled(),
+                        false),
                 booleanProperty(
                         DISTINCT_AGGREGATION_SPILL_ENABLED,
                         "Enable spill for distinct aggregations if spill_enabled and aggregation_spill_enabled",
@@ -737,6 +744,15 @@ public SystemSessionProperties(
                         false,
                         value -> DataSize.valueOf((String) value),
                         DataSize::toString),
+                new PropertyMetadata<>(
+                        TOPN_OPERATOR_UNSPILL_MEMORY_LIMIT,
+                        "How much memory can should be allocated per topN operator in unspilling process",
+                        VARCHAR,
+                        DataSize.class,
+                        featuresConfig.getTopNOperatorUnspillMemoryLimit(),
+                        false,
+                        value -> DataSize.valueOf((String) value),
+                        DataSize::toString),
                 new PropertyMetadata<>(
                         QUERY_MAX_REVOCABLE_MEMORY_PER_NODE,
                         "Maximum amount of revocable memory a query can use",
@@ -1762,6 +1778,11 @@ public static boolean isAggregationSpillEnabled(Session session)
         return session.getSystemProperty(AGGREGATION_SPILL_ENABLED, Boolean.class) && isSpillEnabled(session);
     }
 
+    public static boolean isTopNSpillEnabled(Session session)
+    {
+        return session.getSystemProperty(TOPN_SPILL_ENABLED, Boolean.class) && isSpillEnabled(session);
+    }
+
     public static boolean isDistinctAggregationSpillEnabled(Session session)
     {
         return session.getSystemProperty(DISTINCT_AGGREGATION_SPILL_ENABLED, Boolean.class) && isAggregationSpillEnabled(session);
@@ -1804,6 +1825,13 @@ public static DataSize getAggregationOperatorUnspillMemoryLimit(Session session)
         return memoryLimitForMerge;
     }
 
+    public static DataSize getTopNOperatorUnspillMemoryLimit(Session session)
+    {
+        DataSize unspillMemoryLimit = session.getSystemProperty(TOPN_OPERATOR_UNSPILL_MEMORY_LIMIT, DataSize.class);
+        checkArgument(unspillMemoryLimit.toBytes() >= 0, "%s must be positive", TOPN_OPERATOR_UNSPILL_MEMORY_LIMIT);
+        return unspillMemoryLimit;
+    }
+
     public static DataSize getQueryMaxRevocableMemoryPerNode(Session session)
     {
         return session.getSystemProperty(QUERY_MAX_REVOCABLE_MEMORY_PER_NODE, DataSize.class);

presto-main/src/main/java/com/facebook/presto/operator/GroupByHash.java

@@ -16,10 +16,12 @@
 import com.facebook.presto.Session;
 import com.facebook.presto.common.Page;
 import com.facebook.presto.common.PageBuilder;
+import com.facebook.presto.common.array.IntBigArray;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.spi.function.aggregation.GroupByIdBlock;
 import com.facebook.presto.sql.gen.JoinCompiler;
 import com.google.common.annotations.VisibleForTesting;
+import it.unimi.dsi.fastutil.ints.IntIterator;
 
 import java.util.List;
 import java.util.Optional;
@@ -85,4 +87,34 @@ default boolean contains(int position, Page page, int[] hashChannels, long rawHa
 
     @VisibleForTesting
     int getCapacity();
+
+    default IntIterator getHashSortedGroupIds()
+    {
+        IntBigArray groupIds = new IntBigArray();
+        groupIds.ensureCapacity(getGroupCount());
+        for (int i = 0; i < getGroupCount(); i++) {
+            groupIds.set(i, i);
+        }
+
+        groupIds.sort(0, getGroupCount(), (leftGroupId, rightGroupId) ->
+                Long.compare(getRawHash(leftGroupId), getRawHash(rightGroupId)));
+
+        return new IntIterator()
+        {
+            private final int totalPositions = getGroupCount();
+            private int position;
+
+            @Override
+            public boolean hasNext()
+            {
+                return position < totalPositions;
+            }
+
+            @Override
+            public int nextInt()
+            {
+                return groupIds.get(position++);
+            }
+        };
+    }
 }