From f3910fa5a4976065f6432e133ab91f5f78049278 Mon Sep 17 00:00:00 2001
From: David Navas <davidn@clearstorydata.com>
Date: Mon, 15 Aug 2016 12:10:33 -0700
Subject: [PATCH] Clear just the bits that are necessary

---
 .../apache/spark/util/collection/BitSet.scala | 14 ++++++++
 .../spark/util/collection/BitSetSuite.scala   | 32 +++++++++++++++++++
 .../execution/joins/SortMergeOuterJoin.scala  |  4 +--
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/util/collection/BitSet.scala b/core/src/main/scala/org/apache/spark/util/collection/BitSet.scala
index 7ab67fc3a2de9..dd37bb1b715cd 100644
--- a/core/src/main/scala/org/apache/spark/util/collection/BitSet.scala
+++ b/core/src/main/scala/org/apache/spark/util/collection/BitSet.scala
@@ -57,6 +57,20 @@ class BitSet(numBits: Int) extends Serializable {
     }
   }
 
+  /**
+   * Clear all the bits up to a given index
+   */
+  def clearUntil(bitIndex: Int) {
+    val wordIndex = bitIndex >> 6 // divide by 64
+    var i = 0
+    while(i < wordIndex) { words(i) = 0; i += 1 }
+    if(wordIndex < words.length) {
+      // Set the remaining bits (note that the mask could still be zero)
+      val mask = ~(-1L << (bitIndex & 0x3f))
+      words(wordIndex) &= ~mask
+    }
+  }
+
   /**
    * Compute the bit-wise AND of the two sets returning the
    * result.
diff --git a/core/src/test/scala/org/apache/spark/util/collection/BitSetSuite.scala b/core/src/test/scala/org/apache/spark/util/collection/BitSetSuite.scala
index 69dbfa9cd7141..0169c9926e68f 100644
--- a/core/src/test/scala/org/apache/spark/util/collection/BitSetSuite.scala
+++ b/core/src/test/scala/org/apache/spark/util/collection/BitSetSuite.scala
@@ -152,4 +152,36 @@ class BitSetSuite extends SparkFunSuite {
     assert(bitsetDiff.nextSetBit(85) === 85)
     assert(bitsetDiff.nextSetBit(86) === -1)
   }
+
+  test( "[gs]etUntil" ) {
+    val bitSet = new BitSet(100)
+
+    bitSet.setUntil(bitSet.capacity)
+
+    (0 until bitSet.capacity).foreach { i =>
+      assert(bitSet.get(i))
+    }
+
+    bitSet.clearUntil(bitSet.capacity)
+
+    (0 until bitSet.capacity).foreach { i =>
+      assert(!bitSet.get(i))
+    }
+
+    val setUntil = bitSet.capacity / 2
+    bitSet.setUntil(setUntil)
+
+    val clearUntil = setUntil / 2
+    bitSet.clearUntil(clearUntil)
+
+    (0 until clearUntil).foreach { i =>
+      assert(!bitSet.get(i))
+    }
+    (clearUntil until setUntil).foreach { i =>
+      assert(bitSet.get(i))
+    }
+    (setUntil until bitSet.capacity).foreach { i =>
+      assert(!bitSet.get(i))
+    }
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala
index efaa69c1d3227..5969e6af300f0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala
@@ -371,12 +371,12 @@ private class SortMergeFullOuterJoinScanner(
     }
 
     if (leftMatches.size <= leftMatched.capacity) {
-      leftMatched.clear()
+      leftMatched.clearUntil(leftMatches.size)
     } else {
       leftMatched = new BitSet(leftMatches.size)
     }
     if (rightMatches.size <= rightMatched.capacity) {
-      rightMatched.clear()
+      rightMatched.clearUntil(rightMatches.size)
     } else {
       rightMatched = new BitSet(rightMatches.size)
     }