title: PySpark skip exchange
body: <p>I have two dataframes with the same schema, and both partitioned on the 'id' column.</p>
  <p>Schema:</p>
  <pre><code>id: int
  key: string
  value: string
  </code></pre>
  <p>If I perform a join on just the 'id' column, no additional exchange (shuffle) will occur. However, if I join on both the 'id' and 'key' column, then an exchange happens again. Since all data for a particular 'id' is already in the same partition it could perform the join without an exchange, but it repartitions both dataframes. Is there a way to avoid the exchange for this join with the two keys?</p>
  <p>I have not found any satisfactory answer, so I am asking here.</p>
  <p><strong>Quick example:</strong></p>
  <pre><code>data1_records = [[1, 'abc', 5],
                  [2, 'def', 6],
                  [3, 'ghj', 7]]
  
  data1_df = spark.createDataFrame(data1_records, ['id', 'key', 'value1']).repartition('id').cache()
  
  data2_records = [[1, 'abc', 50],
                  [2, 'def', 60],
                  [3, 'ghj', 70]]
  
  data2_df = spark.createDataFrame(data2_records, ['id', 'key', 'value2']).repartition('id').cache()
  </code></pre>
  <p>Plan for joining on just 'id' column:</p>
  <pre><code>data1_df.join(data2_df, [&quot;id&quot;]).explain()
  
  == Physical Plan ==
  *(3) Project [id#2656L, key#2657, value1#2658L, key#2678, value2#2679L]
  +- *(3) SortMergeJoin [id#2656L], [id#2677L], Inner
     :- *(1) Sort [id#2656L ASC NULLS FIRST], false, 0
     :  +- *(1) Filter isnotnull(id#2656L)
     :     +- InMemoryTableScan [id#2656L, key#2657, value1#2658L], [isnotnull(id#2656L)]
     :           +- InMemoryRelation [id#2656L, key#2657, value1#2658L], StorageLevel(disk, memory, deserialized, 1 replicas)
     :                 +- Exchange hashpartitioning(id#2656L, 200)
     :                    +- Scan ExistingRDD[id#2656L,key#2657,value1#2658L]
     +- *(2) Sort [id#2677L ASC NULLS FIRST], false, 0
        +- *(2) Filter isnotnull(id#2677L)
           +- InMemoryTableScan [id#2677L, key#2678, value2#2679L], [isnotnull(id#2677L)]
                 +- InMemoryRelation [id#2677L, key#2678, value2#2679L], StorageLevel(disk, memory, deserialized, 1 replicas)
                       +- Exchange hashpartitioning(id#2677L, 200)
                          +- Scan ExistingRDD[id#2677L,key#2678,value2#2679L]
  </code></pre>
  <p>Plan for joining on 'id' and 'key' columns (<strong>has additional exchanges</strong>):</p>
  <pre><code>data1_df.join(data2_df, [&quot;id&quot;, &quot;key&quot;]).explain()
  
  == Physical Plan ==
  *(5) Project [id#2580L, key#2581, value1#2582L, value2#2603L]
  +- *(5) SortMergeJoin [id#2580L, key#2581], [id#2601L, key#2602], Inner
     :- *(2) Sort [id#2580L ASC NULLS FIRST, key#2581 ASC NULLS FIRST], false, 0
     :  +- Exchange hashpartitioning(id#2580L, key#2581, 200)
     :     +- *(1) Filter (isnotnull(id#2580L) &amp;&amp; isnotnull(key#2581))
     :        +- InMemoryTableScan [id#2580L, key#2581, value1#2582L], [isnotnull(id#2580L), isnotnull(key#2581)]
     :              +- InMemoryRelation [id#2580L, key#2581, value1#2582L], StorageLevel(disk, memory, deserialized, 1 replicas)
     :                    +- Exchange hashpartitioning(id#2580L, 200)
     :                       +- Scan ExistingRDD[id#2580L,key#2581,value1#2582L]
     +- *(4) Sort [id#2601L ASC NULLS FIRST, key#2602 ASC NULLS FIRST], false, 0
        +- Exchange hashpartitioning(id#2601L, key#2602, 200)
           +- *(3) Filter (isnotnull(id#2601L) &amp;&amp; isnotnull(key#2602))
              +- InMemoryTableScan [id#2601L, key#2602, value2#2603L], [isnotnull(id#2601L), isnotnull(key#2602)]
                    +- InMemoryRelation [id#2601L, key#2602, value2#2603L], StorageLevel(disk, memory, deserialized, 1 replicas)
                          +- Exchange hashpartitioning(id#2601L, 200)
                             +- Scan ExistingRDD[id#2601L,key#2602,value2#2603L]
  </code></pre>
  <p><strong>Edit 1</strong>
  This works for a groupBy operation.</p>
  <p>GroupBy on 'id' no additional exchange (expected)</p>
  <pre><code>data1_df.groupBy(&quot;id&quot;).agg(F.sum(&quot;value1&quot;)).explain()
  
  == Physical Plan ==
  *(1) HashAggregate(keys=[id#2932L], functions=[sum(value1#2934L)])
  +- *(1) HashAggregate(keys=[id#2932L], functions=[partial_sum(value1#2934L)])
     +- InMemoryTableScan [id#2932L, value1#2934L]
           +- InMemoryRelation [id#2932L, key#2933, value1#2934L], StorageLevel(disk, memory, deserialized, 1 replicas)
                 +- Exchange hashpartitioning(id#2932L, 200)
                    +- Scan ExistingRDD[id#2932L,key#2933,value1#2934L]
  </code></pre>
  <p>GroupBy on 'id' and 'key' no additional exchange (desired)</p>
  <pre><code>data1_df.groupBy(&quot;id&quot;, &quot;key&quot;).agg(F.sum(&quot;value1&quot;)).explain()
  
  == Physical Plan ==
  *(1) HashAggregate(keys=[id#2998L, key#2999], functions=[sum(value1#3000L)])
  +- *(1) HashAggregate(keys=[id#2998L, key#2999], functions=[partial_sum(value1#3000L)])
     +- InMemoryTableScan [id#2998L, key#2999, value1#3000L]
           +- InMemoryRelation [id#2998L, key#2999, value1#3000L], StorageLevel(disk, memory, deserialized, 1 replicas)
                 +- Exchange hashpartitioning(id#2998L, 200)
                    +- Scan ExistingRDD[id#2998L,key#2999,value1#3000L]
  </code></pre>
  <p>GroupBy on 'key' requires an additional exchange</p>
  <pre><code>data1_df.groupBy(&quot;key&quot;).agg(F.sum(&quot;value1&quot;)).explain()
  
  == Physical Plan ==
  *(2) HashAggregate(keys=[key#3066], functions=[sum(value1#3067L)])
  +- Exchange hashpartitioning(key#3066, 200)
     +- *(1) HashAggregate(keys=[key#3066], functions=[partial_sum(value1#3067L)])
        +- InMemoryTableScan [key#3066, value1#3067L]
              +- InMemoryRelation [id#3065L, key#3066, value1#3067L], StorageLevel(disk, memory, deserialized, 1 replicas)
                    +- Exchange hashpartitioning(id#3065L, 200)
                       +- Scan ExistingRDD[id#3065L,key#3066,value1#3067L]
  </code></pre>