Support spilling for Presto on Spark

presto-main/src/main/java/com/facebook/presto/memory/QueryContext.java

@@ -393,6 +393,11 @@ public synchronized long getPeakNodeTotalMemory()
         return peakNodeTotalMemory;
     }
 
+    public synchronized void setPeakNodeTotalMemory(long peakNodeTotalMemoryInBytes)
+    {
+        this.peakNodeTotalMemory = peakNodeTotalMemoryInBytes;
+    }
+
     private static class QueryMemoryReservationHandler
             implements MemoryReservationHandler
     {
@@ -473,7 +478,7 @@ private void enforceRevocableMemoryLimit(long allocated, long delta, long maxMem
     }
 
     @GuardedBy("this")
-    private String getAdditionalFailureInfo(long allocated, long delta)
+    public String getAdditionalFailureInfo(long allocated, long delta)
     {
         Map<String, Long> queryAllocations = memoryPool.getTaggedMemoryAllocations(queryId);
 

presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkConfig.java

@@ -17,6 +17,8 @@
 import com.facebook.airlift.configuration.ConfigDescription;
 import io.airlift.units.DataSize;
 
+import javax.validation.constraints.DecimalMax;
+import javax.validation.constraints.DecimalMin;
 import javax.validation.constraints.NotNull;
 
 import static io.airlift.units.DataSize.Unit.GIGABYTE;
@@ -37,6 +39,7 @@ public class PrestoSparkConfig
     private DataSize sparkBroadcastJoinMaxMemoryOverride;
     private boolean smileSerializationEnabled = true;
     private int splitAssignmentBatchSize = 1_000_000;
+    private double memoryRevokingThreshold;
 
     public boolean isSparkPartitionCountAutoTuneEnabled()
     {
@@ -191,4 +194,19 @@ public PrestoSparkConfig setSplitAssignmentBatchSize(int splitAssignmentBatchSiz
         this.splitAssignmentBatchSize = splitAssignmentBatchSize;
         return this;
     }
+
+    @DecimalMin("0.0")
+    @DecimalMax("1.0")
+    public double getMemoryRevokingThreshold()
+    {
+        return memoryRevokingThreshold;
+    }
+
+    @Config("spark.memory-revoking-threshold")
+    @ConfigDescription("Revoke memory when memory pool is filled over threshold")
+    public PrestoSparkConfig setMemoryRevokingThreshold(double memoryRevokingThreshold)
+    {
+        this.memoryRevokingThreshold = memoryRevokingThreshold;
+        return this;
+    }
 }

presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java

@@ -116,14 +116,14 @@
 import com.facebook.presto.spi.relation.DomainTranslator;
 import com.facebook.presto.spi.relation.PredicateCompiler;
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
-import com.facebook.presto.spiller.FileSingleStreamSpillerFactory;
 import com.facebook.presto.spiller.GenericPartitioningSpillerFactory;
 import com.facebook.presto.spiller.GenericSpillerFactory;
 import com.facebook.presto.spiller.NodeSpillConfig;
 import com.facebook.presto.spiller.PartitioningSpillerFactory;
 import com.facebook.presto.spiller.SingleStreamSpillerFactory;
 import com.facebook.presto.spiller.SpillerFactory;
 import com.facebook.presto.spiller.SpillerStats;
+import com.facebook.presto.spiller.TempStorageSingleStreamSpillerFactory;
 import com.facebook.presto.split.PageSinkManager;
 import com.facebook.presto.split.PageSinkProvider;
 import com.facebook.presto.split.PageSourceManager;
@@ -379,7 +379,7 @@ protected void setup(Binder binder)
 
         // spill
         binder.bind(SpillerFactory.class).to(GenericSpillerFactory.class).in(Scopes.SINGLETON);
-        binder.bind(SingleStreamSpillerFactory.class).to(FileSingleStreamSpillerFactory.class).in(Scopes.SINGLETON);
+        binder.bind(SingleStreamSpillerFactory.class).to(TempStorageSingleStreamSpillerFactory.class).in(Scopes.SINGLETON);
         binder.bind(PartitioningSpillerFactory.class).to(GenericPartitioningSpillerFactory.class).in(Scopes.SINGLETON);
         binder.bind(SpillerStats.class).in(Scopes.SINGLETON);
 

presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkSessionProperties.java

@@ -24,6 +24,7 @@
 
 import static com.facebook.presto.spi.session.PropertyMetadata.booleanProperty;
 import static com.facebook.presto.spi.session.PropertyMetadata.dataSizeProperty;
+import static com.facebook.presto.spi.session.PropertyMetadata.doubleProperty;
 import static com.facebook.presto.spi.session.PropertyMetadata.integerProperty;
 
 public class PrestoSparkSessionProperties
@@ -38,6 +39,7 @@ public class PrestoSparkSessionProperties
     public static final String STORAGE_BASED_BROADCAST_JOIN_WRITE_BUFFER_SIZE = "storage_based_broadcast_join_write_buffer_size";
     public static final String SPARK_BROADCAST_JOIN_MAX_MEMORY_OVERRIDE = "spark_broadcast_join_max_memory_override";
     public static final String SPARK_SPLIT_ASSIGNMENT_BATCH_SIZE = "spark_split_assignment_batch_size";
+    public static final String SPARK_MEMORY_REVOKING_THRESHOLD = "spark_memory_revoking_threshold";
 
     private final List<PropertyMetadata<?>> sessionProperties;
 
@@ -94,6 +96,11 @@ public PrestoSparkSessionProperties(PrestoSparkConfig prestoSparkConfig)
                         SPARK_SPLIT_ASSIGNMENT_BATCH_SIZE,
                         "Number of splits are processed in a single iteration",
                         prestoSparkConfig.getSplitAssignmentBatchSize(),
+                        false),
+                doubleProperty(
+                        SPARK_MEMORY_REVOKING_THRESHOLD,
+                        "Revoke memory when memory pool is filled over threshold",
+                        prestoSparkConfig.getMemoryRevokingThreshold(),
                         false));
     }
 
