Draft: Add Spark 4.1 support

.github/workflows/base.yml

@@ -25,6 +25,9 @@ jobs:
           - PYSPARK_VERSION: "3.5"
             PYTHON_VERSION: "3.9"
             JAVA_VERSION: "17"
+          - PYSPARK_VERSION: "4.1.2"
+            PYTHON_VERSION: "3.10"
+            JAVA_VERSION: "17"
 
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2

Dockerfile.spark-4.1

@@ -0,0 +1,29 @@
+FROM ubuntu:22.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update
+RUN apt-get install -y software-properties-common
+RUN add-apt-repository ppa:deadsnakes/ppa
+RUN apt-get install -y python3.10 python3-pip
+RUN apt-get install -y python3.10-distutils
+RUN apt-get install -y openjdk-17-jdk
+
+# Update symlink to point to latest
+RUN rm /usr/bin/python3 && ln -s /usr/bin/python3.10 /usr/bin/python3
+RUN python3 --version
+RUN pip3 --version
+RUN java -version
+RUN pip install poetry==1.7.1
+
+RUN mkdir python-deequ
+COPY pyproject.toml /python-deequ
+COPY poetry.lock /python-deequ
+WORKDIR python-deequ
+
+RUN poetry install -vvv
+RUN poetry add pyspark==4.1.2 -vvv
+
+ENV SPARK_VERSION=4.1
+COPY . /python-deequ
+CMD poetry run python -m pytest -s tests

README.md

@@ -29,6 +29,16 @@ You can install [PyDeequ via pip](https://pypi.org/project/pydeequ/).
 pip install pydeequ
 ```
 
+### Select your Spark version
+
+Set `SPARK_VERSION` to match your Spark runtime before importing PyDeequ. For example, use `3.5` or `4.1`:
+
+```python
+import os
+
+os.environ["SPARK_VERSION"] = "4.1"
+```
+
 ### Set up a PySpark session
 ```python
 from pyspark.sql import SparkSession, Row
@@ -261,4 +271,3 @@ See [CONTRIBUTING](CONTRIBUTING.md#security-issue-notifications) for more inform
 ## License
 
 This library is licensed under the Apache 2.0 License.
-

pydeequ/analyzers.py

@@ -311,7 +311,7 @@ def _analyzer_jvm(self):
             self.instance,
             self.predicate,
             self._jvm.scala.Option.apply(self.where),
-            self._jvm.scala.collection.Seq.empty(),
+            to_scala_seq(self._jvm, []),
             self._jvm.scala.Option.apply(None)
         )
 

pydeequ/checks.py

@@ -566,7 +566,7 @@ def satisfies(self, columnCondition, constraintName, assertion=None, hint=None):
             constraintName,
             assertion_func,
             hint,
-            self._jvm.scala.collection.Seq.empty(),
+            to_scala_seq(self._jvm, []),
             self._jvm.scala.Option.apply(None)
         )
         return self

pydeequ/configs.py

