Fix Spark version detection for automatic Deequ maven coordinate definition (#114)

* fix: revert get spark version behavior

* fix: better spark version extraction

Author: lecardozo

.github/workflows/base.yml

@@ -1,7 +1,7 @@
 on:
   push:
     branches:
-     - "*"
+      - "*"
 
 jobs:
   test:
@@ -23,7 +23,7 @@ jobs:
         name: Setup Java 11
         if: startsWith(matrix.PYSPARK_VERSION, '3')
         with:
-          java-version: '11'
+          java-version: "11"
 
       - name: Running tests with pyspark==${{matrix.PYSPARK_VERSION}}
         env:

pydeequ/configs.py

@@ -1,8 +1,9 @@
 # -*- coding: utf-8 -*-
 from functools import lru_cache
-import subprocess
+import os
 import re
 
+
 SPARK_TO_DEEQU_COORD_MAPPING = {
     "3.2": "com.amazon.deequ:deequ:2.0.1-spark-3.2",
     "3.1": "com.amazon.deequ:deequ:2.0.0-spark-3.1",
@@ -11,17 +12,21 @@
 }
 
 
+def _extract_major_minor_versions(full_version: str):
+    major_minor_pattern = re.compile(r"(\d+\.\d+)\.*")
+    match = re.match(major_minor_pattern, full_version)
+    if match:
+        return match.group(1)
+
+
 @lru_cache(maxsize=None)
 def _get_spark_version() -> str:
-    # Get version from a subprocess so we don't mess up with existing SparkContexts.
-    command = [
-        "python",
-        "-c",
-        "from pyspark import SparkContext; print(SparkContext.getOrCreate()._jsc.version())",
-    ]
-    output = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    spark_version = output.stdout.decode().split("\n")[-2]
-    return spark_version
+    try:
+        spark_version = os.environ["SPARK_VERSION"]
+    except KeyError:
+        raise RuntimeError(f"SPARK_VERSION environment variable is required. Supported values are: {SPARK_TO_DEEQU_COORD_MAPPING.keys()}")
+
+    return _extract_major_minor_versions(spark_version)
 
 
 def _get_deequ_maven_config():
@@ -30,7 +35,7 @@ def _get_deequ_maven_config():
         return SPARK_TO_DEEQU_COORD_MAPPING[spark_version[:3]]
     except KeyError:
         raise RuntimeError(
-            f"Found Incompatible Spark version {spark_version}; Use one of the Supported Spark versions for Deequ: {SPARK_TO_DEEQU_COORD_MAPPING.keys()}"
+            f"Found incompatible Spark version {spark_version}; Use one of the Supported Spark versions for Deequ: {SPARK_TO_DEEQU_COORD_MAPPING.keys()}"
         )
 
 

tests/test_config.py

@@ -0,0 +1,15 @@
+import pytest
+from pydeequ.configs import _extract_major_minor_versions
+
+
+@pytest.parametrize(
+    "full_version, major_minor_version",
+    [
+        ("3.2.1", "3.2"),
+        ("3.1", "3.1"),
+        ("3.10.3", "3.10"),
+        ("3.10", "3.10")
+    ]
+)
+def test_extract_major_minor_versions(full_version, major_minor_version):
+    assert _extract_major_minor_versions(full_version) == major_minor_version