feat: support Spark 4 and numpy 2

.github/workflows/ci.yaml

@@ -17,15 +17,27 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.10", "3.11", "3.12"]
-        pyspark-version: ["3.4.1", "3.5.0"]
+        pyspark-version: ["3.4.1", "3.5.0", "4.0.1"]
         keras-version: ["3.3.0", "3.7.0", "3.10.0", "3.12.0"]
+        include:
+          # numpy is otherwise resolved to 2.x, so keep coverage of numpy 1.x.
+          - python-version: "3.10"
+            pyspark-version: "3.5.0"
+            keras-version: "3.10.0"
+            numpy-version: "1.26.4"
     steps:
       - name: Setup Local Repo
         uses: actions/checkout@v4
       - name: Setup Python
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
+      - name: Setup Java
+        # Spark 4 requires Java 17+, which Spark 3.4 and 3.5 also support.
+        uses: actions/setup-java@v4
+        with:
+          distribution: zulu
+          java-version: "17"
       - name: Install uv
         run: |
           pip install --upgrade pip
@@ -34,7 +46,7 @@ jobs:
         run: |
           uv venv --python ${{ matrix.python-version }}
           uv pip install -e ".[tensorflow]"
-          uv pip install pyspark==${{ matrix.pyspark-version }} keras==${{ matrix.keras-version }}
+          uv pip install pyspark==${{ matrix.pyspark-version }} keras==${{ matrix.keras-version }} ${{ matrix.numpy-version && format('numpy=={0}', matrix.numpy-version) || '' }}
       - name: Run tests
         run: uv run -p ${{ matrix.python-version }} python -m pytest -n auto .
   formatting:

pyproject.toml

@@ -10,14 +10,14 @@ license = "Apache-2.0"
 license-files = ["LICENSE.txt"]
 requires-python = ">=3.10,<3.13"
 dependencies = [
-    "pyspark>=3.4.0,<4.0.0",
+    "pyspark>=3.4.0,<5.0.0",
     "pandas>=1.3.4,<3.0.0",
     "pyarrow>=4.0.0",
     "networkx>=2.6.3,<3.0.0",
     "pyfarmhash>=0.3.2,<0.4.0",
     "keras>=3.0.0,<4.0.0",
     "keras-tuner>=1.4.0,<2.0.0",
-    "numpy>=1.22.0,<2.0.0",
+    "numpy>=1.22.0",
     "tensorflow>=2.16.0,<3.0.0",
     "dill>=0.3.0,<1.0.0",
 ]

src/kamae/spark/transformers/cosine_similarity.py

@@ -133,15 +133,20 @@ def norm(x: Column, col_name: str) -> Column:
             )
 
         # If the norms are zero, then we match the tensorflow behavior and return 0.0.
+        # The zero check is explicit since with ANSI mode enabled (the default in
+        # Spark 4) division by zero raises an error rather than returning null.
+        def cosine_similarity(x: Column) -> Column:
+            norm_product = norm(x, input_col_names[0]) * norm(x, input_col_names[1])
+            return F.coalesce(
+                F.when(norm_product != F.lit(0.0), dot_product(x) / norm_product),
+                F.lit(0.0),
+            )
+
         output_col = multi_input_single_output_array_transform(
             input_cols=input_cols,
             input_col_datatypes=datatypes,
             input_col_names=input_col_names,
-            func=lambda x: F.coalesce(
-                dot_product(x)
-                / (norm(x, input_col_names[0]) * norm(x, input_col_names[1])),
-                F.lit(0.0),
-            ),
+            func=cosine_similarity,
         )
         return dataset.withColumn(self.getOutputCol(), output_col)
 

src/kamae/spark/transformers/divide.py

@@ -121,7 +121,11 @@ def divide_no_nan(column1: Column, column2: Column) -> Column:
             """
             Divide two columns, and if the result is NaN, return 0.0 instead
             """
