Adding ability to access local BitSet and to safely get a value at a given position

1 files changed, 9 insertions, 0 deletions

diff --git a/core/src/main/scala/org/apache/spark/util/hash/OpenHashSet.scala b/core/src/main/scala/org/apache/spark/util/hash/OpenHashSet.scala
index 7aa3f6220c..d083ab26ac 100644
--- a/core/src/main/scala/org/apache/spark/util/hash/OpenHashSet.scala
+++ b/core/src/main/scala/org/apache/spark/util/hash/OpenHashSet.scala
@@ -81,6 +81,8 @@ class OpenHashSet[@specialized(Long, Int) T: ClassManifest](
   protected var _data = classManifest[T].newArray(_capacity)
   protected var _bitset = new BitSet(_capacity)
 
+  def getBitSet = _bitset
+
   /** Number of elements in the set. */
   def size: Int = _size
 
@@ -147,6 +149,13 @@ class OpenHashSet[@specialized(Long, Int) T: ClassManifest](
   /** Return the value at the specified position. */
   def getValue(pos: Int): T = _data(pos)
 
+  /** Return the value at the specified position. */
+  def getValueSafe(pos: Int): T = {
+    assert(_bitset.get(pos))
+    _data(pos)
+  }
+
+
   /**
    * Return the next position with an element stored, starting from the given position inclusively.
    */