Support distribution of broadcast table using disk storage in Presto-on-Spark

presto-main/src/main/java/com/facebook/presto/spiller/LocalTempStorage.java

@@ -18,6 +18,7 @@
 import com.facebook.presto.common.io.DataSink;
 import com.facebook.presto.common.io.OutputStreamDataSink;
 import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.storage.StorageCapabilities;
 import com.facebook.presto.spi.storage.TempDataOperationContext;
 import com.facebook.presto.spi.storage.TempDataSink;
 import com.facebook.presto.spi.storage.TempStorage;
@@ -31,6 +32,8 @@
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.nio.file.DirectoryStream;
 import java.nio.file.FileStore;
 import java.nio.file.Files;
@@ -43,6 +46,7 @@
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.lang.String.format;
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.nio.file.Files.createDirectories;
 import static java.nio.file.Files.delete;
 import static java.nio.file.Files.getFileStore;
@@ -118,6 +122,31 @@ public void remove(TempDataOperationContext context, TempStorageHandle handle)
         Files.delete(((LocalTempStorageHandle) handle).getFilePath());
     }
 
+    @Override
+    public byte[] serializeHandle(TempStorageHandle storageHandle)
+    {
+        URI uri = ((LocalTempStorageHandle) storageHandle).getFilePath().toUri();
+        return uri.toString().getBytes(UTF_8);
+    }
+
+    @Override
+    public TempStorageHandle deserialize(byte[] serializedStorageHandle)
+    {
+        String uriString = new String(serializedStorageHandle, UTF_8);
+        try {
+            return new LocalTempStorageHandle(Paths.get(new URI(uriString)));
+        }
+        catch (URISyntaxException e) {
+            throw new IllegalArgumentException("Invalid URI: " + uriString, e);
+        }
+    }
+
+    @Override
+    public List<StorageCapabilities> getStorageCapabilities()
+    {
+        return ImmutableList.of();
+    }
+
     private static void cleanupOldSpillFiles(Path path)
     {
         try (DirectoryStream<Path> stream = newDirectoryStream(path, SPILL_FILE_GLOB)) {
@@ -178,6 +207,12 @@ public Path getFilePath()
         {
             return filePath;
         }
+
+        @Override
+        public String toString()
+        {
+            return filePath.toString();
+        }
     }
 
     private static class LocalTempDataSink

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

@@ -0,0 +1,30 @@
+/*
+ * 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.spark.classloader_interface.PrestoSparkTaskOutput;
+import org.apache.spark.SparkException;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.broadcast.Broadcast;
+
+import java.util.List;
+import java.util.concurrent.TimeoutException;
+
+public interface PrestoSparkBroadcastDependency<T extends PrestoSparkTaskOutput>
+{
+    Broadcast<List<T>> executeBroadcast(JavaSparkContext sparkContext)
+            throws SparkException, TimeoutException;
+
+    void destroy();
+}

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

@@ -21,6 +21,7 @@
 
 import static io.airlift.units.DataSize.Unit.GIGABYTE;
 import static io.airlift.units.DataSize.Unit.KILOBYTE;
+import static io.airlift.units.DataSize.Unit.MEGABYTE;
 
 public class PrestoSparkConfig
 {
@@ -30,6 +31,9 @@ public class PrestoSparkConfig
     private int initialSparkPartitionCount = 16;
     private DataSize maxSplitsDataSizePerSparkPartition = new DataSize(2, GIGABYTE);
     private DataSize shuffleOutputTargetAverageRowSize = new DataSize(1, KILOBYTE);
+    private boolean storageBasedBroadcastJoinEnabled;
+    private DataSize storageBasedBroadcastJoinWriteBufferSize = new DataSize(24, MEGABYTE);
+    private String storageBasedBroadcastJoinStorage = "local";
 
     public boolean isSparkPartitionCountAutoTuneEnabled()
     {
@@ -109,4 +113,43 @@ public PrestoSparkConfig setShuffleOutputTargetAverageRowSize(DataSize shuffleOu
         this.shuffleOutputTargetAverageRowSize = shuffleOutputTargetAverageRowSize;
         return this;
     }
+
+    public boolean isStorageBasedBroadcastJoinEnabled()
+    {
+        return storageBasedBroadcastJoinEnabled;
+    }
+
+    @Config("spark.storage-based-broadcast-join-enabled")
+    @ConfigDescription("Distribute broadcast hashtable to workers using storage")
+    public PrestoSparkConfig setStorageBasedBroadcastJoinEnabled(boolean storageBasedBroadcastJoinEnabled)
+    {
+        this.storageBasedBroadcastJoinEnabled = storageBasedBroadcastJoinEnabled;
+        return this;
+    }
+
+    public DataSize getStorageBasedBroadcastJoinWriteBufferSize()
+    {
+        return storageBasedBroadcastJoinWriteBufferSize;
+    }
+
+    @Config("spark.storage-based-broadcast-join-write-buffer-size")
+    @ConfigDescription("Maximum size in bytes to buffer before flushing pages to disk")
+    public PrestoSparkConfig setStorageBasedBroadcastJoinWriteBufferSize(DataSize storageBasedBroadcastJoinWriteBufferSize)
+    {
+        this.storageBasedBroadcastJoinWriteBufferSize = storageBasedBroadcastJoinWriteBufferSize;
+        return this;
+    }
+
+    public String getStorageBasedBroadcastJoinStorage()
+    {
+        return storageBasedBroadcastJoinStorage;
+    }
+
+    @Config("spark.storage-based-broadcast-join-storage")
+    @ConfigDescription("TempStorage to use for dumping broadcast table")
+    public PrestoSparkConfig setStorageBasedBroadcastJoinStorage(String storageBasedBroadcastJoinStorage)
+    {
+        this.storageBasedBroadcastJoinStorage = storageBasedBroadcastJoinStorage;
+        return this;
+    }
 }

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

@@ -0,0 +1,85 @@
+/*
+ * 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.spark.classloader_interface.PrestoSparkSerializedPage;
+import io.airlift.units.DataSize;
+import org.apache.spark.SparkException;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.broadcast.Broadcast;
+import scala.Tuple2;
+
+import java.util.List;
+import java.util.concurrent.TimeoutException;
+
+import static com.facebook.presto.ExceededMemoryLimitException.exceededLocalBroadcastMemoryLimit;
+import static com.facebook.presto.spark.util.PrestoSparkUtils.computeNextTimeout;
+import static io.airlift.units.DataSize.succinctBytes;
+import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.stream.Collectors.toList;
+
+public class PrestoSparkMemoryBasedBroadcastDependency
+        implements PrestoSparkBroadcastDependency<PrestoSparkSerializedPage>
+{
+    private final RddAndMore<PrestoSparkSerializedPage> broadcastDependency;
+    private final DataSize maxBroadcastSize;
+    private final long queryCompletionDeadline;
+    private Broadcast<List<PrestoSparkSerializedPage>> broadcastVariable;
+
+    public PrestoSparkMemoryBasedBroadcastDependency(RddAndMore<PrestoSparkSerializedPage> broadcastDependency, DataSize maxBroadcastSize, long queryCompletionDeadline)
+    {
+        this.broadcastDependency = requireNonNull(broadcastDependency, "broadcastDependency cannot be null");
+        this.maxBroadcastSize = requireNonNull(maxBroadcastSize, "maxBroadcastSize cannot be null");
+        this.queryCompletionDeadline = queryCompletionDeadline;
+    }
+
+    @Override
+    public Broadcast<List<PrestoSparkSerializedPage>> executeBroadcast(JavaSparkContext sparkContext)
+            throws SparkException, TimeoutException
+    {
+        List<PrestoSparkSerializedPage> broadcastValue = broadcastDependency.collectAndDestroyDependenciesWithTimeout(computeNextTimeout(queryCompletionDeadline), MILLISECONDS).stream()
+                .map(Tuple2::_2)
+                .collect(toList());
+
+        long compressedBroadcastSizeInBytes = broadcastValue.stream()
+                .mapToInt(page -> page.getBytes().length)
+                .sum();
+        long uncompressedBroadcastSizeInBytes = broadcastValue.stream()
+                .mapToInt(page -> page.getUncompressedSizeInBytes())
+                .sum();
+
+        long maxBroadcastSizeInBytes = maxBroadcastSize.toBytes();
+
+        if (compressedBroadcastSizeInBytes > maxBroadcastSizeInBytes) {
+            throw exceededLocalBroadcastMemoryLimit(maxBroadcastSize, format("Compressed broadcast size: %s", succinctBytes(compressedBroadcastSizeInBytes)));
+        }
+
+        if (uncompressedBroadcastSizeInBytes > maxBroadcastSizeInBytes) {
+            throw exceededLocalBroadcastMemoryLimit(maxBroadcastSize, format("Uncompressed broadcast size: %s", succinctBytes(uncompressedBroadcastSizeInBytes)));
+        }
+
+        broadcastVariable = sparkContext.broadcast(broadcastValue);
+        return broadcastVariable;
+    }
+
+    @Override
+    public void destroy()
+    {
+        if (broadcastVariable != null) {
+            broadcastVariable.destroy();
+        }
+    }
+}

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

@@ -96,6 +96,7 @@
 import com.facebook.presto.server.SessionPropertyDefaults;
 import com.facebook.presto.server.security.ServerSecurityModule;
 import com.facebook.presto.spark.classloader_interface.SparkProcessType;
+import com.facebook.presto.spark.execution.PrestoSparkBroadcastTableCacheManager;
 import com.facebook.presto.spark.execution.PrestoSparkExecutionExceptionFactory;
 import com.facebook.presto.spark.execution.PrestoSparkTaskExecutorFactory;
 import com.facebook.presto.spark.node.PrestoSparkInternalNodeManager;
@@ -423,6 +424,7 @@ protected void setup(Binder binder)
         binder.bind(PrestoSparkTaskExecutorFactory.class).in(Scopes.SINGLETON);
         binder.bind(PrestoSparkQueryExecutionFactory.class).in(Scopes.SINGLETON);
         binder.bind(PrestoSparkService.class).in(Scopes.SINGLETON);
+        binder.bind(PrestoSparkBroadcastTableCacheManager.class).in(Scopes.SINGLETON);
 
         // extra credentials and authenticator for Presto-on-Spark
         newSetBinder(binder, PrestoSparkCredentialsProvider.class);