[SPARK-11084] [ML] [PYTHON] Check if index can contain non-zero value before binary search

At this moment `SparseVector.__getitem__` executes `np.searchsorted` first and checks if result is in an expected range after that. It is possible to check if index can contain non-zero value before executing `np.searchsorted`.

Author: zero323 <[email protected]>

Closes #9098 from zero323/sparse_vector_getitem_improved.

Author: zero323

python/pyspark/mllib/linalg/__init__.py

@@ -770,10 +770,10 @@ def __getitem__(self, index):
         if index < 0:
             index += self.size
 
-        insert_index = np.searchsorted(inds, index)
-        if insert_index >= inds.size:
+        if (inds.size == 0) or (index > inds.item(-1)):
             return 0.
 
+        insert_index = np.searchsorted(inds, index)
         row_ind = inds[insert_index]
         if row_ind == index:
             return vals[insert_index]

python/pyspark/mllib/tests.py

@@ -252,6 +252,16 @@ def test_sparse_vector_indexing(self):
         for ind in [7.8, '1']:
             self.assertRaises(TypeError, sv.__getitem__, ind)
 
+        zeros = SparseVector(4, {})
+        self.assertEqual(zeros[0], 0.0)
+        self.assertEqual(zeros[3], 0.0)
+        for ind in [4, -5]:
+            self.assertRaises(ValueError, zeros.__getitem__, ind)
+
+        empty = SparseVector(0, {})
+        for ind in [-1, 0, 1]:
+            self.assertRaises(ValueError, empty.__getitem__, ind)
+
     def test_matrix_indexing(self):
         mat = DenseMatrix(3, 2, [0, 1, 4, 6, 8, 10])
         expected = [[0, 6], [1, 8], [4, 10]]