ae skew join (apache#117)

* ae skew join

* add metircs querystage

* No support that a union node with differing number of partitions if we explicitly repartition them.

* fix style

* release r40

Author: 7mming7

assembly/pom.xml

@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.11</artifactId>
-    <version>2.4.1-kylin-r35</version>
+    <version>2.4.1-kylin-r40</version>
     <relativePath>../pom.xml</relativePath>
   </parent>
 

common/kvstore/pom.xml

@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.11</artifactId>
-    <version>2.4.1-kylin-r35</version>
+    <version>2.4.1-kylin-r40</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

common/network-common/pom.xml

@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.11</artifactId>
-    <version>2.4.1-kylin-r35</version>
+    <version>2.4.1-kylin-r40</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

common/network-shuffle/pom.xml

@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.11</artifactId>
-    <version>2.4.1-kylin-r35</version>
+    <version>2.4.1-kylin-r40</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

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

@@ -197,48 +197,60 @@ public Map<String, Metric> getMetrics() {
     }
   }
 
+  private boolean isShuffleBlock(String[] blockIdParts) {
+    // length == 4: ShuffleBlockId
+    // length == 5: ContinuousShuffleBlockId
+    return (blockIdParts.length == 4 || blockIdParts.length == 5) &&
+            blockIdParts[0].equals("shuffle");
+  }
+
   private class ManagedBufferIterator implements Iterator<ManagedBuffer> {
 
     private int index = 0;
     private final String appId;
     private final String execId;
     private final int shuffleId;
-    // An array containing mapId and reduceId pairs.
-    private final int[] mapIdAndReduceIds;
+    // An array containing mapId, reduceId and numBlocks tuple
+    private final int[] shuffleBlockIds;
 
     ManagedBufferIterator(String appId, String execId, String[] blockIds) {
       this.appId = appId;
       this.execId = execId;
       String[] blockId0Parts = blockIds[0].split("_");
-      if (blockId0Parts.length != 4 || !blockId0Parts[0].equals("shuffle")) {
+      if (!isShuffleBlock(blockId0Parts)) {
         throw new IllegalArgumentException("Unexpected shuffle block id format: " + blockIds[0]);
       }
       this.shuffleId = Integer.parseInt(blockId0Parts[1]);
-      mapIdAndReduceIds = new int[2 * blockIds.length];
+      shuffleBlockIds = new int[3 * blockIds.length];
       for (int i = 0; i < blockIds.length; i++) {
         String[] blockIdParts = blockIds[i].split("_");
-        if (blockIdParts.length != 4 || !blockIdParts[0].equals("shuffle")) {
+        if (!isShuffleBlock(blockIdParts)) {
           throw new IllegalArgumentException("Unexpected shuffle block id format: " + blockIds[i]);
         }
         if (Integer.parseInt(blockIdParts[1]) != shuffleId) {
           throw new IllegalArgumentException("Expected shuffleId=" + shuffleId +
-            ", got:" + blockIds[i]);
+                  ", got:" + blockIds[i]);
+        }
+        shuffleBlockIds[3 * i] = Integer.parseInt(blockIdParts[2]);
+        shuffleBlockIds[3 * i + 1] = Integer.parseInt(blockIdParts[3]);
+        if (blockIdParts.length == 4) {
+          shuffleBlockIds[3 * i + 2] = 1;
+        } else {
+          shuffleBlockIds[3 * i + 2] = Integer.parseInt(blockIdParts[4]);
         }
-        mapIdAndReduceIds[2 * i] = Integer.parseInt(blockIdParts[2]);
-        mapIdAndReduceIds[2 * i + 1] = Integer.parseInt(blockIdParts[3]);
       }
     }
 
     @Override
     public boolean hasNext() {
-      return index < mapIdAndReduceIds.length;
+      return index < shuffleBlockIds.length;
     }
 
     @Override
     public ManagedBuffer next() {
       final ManagedBuffer block = blockManager.getBlockData(appId, execId, shuffleId,
-        mapIdAndReduceIds[index], mapIdAndReduceIds[index + 1]);
-      index += 2;
+              shuffleBlockIds[index], shuffleBlockIds[index + 1], shuffleBlockIds[index + 2]);
+      index += 3;
       metrics.blockTransferRateBytes.mark(block != null ? block.size() : 0);
       return block;
     }

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

@@ -162,21 +162,22 @@ public void registerExecutor(
   }
 
   /**
-   * Obtains a FileSegmentManagedBuffer from (shuffleId, mapId, reduceId). We make assumptions
+   * Obtains a FileSegmentManagedBuffer from (shuffleId, mapId, reduceId, numBlocks). We make assumptions
    * about how the hash and sort based shuffles store their data.
    */
   public ManagedBuffer getBlockData(
-      String appId,
-      String execId,
-      int shuffleId,
-      int mapId,
-      int reduceId) {
+          String appId,
+          String execId,
+          int shuffleId,
+          int mapId,
+          int reduceId,
+          int numBlocks) {
     ExecutorShuffleInfo executor = executors.get(new AppExecId(appId, execId));
     if (executor == null) {
       throw new RuntimeException(
-        String.format("Executor is not registered (appId=%s, execId=%s)", appId, execId));
+              String.format("Executor is not registered (appId=%s, execId=%s)", appId, execId));
     }
-    return getSortBasedShuffleBlockData(executor, shuffleId, mapId, reduceId);
+    return getSortBasedShuffleBlockData(executor, shuffleId, mapId, reduceId, numBlocks);
   }
 
   /**
@@ -280,19 +281,19 @@ public boolean accept(File dir, String name) {
    * and the block id format is from ShuffleDataBlockId and ShuffleIndexBlockId.
    */
   private ManagedBuffer getSortBasedShuffleBlockData(
-    ExecutorShuffleInfo executor, int shuffleId, int mapId, int reduceId) {
+          ExecutorShuffleInfo executor, int shuffleId, int mapId, int reduceId, int numBlocks) {
     File indexFile = getFile(executor.localDirs, executor.subDirsPerLocalDir,
-      "shuffle_" + shuffleId + "_" + mapId + "_0.index");
+            "shuffle_" + shuffleId + "_" + mapId + "_0.index");
 
     try {
       ShuffleIndexInformation shuffleIndexInformation = shuffleIndexCache.get(indexFile);
-      ShuffleIndexRecord shuffleIndexRecord = shuffleIndexInformation.getIndex(reduceId);
+      ShuffleIndexRecord shuffleIndexRecord = shuffleIndexInformation.getIndex(reduceId, numBlocks);
       return new FileSegmentManagedBuffer(
-        conf,
-        getFile(executor.localDirs, executor.subDirsPerLocalDir,
-          "shuffle_" + shuffleId + "_" + mapId + "_0.data"),
-        shuffleIndexRecord.getOffset(),
-        shuffleIndexRecord.getLength());
+              conf,
+              getFile(executor.localDirs, executor.subDirsPerLocalDir,
+                      "shuffle_" + shuffleId + "_" + mapId + "_0.data"),
+              shuffleIndexRecord.getOffset(),
+              shuffleIndexRecord.getLength());
     } catch (ExecutionException e) {
       throw new RuntimeException("Failed to open file: " + indexFile, e);
     }

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