@@ -151,4 +158,9 @@ public static int getSplitAssignmentBatchSize(Session session)
     {
         return session.getSystemProperty(SPARK_SPLIT_ASSIGNMENT_BATCH_SIZE, Integer.class);
     }
+
+    public static double getMemoryRevokingThreshold(Session session)
+    {
+        return session.getSystemProperty(SPARK_MEMORY_REVOKING_THRESHOLD, Double.class);
+    }
 }

presto-spark-base/src/main/java/com/facebook/presto/spark/execution/PrestoSparkTaskExecutorFactory.java

@@ -38,8 +38,10 @@
 import com.facebook.presto.memory.MemoryPool;
 import com.facebook.presto.memory.NodeMemoryConfig;
 import com.facebook.presto.memory.QueryContext;
+import com.facebook.presto.memory.VoidTraversingQueryContextVisitor;
 import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.metadata.SessionPropertyManager;
+import com.facebook.presto.operator.OperatorContext;
 import com.facebook.presto.operator.OutputFactory;
 import com.facebook.presto.operator.TaskContext;
 import com.facebook.presto.operator.TaskStats;
@@ -109,16 +111,20 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.zip.CRC32;
 
+import static com.facebook.presto.ExceededMemoryLimitException.exceededLocalTotalMemoryLimit;
 import static com.facebook.presto.SystemSessionProperties.getHashPartitionCount;
 import static com.facebook.presto.SystemSessionProperties.getQueryMaxBroadcastMemory;
 import static com.facebook.presto.SystemSessionProperties.getQueryMaxMemoryPerNode;
 import static com.facebook.presto.SystemSessionProperties.getQueryMaxTotalMemoryPerNode;
+import static com.facebook.presto.SystemSessionProperties.isSpillEnabled;
 import static com.facebook.presto.execution.TaskState.FAILED;
 import static com.facebook.presto.execution.TaskStatus.STARTING_VERSION;
 import static com.facebook.presto.execution.buffer.BufferState.FINISHED;
 import static com.facebook.presto.metadata.MetadataUpdates.DEFAULT_METADATA_UPDATES;
+import static com.facebook.presto.spark.PrestoSparkSessionProperties.getMemoryRevokingThreshold;
 import static com.facebook.presto.spark.PrestoSparkSessionProperties.getShuffleOutputTargetAverageRowSize;
 import static com.facebook.presto.spark.PrestoSparkSessionProperties.getSparkBroadcastJoinMaxMemoryOverride;
 import static com.facebook.presto.spark.PrestoSparkSessionProperties.getStorageBasedBroadcastJoinWriteBufferSize;
