Merge branch 'SPARK-32512-new' of https://github.com/stczwd/spark into SPARK-32512-new

Author: Jacky Lee

.github/workflows/build_and_test.yml

@@ -111,13 +111,13 @@ jobs:
         key: ${{ matrix.java }}-${{ matrix.hadoop }}-maven-${{ hashFiles('**/pom.xml') }}
         restore-keys: |
           ${{ matrix.java }}-${{ matrix.hadoop }}-maven-
-    - name: Cache Ivy local repository
+    - name: Cache Coursier local repository
       uses: actions/cache@v2
       with:
-        path: ~/.ivy2/cache
-        key: ${{ matrix.java }}-${{ matrix.hadoop }}-ivy-${{ hashFiles('**/pom.xml', '**/plugins.sbt') }}
+        path: ~/.cache/coursier
+        key: ${{ matrix.java }}-${{ matrix.hadoop }}-coursier-${{ hashFiles('**/pom.xml', '**/plugins.sbt') }}
         restore-keys: |
-          ${{ matrix.java }}-${{ matrix.hadoop }}-ivy-
+          ${{ matrix.java }}-${{ matrix.hadoop }}-coursier-
     - name: Install JDK ${{ matrix.java }}
       uses: actions/setup-java@v1
       with:
@@ -206,13 +206,13 @@ jobs:
         key: pyspark-maven-${{ hashFiles('**/pom.xml') }}
         restore-keys: |
           pyspark-maven-
-    - name: Cache Ivy local repository
+    - name: Cache Coursier local repository
       uses: actions/cache@v2
       with:
-        path: ~/.ivy2/cache
-        key: pyspark-ivy-${{ hashFiles('**/pom.xml', '**/plugins.sbt') }}
+        path: ~/.cache/coursier
+        key: pyspark-coursier-${{ hashFiles('**/pom.xml', '**/plugins.sbt') }}
         restore-keys: |
-          pyspark-ivy-
+          pyspark-coursier-
     - name: Install Python 3.6
       uses: actions/setup-python@v2
       with:
@@ -282,13 +282,13 @@ jobs:
         key: sparkr-maven-${{ hashFiles('**/pom.xml') }}
         restore-keys: |
           sparkr-maven-
-    - name: Cache Ivy local repository
+    - name: Cache Coursier local repository
       uses: actions/cache@v2
       with:
-        path: ~/.ivy2/cache
-        key: sparkr-ivy-${{ hashFiles('**/pom.xml', '**/plugins.sbt') }}
+        path: ~/.cache/coursier
+        key: sparkr-coursier-${{ hashFiles('**/pom.xml', '**/plugins.sbt') }}
         restore-keys: |
-          sparkr-ivy-
+          sparkr-coursier-
     - name: Run tests
       run: |
         mkdir -p ~/.m2
@@ -404,13 +404,13 @@ jobs:
     steps:
     - name: Checkout Spark repository
       uses: actions/checkout@v2
-    - name: Cache Ivy local repository
+    - name: Cache Coursier local repository
       uses: actions/cache@v2
       with:
-        path: ~/.ivy2/cache
-        key: scala-213-ivy-${{ hashFiles('**/pom.xml', '**/plugins.sbt') }}
+        path: ~/.cache/coursier
+        key: scala-213-coursier-${{ hashFiles('**/pom.xml', '**/plugins.sbt') }}
         restore-keys: |
-          scala-213-ivy-
+          scala-213-coursier-
     - name: Install Java 11
       uses: actions/setup-java@v1
       with:

R/pkg/inst/profile/shell.R

@@ -43,5 +43,7 @@
   cat("      /_/", "\n")
   cat("\n")
 
-  cat("\nSparkSession available as 'spark'.\n")
+  cat("\nSparkSession Web UI available at", SparkR::sparkR.uiWebUrl())
+  cat("\nSparkSession available as 'spark'(master = ", unlist(SparkR::sparkR.conf("spark.master")),
+    ", app id = ", unlist(SparkR::sparkR.conf("spark.app.id")), ").", "\n", sep = "")
 }