@@ -59,9 +59,9 @@ public int getSize() {
   /**
    * Get index offset for a particular reducer.
    */
-  public ShuffleIndexRecord getIndex(int reduceId) {
+  public ShuffleIndexRecord getIndex(int reduceId, int numBlocks) {
     long offset = offsets.get(reduceId);
-    long nextOffset = offsets.get(reduceId + 1);
+    long nextOffset = offsets.get(reduceId + numBlocks);
     return new ShuffleIndexRecord(offset, nextOffset - offset);
   }
 }

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandlerSuite.java

@@ -83,8 +83,8 @@ public void testOpenShuffleBlocks() {
 
     ManagedBuffer block0Marker = new NioManagedBuffer(ByteBuffer.wrap(new byte[3]));
     ManagedBuffer block1Marker = new NioManagedBuffer(ByteBuffer.wrap(new byte[7]));
-    when(blockResolver.getBlockData("app0", "exec1", 0, 0, 0)).thenReturn(block0Marker);
-    when(blockResolver.getBlockData("app0", "exec1", 0, 0, 1)).thenReturn(block1Marker);
+    when(blockResolver.getBlockData("app0", "exec1", 0, 0, 0, 1)).thenReturn(block0Marker);
+    when(blockResolver.getBlockData("app0", "exec1", 0, 0, 1, 1)).thenReturn(block1Marker);
     ByteBuffer openBlocks = new OpenBlocks("app0", "exec1",
       new String[] { "shuffle_0_0_0", "shuffle_0_0_1" })
       .toByteBuffer();
@@ -106,8 +106,8 @@ public void testOpenShuffleBlocks() {
     assertEquals(block0Marker, buffers.next());
     assertEquals(block1Marker, buffers.next());
     assertFalse(buffers.hasNext());
-    verify(blockResolver, times(1)).getBlockData("app0", "exec1", 0, 0, 0);
-    verify(blockResolver, times(1)).getBlockData("app0", "exec1", 0, 0, 1);
+    verify(blockResolver, times(1)).getBlockData("app0", "exec1", 0, 0, 0, 1);
+    verify(blockResolver, times(1)).getBlockData("app0", "exec1", 0, 0, 1, 1);
 
     // Verify open block request latency metrics
     Timer openBlockRequestLatencyMillis = (Timer) ((ExternalShuffleBlockHandler) handler)

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockResolverSuite.java

@@ -66,7 +66,7 @@ public void testBadRequests() throws IOException {
     ExternalShuffleBlockResolver resolver = new ExternalShuffleBlockResolver(conf, null);
     // Unregistered executor
     try {
-      resolver.getBlockData("app0", "exec1", 1, 1, 0);
+      resolver.getBlockData("app0", "exec1", 1, 1, 0, 1);
       fail("Should have failed");
     } catch (RuntimeException e) {
       assertTrue("Bad error message: " + e, e.getMessage().contains("not registered"));
@@ -75,7 +75,7 @@ public void testBadRequests() throws IOException {
     // Invalid shuffle manager
     try {
       resolver.registerExecutor("app0", "exec2", dataContext.createExecutorInfo("foobar"));
-      resolver.getBlockData("app0", "exec2", 1, 1, 0);
+      resolver.getBlockData("app0", "exec2", 1, 1, 0, 1);
       fail("Should have failed");
     } catch (UnsupportedOperationException e) {
       // pass
@@ -85,7 +85,7 @@ public void testBadRequests() throws IOException {
     resolver.registerExecutor("app0", "exec3",
       dataContext.createExecutorInfo(SORT_MANAGER));
     try {
-      resolver.getBlockData("app0", "exec3", 1, 1, 0);
+      resolver.getBlockData("app0", "exec3", 1, 1, 0, 1);
       fail("Should have failed");
     } catch (Exception e) {
       // pass
@@ -99,18 +99,25 @@ public void testSortShuffleBlocks() throws IOException {
       dataContext.createExecutorInfo(SORT_MANAGER));
 
     InputStream block0Stream =
-      resolver.getBlockData("app0", "exec0", 0, 0, 0).createInputStream();
+      resolver.getBlockData("app0", "exec0", 0, 0, 0, 1).createInputStream();
     String block0 = CharStreams.toString(
         new InputStreamReader(block0Stream, StandardCharsets.UTF_8));
     block0Stream.close();
     assertEquals(sortBlock0, block0);
 
     InputStream block1Stream =
-      resolver.getBlockData("app0", "exec0", 0, 0, 1).createInputStream();
+      resolver.getBlockData("app0", "exec0", 0, 0, 1, 1).createInputStream();
     String block1 = CharStreams.toString(
         new InputStreamReader(block1Stream, StandardCharsets.UTF_8));
     block1Stream.close();
     assertEquals(sortBlock1, block1);
+
+    InputStream block01Stream =
+            resolver.getBlockData("app0", "exec0", 0, 0, 0, 2).createInputStream();
+    String block01 = CharStreams.toString(
+            new InputStreamReader(block01Stream, StandardCharsets.UTF_8));
+    block01Stream.close();
+    assertEquals(sortBlock0 + sortBlock1, block01);
   }
 
   @Test

common/network-yarn/pom.xml

@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.11</artifactId>
-    <version>2.4.1-kylin-r35</version>
+    <version>2.4.1-kylin-r40</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>