@@ -134,7 +140,9 @@
 import static com.google.common.base.Throwables.propagateIfPossible;
 import static com.google.common.collect.Iterables.getFirst;
 import static io.airlift.units.DataSize.Unit.BYTE;
+import static io.airlift.units.DataSize.succinctBytes;
 import static java.lang.Math.min;
+import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 import static java.util.UUID.randomUUID;
 
@@ -176,6 +184,8 @@ public class PrestoSparkTaskExecutorFactory
     private final PrestoSparkBroadcastTableCacheManager prestoSparkBroadcastTableCacheManager;
     private final String storageBasedBroadcastJoinStorage;
 
+    private AtomicBoolean memoryRevokePending = new AtomicBoolean();
+
     @Inject
     public PrestoSparkTaskExecutorFactory(
             SessionPropertyManager sessionPropertyManager,
@@ -362,6 +372,7 @@ public <T extends PrestoSparkTaskOutput> IPrestoSparkTaskExecutor<T> doCreate(
         }
 
         MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("spark-executor-memory-pool"), maxTotalMemory);
+
         SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(maxSpillMemory);
 
         QueryContext queryContext = new QueryContext(
@@ -387,6 +398,42 @@ public <T extends PrestoSparkTaskOutput> IPrestoSparkTaskExecutor<T> doCreate(
                 allocationTrackingEnabled,
                 false);
 
+        final double memoryRevokingThreshold = getMemoryRevokingThreshold(session);
+        if (isSpillEnabled(session)) {
+            memoryPool.addListener((pool, queryId, totalMemoryReservationBytes) -> {
+                if (totalMemoryReservationBytes > queryContext.getPeakNodeTotalMemory()) {
+                    queryContext.setPeakNodeTotalMemory(totalMemoryReservationBytes);
+                }
+                if (totalMemoryReservationBytes > maxTotalMemory.toBytes()) {
+                    throw exceededLocalTotalMemoryLimit(
+                            maxTotalMemory,
+                            queryContext.getAdditionalFailureInfo(totalMemoryReservationBytes, 0) +
+                                    format("Total reserved memory: %s, Total revocable memory: %s",
+                                            succinctBytes(pool.getQueryMemoryReservation(queryId)),
+                                            succinctBytes(pool.getQueryRevocableMemoryReservation(queryId))));
+                }
+                if (totalMemoryReservationBytes > pool.getMaxBytes() * memoryRevokingThreshold && memoryRevokePending.compareAndSet(false, true)) {
+                    memoryUpdateExecutor.execute(() -> {
+                        try {
+                            taskContext.accept(new VoidTraversingQueryContextVisitor<Void>()
+                            {
+                                @Override
+                                public Void visitOperatorContext(OperatorContext operatorContext, Void nothing)
+                                {
+                                    operatorContext.requestMemoryRevoking();
+                                    return null;
+                                }
+                            }, null);
+                            memoryRevokePending.set(false);
+                        }
+                        catch (Exception e) {
+                            log.error(e, "Error requesting memory revoking");
+                        }
+                    });
+                }
+            });
+        }
+
         ImmutableMap.Builder<PlanNodeId, List<PrestoSparkShuffleInput>> shuffleInputs = ImmutableMap.builder();
         ImmutableMap.Builder<PlanNodeId, List<java.util.Iterator<PrestoSparkSerializedPage>>> pageInputs = ImmutableMap.builder();
         ImmutableMap.Builder<PlanNodeId, List<?>> broadcastInputs = ImmutableMap.builder();
@@ -547,7 +594,7 @@ private List<TaskSource> getTaskSources(Iterator<SerializedPrestoSparkTaskSource
             totalSerializedSizeInBytes += serializedTaskSource.getBytes().length;
             result.add(deserializeZstdCompressed(taskSourceCodec, serializedTaskSource.getBytes()));
         }
-        log.info("Total serialized size of all task sources: %s", DataSize.succinctBytes(totalSerializedSizeInBytes));
+        log.info("Total serialized size of all task sources: %s", succinctBytes(totalSerializedSizeInBytes));
         return result.build();
     }
 