-            return F.coalesce(column1 / column2, F.lit(0.0))
+            # Guard the divisor explicitly, since with ANSI mode enabled (the default
+            # in Spark 4) division by zero raises an error rather than returning null.
+            return F.coalesce(
+                F.when(column2 != F.lit(0), column1 / column2), F.lit(0.0)
+            )
 
         output_col = multi_input_single_output_scalar_transform(
             input_cols=input_cols,

src/kamae/spark/transformers/list_median.py

@@ -162,12 +162,17 @@ def _transform(self, dataset: DataFrame) -> DataFrame:
         dataset = dataset.withColumn("array_size", F.size(F.col("sorted_values")))
         mid_index_1 = (F.col("array_size") / 2).cast("int")
         mid_index_2 = (F.col("array_size") / 2 - 1).cast("int")
+        # Empty arrays give a null median (filled below). This is explicit since with
+        # ANSI mode enabled (the default in Spark 4) element_at raises an error for
+        # an out of bounds index rather than returning null.
         dataset = dataset.withColumn(
             self.getOutputCol(),
-            F.when(
+            F.when(F.col("array_size") <= 0, F.lit(None))
+            .when(
                 F.col("array_size") % 2 == 1,
                 F.element_at(F.col("sorted_values"), mid_index_1 + 1),
-            ).otherwise(
+            )
+            .otherwise(
                 (
                     F.element_at(F.col("sorted_values"), mid_index_1 + 1)
                     + F.element_at(F.col("sorted_values"), mid_index_2 + 1)

src/kamae/spark/transformers/modulo.py

@@ -179,14 +179,20 @@ def _transform(self, dataset: DataFrame) -> DataFrame:
             input_cols=input_cols,
             input_col_names=input_col_names,
             input_col_datatypes=input_col_datatypes,
+            # A zero divisor returns null. This is explicit since with ANSI mode
+            # enabled (the default in Spark 4) modulo by zero raises an error.
             func=lambda x: F.when(
-                # If the result of the modulo operation is negative,
-                # add the divisor to the result. This is to match the behavior of the
-                # tensorflow modulo layer.
-                x[input_col_names[0]] % x[input_col_names[1]] >= 0,
-                x[input_col_names[0]] % x[input_col_names[1]],
-            ).otherwise(
-                (x[input_col_names[0]] % x[input_col_names[1]]) + x[input_col_names[1]]
+                x[input_col_names[1]] != F.lit(0),
+                F.when(
+                    # If the result of the modulo operation is negative,
+                    # add the divisor to the result. This is to match the behavior of
+                    # the tensorflow modulo layer.
+                    x[input_col_names[0]] % x[input_col_names[1]] >= 0,
+                    x[input_col_names[0]] % x[input_col_names[1]],
+                ).otherwise(
+                    (x[input_col_names[0]] % x[input_col_names[1]])
+                    + x[input_col_names[1]]
+                ),
             ),
         )
         return dataset.withColumn(self.getOutputCol(), output_col)

src/kamae/spark/transformers/pairwise_cosine_similarity.py

@@ -126,7 +126,12 @@ def cosine_sim_at_index(idx: Column) -> Column:
                     lambda acc, x: acc + (x * x).cast("double"),
                 )
             )
-            return F.coalesce(dot / (query_norm * cand_norm), F.lit(0.0))
+            # Guard against zero norms explicitly, since with ANSI mode enabled (the
+            # default in Spark 4) division by zero raises rather than returning null.
+            norm_product = query_norm * cand_norm
+            return F.coalesce(
+                F.when(norm_product != F.lit(0.0), dot / norm_product), F.lit(0.0)
+            )
 
         similarities = F.transform(indices, cosine_sim_at_index)
         return dataset.withColumn(self.getOutputCol(), similarities)

src/kamae/spark/transformers/sub_string_delim_at_index.py

@@ -23,7 +23,7 @@
 import tensorflow as tf
 from pyspark import keyword_only
 from pyspark.ml.param import Param, Params, TypeConverters
-from pyspark.sql import DataFrame
+from pyspark.sql import Column, DataFrame
 from pyspark.sql.types import DataType, StringType
 
 from kamae.keras.core.backend import TENSORFLOW_ONLY
@@ -198,13 +198,24 @@ def _transform(self, dataset: DataFrame) -> DataFrame:
         # Since element_at is a 1-based index , we need to add 1 to the index if it
         # is non-negative.
         one_based_index = index + 1 if index >= 0 else index
+
+        # The bounds check is explicit since with ANSI mode enabled (the default in
+        # Spark 4) element_at raises an error for an out of bounds index rather than
+        # returning null.
+        def sub_string_at_index(x: Column) -> Column:
+            split_col = F.split(x, delimiter)
+            return F.coalesce(
+                F.when(
+                    F.size(split_col) >= F.lit(abs(one_based_index)),
+                    F.element_at(split_col, one_based_index),
+                ),
+                F.lit(default_value),
+            )
+
         output_col = single_input_single_output_scalar_transform(
             input_col=F.col(self.getInputCol()),
             input_col_datatype=input_datatype,
-            func=lambda x: F.coalesce(
-                F.element_at(F.split(x, delimiter), one_based_index),
-                F.lit(default_value),
-            ),
+            func=sub_string_at_index,
         )
         return dataset.withColumn(self.getOutputCol(), output_col)
 

tests/kamae/keras/test_jit_compatibility.py

@@ -516,12 +516,12 @@ def jit_call(*inputs):
             return layer(inputs[0])
         return layer(list(inputs))
 
-    # Must raise Exception when trying to JIT compile
+    # Must raise Exception when trying to JIT compile. We lower to HLO rather than
+    # executing, since executing goes through the XLA compilation cache, which in
+    # TF 2.21 aborts the process when hashing string constants after a prior failed
+    # compile (Check failed: DataTypeCanUseMemcpy(dtype())).
     with pytest.raises(Exception):
-        result = jit_call(*input_tensors)
-        # Force evaluation if result is symbolic
-        if hasattr(result, "numpy"):
-            result.numpy()
+        jit_call.experimental_get_compiler_ir(*input_tensors)(stage="hlo")
 
 
 def test_all_layers_define_jit_compatible_and_supported_backends():

tests/kamae/spark/pipeline/test_pipeline.py

@@ -18,7 +18,6 @@
 
 import pytest
 import tensorflow as tf
-from pyspark.sql import DataFrame
 from pyspark.sql.types import DoubleType
 
 from kamae.spark.estimators import (
@@ -619,10 +618,13 @@ def test_spark_pipeline_checkpoint_invocation(
         """
         checkpoint must be invoked during fit only when checkpointInterval > 0.
         """
-        original_checkpoint = DataFrame.checkpoint
+        # Patch the concrete class: on Spark 4 pyspark.sql.DataFrame is an abstract
+        # parent of the classic DataFrame, so patching it would not intercept calls.
+        dataframe_cls = type(example_dataframe)
+        original_checkpoint = dataframe_cls.checkpoint
 
         with patch.object(
-            DataFrame,
+            dataframe_cls,
             "checkpoint",
             autospec=True,
             side_effect=original_checkpoint,

tests/kamae/spark/test_helpers.py

@@ -17,10 +17,10 @@ def tensor_to_python_type(x):
     if x.dtype.name == "string":
         return x.numpy().decode("utf-8")
     elif x.dtype.name == "bool":
-        return bool(x.numpy())
+        return bool(x.numpy().item())
     elif x.dtype.is_floating:
-        return float(x.numpy())
+        return float(x.numpy().item())
     elif x.dtype.is_integer:
-        return int(x.numpy())
+        return int(x.numpy().item())
     else:
         raise ValueError(f"Unknown dtype: {x.dtype}")

tests/kamae/spark/transformers/test_ansi_mode.py

@@ -0,0 +1,106 @@
+# Copyright [2024] Expedia, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Transformers must give the same results whether or not ANSI mode is enabled, since
+it is enabled by default from Spark 4 and raises on division by zero and out of
+bounds array access, where non-ANSI mode returns null."""
+
+import pytest
+
+from kamae.spark.transformers import (
+    CosineSimilarityTransformer,
+    DivideTransformer,
+    ListMedianTransformer,
+    ModuloTransformer,
+    PairwiseCosineSimilarityTransformer,
+    SubStringDelimAtIndexTransformer,
+)
+
+
+@pytest.fixture(params=["true", "false"], ids=["ansi", "non_ansi"])
+def ansi_spark_session(request, spark_session):
+    original = spark_session.conf.get("spark.sql.ansi.enabled")
+    spark_session.conf.set("spark.sql.ansi.enabled", request.param)
+    yield spark_session
+    spark_session.conf.set("spark.sql.ansi.enabled", original)
+
+
+class TestAnsiMode:
+    def test_divide_by_zero(self, ansi_spark_session):
+        df = ansi_spark_session.createDataFrame(
+            [(1.0, 0.0), (0.0, 0.0), (6.0, 2.0)], ["col1", "col2"]
+        )
+        transformer = DivideTransformer(inputCols=["col1", "col2"], outputCol="out")
+        actual = [r.out for r in transformer.transform(df).collect()]
+        assert actual == [0.0, 0.0, 3.0]
+
+    def test_modulo_by_zero(self, ansi_spark_session):
+        df = ansi_spark_session.createDataFrame(
+            [(10, 0), (-7, 3), (7, 3)], ["col1", "col2"]
+        )
+        transformer = ModuloTransformer(inputCols=["col1", "col2"], outputCol="out")
+        actual = [r.out for r in transformer.transform(df).collect()]
+        assert actual == [None, 2, 1]
+
+    def test_cosine_similarity_zero_norm(self, ansi_spark_session):
+        df = ansi_spark_session.createDataFrame(
+            [([0.0, 0.0], [1.0, 2.0]), ([1.0, 0.0], [2.0, 0.0])], ["col1", "col2"]
+        )
+        transformer = CosineSimilarityTransformer(
+            inputCols=["col1", "col2"], outputCol="out"
+        )
+        actual = [r.out for r in transformer.transform(df).collect()]
+        assert actual == [0.0, 1.0]
+
+    def test_pairwise_cosine_similarity_zero_norm(self, ansi_spark_session):
+        df = ansi_spark_session.createDataFrame(
+            [([1.0, 0.0], [0.0, 0.0, 2.0, 0.0])], ["query", "candidates"]
+        )
+        transformer = PairwiseCosineSimilarityTransformer(
+            inputCols=["query", "candidates"], outputCol="out", embeddingDim=2
+        )
+        actual = transformer.transform(df).collect()[0].out
+        assert actual == [0.0, 1.0]
+
+    @pytest.mark.parametrize(
+        "index, expected",
+        [(1, ["b", "default"]), (-1, ["c", "a"]), (5, ["default", "default"])],
+    )
+    def test_sub_string_delim_at_index_out_of_bounds(
+        self, ansi_spark_session, index, expected
+    ):
+        df = ansi_spark_session.createDataFrame([("a_b_c",), ("a",)], ["col1"])
+        transformer = SubStringDelimAtIndexTransformer(
+            inputCol="col1",
+            outputCol="out",
+            delimiter="_",
+            index=index,
+            defaultValue="default",
+        )
+        actual = [r.out for r in transformer.transform(df).collect()]
+        assert actual == expected
+
+    def test_list_median_all_values_filtered(self, ansi_spark_session):
+        df = ansi_spark_session.createDataFrame(
+            [(1, 5.0), (1, 3.0), (2, -1.0), (2, -1.0)], ["search_id", "value"]
+        )
+        transformer = ListMedianTransformer(
+            inputCol="value",
+            outputCol="out",
+            queryIdCol="search_id",
+            minFilterValue=0.0,
+            nanFillValue=-99.0,
+        )
+        actual = {(r.search_id, r.out) for r in transformer.transform(df).collect()}
+        assert actual == {(1, 4.0), (2, -99.0)}