@@ -5,6 +5,7 @@
 
 
 SPARK_TO_DEEQU_COORD_MAPPING = {
+    "4.1": "com.amazon.deequ:deequ:2.0.18-spark-4.1",
     "3.5": "com.amazon.deequ:deequ:2.0.8-spark-3.5",
     "3.3": "com.amazon.deequ:deequ:2.0.8-spark-3.3",
     "3.2": "com.amazon.deequ:deequ:2.0.8-spark-3.2",

pydeequ/profiles.py

@@ -254,7 +254,7 @@ def _columnProfilesFromColumnRunBuilderRun(self, run):
         :return self: a setter for columnProfilerRunner result
         """
         self._run_result = run
-        profile_map = self._jvm.scala.collection.JavaConversions.mapAsJavaMap(run.profiles())  # TODO from ScalaUtils
+        profile_map = self._jvm.scala.collection.JavaConverters.mapAsJavaMap(run.profiles())  # TODO from ScalaUtils
         self._profiles = {column: self._columnProfileBuilder(column, profile_map[column]) for column in profile_map}
         self._numRecords = run.numRecords()
         return self

pydeequ/scala_utils.py

@@ -77,7 +77,8 @@ def to_scala_seq(jvm, iterable):
     Returns:
         Scala sequence
     """
-    return jvm.scala.collection.JavaConversions.iterableAsScalaIterable(iterable).toSeq()
+    # toSeq yields Stream on Scala 2.12 and List on 2.13; List is a compatible Seq on both.
+    return jvm.scala.collection.JavaConverters.iterableAsScalaIterable(iterable).toList()
 
 
 def to_scala_map(spark_session, d):
@@ -93,11 +94,11 @@ def to_scala_map(spark_session, d):
 
 
 def scala_map_to_dict(jvm, scala_map):
-    return dict(jvm.scala.collection.JavaConversions.mapAsJavaMap(scala_map))
+    return dict(jvm.scala.collection.JavaConverters.mapAsJavaMap(scala_map))
 
 
 def scala_map_to_java_map(jvm, scala_map):
-    return jvm.scala.collection.JavaConversions.mapAsJavaMap(scala_map)
+    return jvm.scala.collection.JavaConverters.mapAsJavaMap(scala_map)
 
 
 def java_list_to_python_list(java_list: str, datatype):

pyproject.toml

@@ -29,9 +29,15 @@ classifiers = [
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4"
-numpy = ">=1.14.1"
-pandas = ">=0.23.0"
-pyspark = { version = ">=2.4.7,<4.0.0", optional = true }
+numpy = [
+    { version = ">=1.14.1,<1.24", python = ">=3.9,<3.10" },
+    { version = ">=1.22", python = ">=3.10" },
+]
+pandas = [
+    { version = ">=0.23.0,<2", python = ">=3.9,<3.10" },
+    { version = ">=2.2", python = ">=3.10" },
+]
+pyspark = { version = ">=2.4.7,<4.2", optional = true }
 
 [tool.poetry.dev-dependencies]
 pytest = "^8.0"

tests/test_analyzers.py

@@ -486,8 +486,14 @@ def test_fail_MinLength(self):
         self.assertEqual(self.MinLength("a"), [])
 
     def test_MutualInformation(self):
-        self.assertEqual(self.MutualInformation(["b", "c"]), [Row(value=0.7324081924454064)])
-        self.assertEqual(self.MutualInformation(["b", "d"]), [Row(value=0.6365141682948128)])
+        # Spark 3.5 and 4.1 can differ by one ULP in mutual information results.
+        for columns, expected in [
+            (["b", "c"], 0.7324081924454064),
+            (["b", "d"], 0.6365141682948128),
+        ]:
+            result = self.MutualInformation(columns)
+            self.assertEqual(len(result), 1)
+            self.assertAlmostEqual(result[0].value, expected, delta=1e-15)
 
     @pytest.mark.xfail(reason="@unittest.expectedFailure")
     def test_fail_MutualInformation(self):

tests/test_checks.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import math
 import unittest
 from typing import List, Union
 
@@ -683,7 +684,12 @@ def test_hasMutualInformation(self):
             [Row(constraint_status="Success")],
         )
         self.assertEqual(
-            self.hasMutualInformation("c", "b", lambda x: x == 0.7324081924454064),
+            # Spark 3.5 and 4.1 can differ by one ULP in mutual information results.
+            self.hasMutualInformation(
+                "c",
+                "b",
+                lambda x: math.isclose(x, 0.7324081924454064, rel_tol=0.0, abs_tol=1e-15),
+            ),
             [Row(constraint_status="Success")],
         )
 

tests/test_config.py

@@ -1,5 +1,10 @@
 import pytest
-from pydeequ.configs import _extract_major_minor_versions
+from pydeequ.configs import (
+    SPARK_TO_DEEQU_COORD_MAPPING,
+    _extract_major_minor_versions,
+    _get_deequ_maven_config,
+    _get_spark_version,
+)
 
 
 @pytest.mark.parametrize(
@@ -13,3 +18,14 @@
 )
 def test_extract_major_minor_versions(full_version, major_minor_version):
     assert _extract_major_minor_versions(full_version) == major_minor_version
+
+
+def test_spark_4_1_deequ_coordinate(monkeypatch):
+    monkeypatch.setenv("SPARK_VERSION", "4.1.2")
+    _get_spark_version.cache_clear()
+
+    try:
+        assert _get_deequ_maven_config() == "com.amazon.deequ:deequ:2.0.18-spark-4.1"
+        assert SPARK_TO_DEEQU_COORD_MAPPING["4.1"] == "com.amazon.deequ:deequ:2.0.18-spark-4.1"
+    finally:
+        _get_spark_version.cache_clear()