presto-spark-base/src/test/java/com/facebook/presto/spark/TestPrestoSparkAbstractTestAggregations.java → presto-spark-base/src/test/java/com/facebook/presto/spark/TestPrestoSparkAggregations.java

@@ -16,13 +16,14 @@
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestAggregations;
 
-public class TestPrestoSparkAbstractTestAggregations
+import static com.facebook.presto.spark.PrestoSparkQueryRunner.createHivePrestoSparkQueryRunner;
+
+public class TestPrestoSparkAggregations
         extends AbstractTestAggregations
 {
     @Override
     protected QueryRunner createQueryRunner()
-            throws Exception
     {
-        return PrestoSparkQueryRunner.createHivePrestoSparkQueryRunner();
+        return createHivePrestoSparkQueryRunner();
     }
 }

presto-spark-base/src/test/java/com/facebook/presto/spark/TestPrestoSparkConfig.java

@@ -43,7 +43,8 @@ public void testDefaults()
                 .setStorageBasedBroadcastJoinWriteBufferSize(new DataSize(24, MEGABYTE))
                 .setSparkBroadcastJoinMaxMemoryOverride(null)
                 .setSmileSerializationEnabled(true)
-                .setSplitAssignmentBatchSize(1_000_000));
+                .setSplitAssignmentBatchSize(1_000_000)
+                .setMemoryRevokingThreshold(0));
     }
 
     @Test
@@ -62,6 +63,7 @@ public void testExplicitPropertyMappings()
                 .put("spark.broadcast-join-max-memory-override", "1GB")
                 .put("spark.smile-serialization-enabled", "false")
                 .put("spark.split-assignment-batch-size", "420")
+                .put("spark.memory-revoking-threshold", "0.5")
                 .build();
         PrestoSparkConfig expected = new PrestoSparkConfig()
                 .setSparkPartitionCountAutoTuneEnabled(false)
@@ -75,7 +77,8 @@ public void testExplicitPropertyMappings()
                 .setStorageBasedBroadcastJoinWriteBufferSize(new DataSize(4, MEGABYTE))
                 .setSparkBroadcastJoinMaxMemoryOverride(new DataSize(1, GIGABYTE))
                 .setSmileSerializationEnabled(false)
-                .setSplitAssignmentBatchSize(420);
+                .setSplitAssignmentBatchSize(420)
+                .setMemoryRevokingThreshold(0.5);
         assertFullMapping(properties, expected);
     }
 }

presto-spark-base/src/test/java/com/facebook/presto/spark/TestPrestoSparkSpilledAggregations.java

@@ -0,0 +1,38 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.spark;
+
+import com.facebook.presto.testing.QueryRunner;
+import com.google.common.collect.ImmutableMap;
+
+import java.nio.file.Paths;
+
+import static com.facebook.presto.spark.PrestoSparkQueryRunner.createHivePrestoSparkQueryRunner;
+import static io.airlift.tpch.TpchTable.getTables;
+
+public class TestPrestoSparkSpilledAggregations
+        extends TestPrestoSparkAggregations
+{
+    @Override
+    protected QueryRunner createQueryRunner()
+    {
+        ImmutableMap.Builder<String, String> configProperties = ImmutableMap.builder();
+        configProperties.put("experimental.spill-enabled", "true");
+        configProperties.put("experimental.join-spill-enabled", "true");
+        configProperties.put("experimental.temp-storage-buffer-size", "1MB");
+        configProperties.put("spark.memory-revoking-threshold", "0.0");
+        configProperties.put("experimental.spiller-spill-path", Paths.get(System.getProperty("java.io.tmpdir"), "presto", "spills").toString());
+        return createHivePrestoSparkQueryRunner(getTables(), configProperties.build());
+    }
+}

presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java

@@ -160,7 +160,7 @@ public void testSelectLargeInterval()
         assertEquals(result.getMaterializedRows().get(0).getField(0), new SqlIntervalYearMonth(Short.MAX_VALUE, 0));
     }
 
-    @Test
+    @Test(enabled = false)
     public void testEmptyJoins()
     {
         Session sessionWithEmptyJoin = Session.builder(getSession())