[SPARK-2909] [MLlib] [PySpark] SparseVector in pyspark now supports indexing

Slightly different than the scala code which converts the sparsevector into a densevector and then checks the index.

I also hope I've added tests in the right place.

Author: MechCoder <manojkumarsivaraj334@gmail.com>

Closes #4025 from MechCoder/spark-2909 and squashes the following commits:

07d0f26 [MechCoder] STY: Rename item to index
f02148b [MechCoder] [SPARK-2909] [Mlib] SparseVector in pyspark now supports indexing

1 files changed, 17 insertions, 0 deletions

diff --git a/python/pyspark/mllib/linalg.py b/python/pyspark/mllib/linalg.py
index 4f8491f43e..7f21190ed8 100644
--- a/python/pyspark/mllib/linalg.py
+++ b/python/pyspark/mllib/linalg.py
@@ -510,6 +510,23 @@ class SparseVector(Vector):
                 and np.array_equal(other.indices, self.indices)
                 and np.array_equal(other.values, self.values))
 
+    def __getitem__(self, index):
+        inds = self.indices
+        vals = self.values
+        if not isinstance(index, int):
+            raise ValueError(
+                "Indices must be of type integer, got type %s" % type(index))
+        if index < 0:
+            index += self.size
+        if index >= self.size or index < 0:
+            raise ValueError("Index %d out of bounds." % index)
+
+        insert_index = np.searchsorted(inds, index)
+        row_ind = inds[insert_index]
+        if row_ind == index:
+            return vals[insert_index]
+        return 0.
+
     def __ne__(self, other):
         return not self.__eq__(other)