Add tests and throw exception when we can't parse input.

python/pyspark/tests.py

@@ -2303,6 +2303,15 @@ def test_py4j_exception_message(self):
 
         self.assertTrue('NullPointerException' in _exception_message(context.exception))
 
+    def test_parsing_version_string(self):
+        from pyspark.util import VersionUtils
+
+        (major, minor) = VersionUtils.majorMinorVersion("2.4.0")
+        self.assertEqual(major, 2)
+        self.assertEqual(minor, 4)
+
+        self.assertRaises(ValueError, lambda: VersionUtils.majorMinorVersion("abced"))
+
 
 @unittest.skipIf(not _have_scipy, "SciPy not installed")
 class SciPyTests(PySparkTestCase):

python/pyspark/util.py

@@ -62,24 +62,28 @@ def _get_argspec(f):
     return argspec
 
 
-def majorMinorVersion(version):
+class VersionUtils(object):
     """
-    Get major and minor version numbers for given Spark version string.
-
-    >>> version = "2.4.0"
-    >>> majorMinorVersion(version)
-    (2, 4)
+    Provides utility method to determine Spark versions with given input string.
+    """
+    @staticmethod
+    def majorMinorVersion(version):
+        """
+        Get major and minor version numbers for given Spark version string.
 
-    >>> version = "abc"
-    >>> majorMinorVersion(version) is None
-    True
+        >>> version = "2.4.0"
+        >>> majorMinorVersion(version)
+        (2, 4)
+        >>> version = "2.3.0-SNAPSHOT"
+        >>> majorMinorVersion(version)
+        (2, 3)
 
-    """
-    m = re.search('^(\d+)\.(\d+)(\..*)?$', version)
-    if m is None:
-        return None
-    else:
-        return (int(m.group(1)), int(m.group(2)))
+        """
+        m = re.search('^(\d+)\.(\d+)(\..*)?$', version)
+        if m is None:
+            raise ValueError("invalid version string: " + version)
+        else:
+            return (int(m.group(1)), int(m.group(2)))
 
 
 if __name__ == "__main__":