common/network-common/src/main/java/org/apache/spark/network/protocol/Encoders.java

@@ -18,6 +18,7 @@
 package org.apache.spark.network.protocol;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 
 import io.netty.buffer.ByteBuf;
@@ -46,7 +47,11 @@ public static String decode(ByteBuf buf) {
     }
   }
 
-  /** Bitmaps are encoded with their serialization length followed by the serialization bytes. */
+  /**
+   * Bitmaps are encoded with their serialization length followed by the serialization bytes.
+   *
+   * @since 3.1.0
+   */
   public static class Bitmaps {
     public static int encodedLength(RoaringBitmap b) {
       // Compress the bitmap before serializing it. Note that since BlockTransferMessage
@@ -57,13 +62,20 @@ public static int encodedLength(RoaringBitmap b) {
       return b.serializedSizeInBytes();
     }
 
+    /**
+     * The input ByteBuf for this encoder should have enough write capacity to fit the serialized
+     * bitmap. Other encoders which use {@link io.netty.buffer.AbstractByteBuf#writeBytes(byte[])}
+     * to write can expand the buf as writeBytes calls {@link ByteBuf#ensureWritable} internally.
+     * However, this encoder doesn't rely on netty's writeBytes and will fail if the input buf
+     * doesn't have enough write capacity.
+     */
     public static void encode(ByteBuf buf, RoaringBitmap b) {
-      int encodedLength = b.serializedSizeInBytes();
       // RoaringBitmap requires nio ByteBuffer for serde. We expose the netty ByteBuf as a nio
       // ByteBuffer. Here, we need to explicitly manage the index so we can write into the
       // ByteBuffer, and the write is reflected in the underneath ByteBuf.
-      b.serialize(buf.nioBuffer(buf.writerIndex(), encodedLength));
-      buf.writerIndex(buf.writerIndex() + encodedLength);
+      ByteBuffer byteBuffer = buf.nioBuffer(buf.writerIndex(), buf.writableBytes());
+      b.serialize(byteBuffer);
+      buf.writerIndex(buf.writerIndex() + byteBuffer.position());
     }
 
     public static RoaringBitmap decode(ByteBuf buf) {
@@ -172,7 +184,11 @@ public static long[] decode(ByteBuf buf) {
     }
   }
 
-  /** Bitmap arrays are encoded with the number of bitmaps followed by per-Bitmap encoding. */
+  /**
+   * Bitmap arrays are encoded with the number of bitmaps followed by per-Bitmap encoding.
+   *
+   * @since 3.1.0
+   */
   public static class BitmapArrays {
     public static int encodedLength(RoaringBitmap[] bitmaps) {
       int totalLength = 4;

common/network-common/src/main/java/org/apache/spark/network/util/TransportConf.java

@@ -363,4 +363,39 @@ public boolean useOldFetchProtocol() {
     return conf.getBoolean("spark.shuffle.useOldFetchProtocol", false);
   }
 
+  /**
+   * Class name of the implementation of MergedShuffleFileManager that merges the blocks
+   * pushed to it when push-based shuffle is enabled. By default, push-based shuffle is disabled at
+   * a cluster level because this configuration is set to
+   * 'org.apache.spark.network.shuffle.ExternalBlockHandler$NoOpMergedShuffleFileManager'.
+   * To turn on push-based shuffle at a cluster level, set the configuration to
+   * 'org.apache.spark.network.shuffle.RemoteBlockPushResolver'.
+   */
+  public String mergedShuffleFileManagerImpl() {
+    return conf.get("spark.shuffle.server.mergedShuffleFileManagerImpl",
+      "org.apache.spark.network.shuffle.ExternalBlockHandler$NoOpMergedShuffleFileManager");
+  }
+
+  /**
+   * The minimum size of a chunk when dividing a merged shuffle file into multiple chunks during
+   * push-based shuffle.
+   * A merged shuffle file consists of multiple small shuffle blocks. Fetching the
+   * complete merged shuffle file in a single response increases the memory requirements for the
+   * clients. Instead of serving the entire merged file, the shuffle service serves the
+   * merged file in `chunks`. A `chunk` constitutes few shuffle blocks in entirety and this
+   * configuration controls how big a chunk can get. A corresponding index file for each merged
+   * shuffle file will be generated indicating chunk boundaries.
+   */
+  public int minChunkSizeInMergedShuffleFile() {
+    return Ints.checkedCast(JavaUtils.byteStringAsBytes(
+      conf.get("spark.shuffle.server.minChunkSizeInMergedShuffleFile", "2m")));
+  }
+
+  /**
+   * The size of cache in memory which is used in push-based shuffle for storing merged index files.
+   */
+  public long mergedIndexCacheSize() {
+    return JavaUtils.byteStringAsBytes(
+      conf.get("spark.shuffle.server.mergedIndexCacheSize", "100m"));
+  }
 }

common/network-common/src/test/java/org/apache/spark/network/protocol/EncodersSuite.java

@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.network.protocol;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import org.junit.Test;
+import org.roaringbitmap.RoaringBitmap;
+
+import static org.junit.Assert.*;
+
+/**
+ * Tests for {@link Encoders}.
+ */
+public class EncodersSuite {
+
+  @Test
+  public void testRoaringBitmapEncodeDecode() {
+    RoaringBitmap bitmap = new RoaringBitmap();
+    bitmap.add(1, 2, 3);
+    ByteBuf buf = Unpooled.buffer(Encoders.Bitmaps.encodedLength(bitmap));
+    Encoders.Bitmaps.encode(buf, bitmap);
+    RoaringBitmap decodedBitmap = Encoders.Bitmaps.decode(buf);
+    assertEquals(bitmap, decodedBitmap);
+  }
+
+  @Test (expected = java.nio.BufferOverflowException.class)
+  public void testRoaringBitmapEncodeShouldFailWhenBufferIsSmall() {
+    RoaringBitmap bitmap = new RoaringBitmap();
+    bitmap.add(1, 2, 3);
+    ByteBuf buf = Unpooled.buffer(4);
+    Encoders.Bitmaps.encode(buf, bitmap);
+  }
+
+  @Test
+  public void testBitmapArraysEncodeDecode() {
+    RoaringBitmap[] bitmaps = new RoaringBitmap[] {
+      new RoaringBitmap(),
+      new RoaringBitmap(),
+      new RoaringBitmap(), // empty
+      new RoaringBitmap(),
+      new RoaringBitmap()
+    };
+    bitmaps[0].add(1, 2, 3);
+    bitmaps[1].add(1, 2, 4);
+    bitmaps[3].add(7L, 9L);
+    bitmaps[4].add(1L, 100L);
+    ByteBuf buf = Unpooled.buffer(Encoders.BitmapArrays.encodedLength(bitmaps));
+    Encoders.BitmapArrays.encode(buf, bitmaps);
+    RoaringBitmap[] decodedBitmaps = Encoders.BitmapArrays.decode(buf);
+    assertArrayEquals(bitmaps, decodedBitmaps);
+  }
+}

common/network-shuffle/pom.xml

@@ -47,6 +47,11 @@
       <artifactId>metrics-core</artifactId>
     </dependency>
 
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-tags_${scala.binary.version}</artifactId>
+    </dependency>
+
     <!-- Provided dependencies -->
     <dependency>
       <groupId>org.slf4j</groupId>
@@ -70,11 +75,6 @@
       <type>test-jar</type>
       <scope>test</scope>
     </dependency>
-    <dependency>
-      <groupId>org.apache.spark</groupId>
-      <artifactId>spark-tags_${scala.binary.version}</artifactId>
-      <scope>test</scope>
-    </dependency>
 
     <!--
       This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ErrorHandler.java

@@ -21,14 +21,18 @@
 
 import com.google.common.base.Throwables;
 
+import org.apache.spark.annotation.Evolving;
+
 /**
  * Plugs into {@link RetryingBlockFetcher} to further control when an exception should be retried
  * and logged.
  * Note: {@link RetryingBlockFetcher} will delegate the exception to this handler only when
  * - remaining retries < max retries
  * - exception is an IOException
+ *
+ * @since 3.1.0
  */
-
+@Evolving
 public interface ErrorHandler {
 
   boolean shouldRetryError(Throwable t);
@@ -44,6 +48,8 @@ default boolean shouldLogError(Throwable t) {
 
   /**
    * The error handler for pushing shuffle blocks to remote shuffle services.
+   *
+   * @since 3.1.0
    */
   class BlockPushErrorHandler implements ErrorHandler {
     /**

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockHandler.java

@@ -68,7 +68,7 @@ public ExternalBlockHandler(TransportConf conf, File registeredExecutorFile)
     throws IOException {
     this(new OneForOneStreamManager(),
       new ExternalShuffleBlockResolver(conf, registeredExecutorFile),
-      new NoOpMergedShuffleFileManager());
+      new NoOpMergedShuffleFileManager(conf));
   }
 
   public ExternalBlockHandler(
@@ -89,7 +89,7 @@ public ExternalShuffleBlockResolver getBlockResolver() {
   public ExternalBlockHandler(
       OneForOneStreamManager streamManager,
       ExternalShuffleBlockResolver blockManager) {
-    this(streamManager, blockManager, new NoOpMergedShuffleFileManager());
+    this(streamManager, blockManager, new NoOpMergedShuffleFileManager(null));
   }
 
   /** Enables mocking out the StreamManager, BlockManager, and MergeManager. */
@@ -175,7 +175,7 @@ protected void handleMessage(
         RegisterExecutor msg = (RegisterExecutor) msgObj;
         checkAuth(client, msg.appId);
         blockManager.registerExecutor(msg.appId, msg.execId, msg.executorInfo);
-        mergeManager.registerExecutor(msg.appId, msg.executorInfo.localDirs);
+        mergeManager.registerExecutor(msg.appId, msg.executorInfo);
         callback.onSuccess(ByteBuffer.wrap(new byte[0]));
       } finally {
         responseDelayContext.stop();
@@ -232,6 +232,7 @@ public StreamManager getStreamManager() {
    */
   public void applicationRemoved(String appId, boolean cleanupLocalDirs) {
     blockManager.applicationRemoved(appId, cleanupLocalDirs);
+    mergeManager.applicationRemoved(appId, cleanupLocalDirs);
   }
 
   /**
@@ -430,8 +431,15 @@ public ManagedBuffer next() {
   /**
    * Dummy implementation of merged shuffle file manager. Suitable for when push-based shuffle
    * is not enabled.
+   *
+   * @since 3.1.0
    */
-  private static class NoOpMergedShuffleFileManager implements MergedShuffleFileManager {
+  public static class NoOpMergedShuffleFileManager implements MergedShuffleFileManager {
+
+    // This constructor is needed because we use this constructor to instantiate an implementation
+    // of MergedShuffleFileManager using reflection.
+    // See YarnShuffleService#newMergedShuffleFileManagerInstance.
+    public NoOpMergedShuffleFileManager(TransportConf transportConf) {}
 
     @Override
     public StreamCallbackWithID receiveBlockDataAsStream(PushBlockStream msg) {
@@ -444,18 +452,13 @@ public MergeStatuses finalizeShuffleMerge(FinalizeShuffleMerge msg) throws IOExc
     }
 
     @Override
-    public void registerApplication(String appId, String user) {
-      // No-op. Do nothing.
-    }
-
-    @Override
-    public void registerExecutor(String appId, String[] localDirs) {
+    public void registerExecutor(String appId, ExecutorShuffleInfo executorInfo) {
       // No-Op. Do nothing.
     }
 
     @Override
     public void applicationRemoved(String appId, boolean cleanupLocalDirs) {
-      throw new UnsupportedOperationException("Cannot handle shuffle block merge");
+      // No-Op. Do nothing.
     }
 
     @Override

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/MergedBlockMeta.java

@@ -34,6 +34,8 @@
  * 1. Number of chunks in a merged shuffle block.
  * 2. Bitmaps for each chunk in the merged block. A chunk bitmap contains all the mapIds that were
  *    merged to that merged block chunk.
+ *
+ * @since 3.1.0
  */
 public class MergedBlockMeta {
   private final int numChunks;