diff --git a/README.md b/README.md index 88ad70ba..c1fa7879 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,7 @@ torch_transformers = kamae.get_compatible_transformers('torch') |:-------------------------------:|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------:|:------------------------------------------------------------------:|:----------------:|:-------------------------------------------------------------------------:| | AbsoluteValue | Applies the `abs(x)` transform. | [Link](src/kamae/keras/core/layers/absolute_value.py) | Multi-backend | [Link](src/kamae/spark/transformers/absolute_value.py) | | ArrayConcatenate | Assembles multiple features into a single array. | [Link](src/kamae/keras/core/layers/array_concatenate.py) | Multi-backend | [Link](src/kamae/spark/transformers/array_concatenate.py) | +| ArrayContains | Checks whether a scalar value is contained in an array feature. | [Link](src/kamae/keras/core/layers/array_contains.py) | Multi-backend | [Link](src/kamae/spark/transformers/array_contains.py) | | ArrayCrop | Crops or pads a feature array to a consistent size. | [Link](src/kamae/keras/core/layers/array_crop.py) | Multi-backend | [Link](src/kamae/spark/transformers/array_crop.py) | | ArrayReduceMax | Reduces the last dimension of a tensor by taking the maximum. | [Link](src/kamae/keras/core/layers/array_reduce_max.py) | Multi-backend | [Link](src/kamae/spark/transformers/array_reduce_max.py) | | ArraySplit | Splits a feature array into multiple features. | [Link](src/kamae/keras/core/layers/array_split.py) | Multi-backend | [Link](src/kamae/spark/transformers/array_split.py) | diff --git a/src/kamae/keras/core/layers/__init__.py b/src/kamae/keras/core/layers/__init__.py index 474df48c..ab7e2308 100644 --- a/src/kamae/keras/core/layers/__init__.py +++ b/src/kamae/keras/core/layers/__init__.py @@ -20,6 +20,7 @@ from .absolute_value import AbsoluteValueLayer from .array_concatenate import ArrayConcatenateLayer +from .array_contains import ArrayContainsLayer from .array_crop import ArrayCropLayer from .array_reduce_max import ArrayReduceMaxLayer from .array_split import ArraySplitLayer @@ -73,6 +74,7 @@ "LogicalNotLayer", "NumericalIfStatementLayer", "ArrayConcatenateLayer", + "ArrayContainsLayer", "ArrayReduceMaxLayer", "ArraySplitLayer", "ArrayCropLayer", diff --git a/src/kamae/keras/core/layers/array_contains.py b/src/kamae/keras/core/layers/array_contains.py new file mode 100644 index 00000000..a83ba1ed --- /dev/null +++ b/src/kamae/keras/core/layers/array_contains.py @@ -0,0 +1,137 @@ +# 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. + +from typing import Any, Dict, Iterable, List, Optional, Union + +import keras +from keras import KerasTensor, ops + +import kamae +from kamae.keras.core.backend import ALL_BACKENDS +from kamae.keras.core.base import BaseLayer +from kamae.keras.core.utils.input_utils import allow_single_or_multiple_tensor_input + + +@keras.saving.register_keras_serializable(package=kamae.__name__) +class ArrayContainsLayer(BaseLayer): + """ + Computes whether a value is contained in an array along a given axis. + + With `value_constant` set, takes a single `array` input and checks whether + that constant is present. Otherwise takes two inputs `(array, value)`. + Returns a boolean; set `output_dtype` to cast the result. + + Example: + >>> layer = ArrayContainsLayer(value_constant=2) + >>> layer([[1, 2, 3], [4, 5, 6]]) + [[True], [False]] + """ + + supported_backends = ALL_BACKENDS + jit_compatible = True + + def __init__( + self, + name: Optional[str] = None, + input_dtype: Optional[str] = None, + output_dtype: Optional[str] = None, + axis: int = -1, + keepdims: bool = True, + value_constant: Optional[float] = None, + **kwargs: Any, + ) -> None: + """ + Initializes the ArrayContainsLayer layer. + + :param name: Name of the layer, defaults to `None`. + :param input_dtype: The dtype to cast the input to. Defaults to `None`. + :param output_dtype: The dtype to cast the output to. Defaults to `None`. + :param axis: The axis along which to check for the value. Defaults to + `-1`. + :param keepdims: Whether to keep the reduced axis as a size 1 dimension + in the output. Defaults to `True`. + :param value_constant: Optional constant to check for. If set, the layer + takes a single `array` input. If `None`, two inputs `(array, value)` are + required. Defaults to `None`. + """ + super().__init__( + name=name, input_dtype=input_dtype, output_dtype=output_dtype, **kwargs + ) + self.axis = axis + self.keepdims = keepdims + self.value_constant = value_constant + + @property + def compatible_dtypes(self) -> Optional[List[str]]: + """ + Returns the compatible dtypes of the layer. + + :returns: List of compatible dtype names. + """ + return [ + "bfloat16", + "float16", + "float32", + "float64", + "int8", + "uint8", + "int16", + "uint16", + "int32", + "uint32", + "int64", + "uint64", + ] + + @allow_single_or_multiple_tensor_input + def _call( + self, inputs: Union[KerasTensor, Iterable[KerasTensor]], **kwargs: Any + ) -> KerasTensor: + """ + Computes whether `value` is present in `array` along the given axis. + + :param inputs: Single `array` tensor (when `value_constant` is set) or a + list of two tensors `(array, value)`. + :returns: A boolean tensor, `True` where the value is found. + """ + if self.value_constant is not None and len(inputs) == 1: + # Constant value given upon initialization + array, value = inputs[0], self.value_constant + array, value = self._force_cast_to_compatible_numeric_type(array, value) + elif self.value_constant is None and len(inputs) == 2: + # Dynamic value given + array, value = inputs + elif self.value_constant is not None: + raise ValueError("Expected 1 input when `value_constant` is set") + else: + raise ValueError(f"Expected 2 inputs, got {len(inputs)} inputs instead") + + return ops.any(ops.equal(array, value), axis=self.axis, keepdims=self.keepdims) + + def get_config(self) -> Dict[str, Any]: + """ + Gets the configuration of the ArrayContains layer. + Used for saving and loading from a model. + + :returns: Dictionary of the configuration of the layer. + """ + config = super().get_config() + config.update( + { + "axis": self.axis, + "keepdims": self.keepdims, + "value_constant": self.value_constant, + } + ) + return config diff --git a/src/kamae/spark/transformers/__init__.py b/src/kamae/spark/transformers/__init__.py index de563832..a61ae8ea 100644 --- a/src/kamae/spark/transformers/__init__.py +++ b/src/kamae/spark/transformers/__init__.py @@ -14,6 +14,7 @@ from .absolute_value import AbsoluteValueTransformer # noqa: F401 from .array_concatenate import ArrayConcatenateTransformer # noqa: F401 +from .array_contains import ArrayContainsTransformer # noqa: F401 from .array_crop import ArrayCropTransformer # noqa: F401 from .array_reduce_max import ArrayReduceMaxTransformer # noqa: F401 from .array_split import ArraySplitTransformer # noqa: F401 diff --git a/src/kamae/spark/transformers/array_contains.py b/src/kamae/spark/transformers/array_contains.py new file mode 100644 index 00000000..6e818e84 --- /dev/null +++ b/src/kamae/spark/transformers/array_contains.py @@ -0,0 +1,188 @@ +# 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. + +# pylint: disable=unused-argument +# pylint: disable=invalid-name +# pylint: disable=too-many-ancestors +# pylint: disable=no-member +from typing import List, Optional + +import keras +import pyspark.sql.functions as F +from pyspark import keyword_only +from pyspark.sql import DataFrame +from pyspark.sql.types import ( + ByteType, + DataType, + DoubleType, + FloatType, + FractionalType, + IntegerType, + LongType, + ShortType, +) + +from kamae.keras.core.backend import ALL_BACKENDS +from kamae.keras.core.layers import ArrayContainsLayer +from kamae.spark.params import ( + MathFloatConstantParams, + MultiInputSingleOutputParams, + SingleInputSingleOutputParams, +) +from kamae.spark.utils import ( + get_element_type, + single_input_single_output_array_transform, +) + +from .base import BaseTransformer + + +class ArrayContainsTransformer( + BaseTransformer, + SingleInputSingleOutputParams, + MultiInputSingleOutputParams, + MathFloatConstantParams, +): + """ + ArrayContainsLayer Spark Transformer for use in Spark pipelines. + + Checks whether a scalar value is present in a (possibly nested) numeric + array. The value is either a second input column (`inputCols`) or the + `mathFloatConstant` (`inputCol`). Outputs a boolean; set `outputDtype` to + cast the result. + + Example: + + >>> df.show() + +---------+ + | a| + +---------+ + |[1, 2, 3]| + |[4, 5, 6]| + +---------+ + >>> t = ArrayContainsTransformer(inputCol="a", outputCol="b", mathFloatConstant=2) + >>> t.transform(df).show() + +---------+-----+ + | a| b| + +---------+-----+ + |[1, 2, 3]| true| + |[4, 5, 6]|false| + +---------+-----+ + """ + + supported_backends = ALL_BACKENDS + jit_compatible = True + + @keyword_only + def __init__( + self, + inputCol: Optional[str] = None, + inputCols: Optional[List[str]] = None, + outputCol: Optional[str] = None, + inputDtype: Optional[str] = None, + outputDtype: Optional[str] = None, + layerName: Optional[str] = None, + mathFloatConstant: Optional[float] = None, + ) -> None: + """ + Initializes an ArrayContainsTransformer transformer. + + :param inputCol: Input array column name. Only used if inputCols is not + specified. If specified, we check whether `mathFloatConstant` is contained + in this array column. + :param inputCols: Input column names, given as `[arrayCol, valueCol]`. + :param outputCol: Output column name. + :param inputDtype: Input data type to cast input column(s) to before + transforming. + :param outputDtype: Output data type to cast the output column to after + transforming. + :param layerName: Name of the layer. Used as the name of the Keras layer + in the keras model. If not set, we use the uid of the Spark transformer. + :param mathFloatConstant: Optional constant value to check for. Used with + `inputCol`. If not provided, then `inputCols` is required. + :returns: None - class instantiated. + """ + super().__init__() + self._setDefault(mathFloatConstant=None) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @property + def compatible_dtypes(self) -> Optional[List[DataType]]: + """ + List of compatible data types for the layer. + If the computation can be performed on any data type, return None. + + :returns: List of compatible data types for the layer. + """ + return [ + FloatType(), + DoubleType(), + ByteType(), + ShortType(), + IntegerType(), + LongType(), + ] + + def setInputCols(self, value: List[str]) -> "ArrayContainsTransformer": + """ + Sets the input columns, ensuring exactly two are provided: + `[arrayCol, valueCol]`. + + :param value: List of two input column names. + :returns: Instance of class with input columns set. + """ + if len(value) != 2: + raise ValueError(f"Expected 2 input cols, received {len(value)} instead.") + + return self._set(inputCols=value) + + def _transform(self, dataset: DataFrame) -> DataFrame: + """ + Adds `outputCol`, `True` where the value is present in the innermost + array. The value is a second input column or `mathFloatConstant`. + + :param dataset: Pyspark dataframe to transform. + :returns: Transformed pyspark dataframe. + """ + array_col, value_col = self.get_multiple_input_cols("mathFloatConstant", 2) + df = dataset.select(array_col, value_col) + c_dt = self.get_column_datatype(df, df.columns[0]) + e_dt = get_element_type(c_dt) + + if self.getMathFloatConstant() is not None and isinstance(e_dt, FractionalType): + value_col = value_col.cast(e_dt) + + output_col = single_input_single_output_array_transform( + input_col=array_col, + input_col_datatype=c_dt, + func=lambda x: F.coalesce(F.array_contains(x, value_col), F.lit(False)), + ) + return dataset.withColumn(self.getOutputCol(), output_col) + + def get_keras_layer(self) -> keras.layers.Layer: + """ + Gets the Keras layer for the array contains transformer. + + :returns: Keras layer with name equal to the layerName parameter that + performs the array contains operation. + """ + return ArrayContainsLayer( + name=self.getLayerName(), + input_dtype=self.getInputKerasDtype(), + output_dtype=self.getOutputKerasDtype(), + axis=-1, + keepdims=True, + value_constant=self.getMathFloatConstant(), + ) diff --git a/tests/kamae/keras/tensorflow/layers/test_array_contains.py b/tests/kamae/keras/tensorflow/layers/test_array_contains.py new file mode 100644 index 00000000..1e7c8d50 --- /dev/null +++ b/tests/kamae/keras/tensorflow/layers/test_array_contains.py @@ -0,0 +1,270 @@ +# 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. + +import pytest +import tensorflow as tf + +from kamae.keras.core.layers import ArrayContainsLayer + + +class TestArrayContains: + @pytest.mark.parametrize( + "input_tensors, input_name, input_dtype, output_dtype, expected_output", + [ + ( + [ + tf.constant([[[1, 2, 3]]]), + tf.constant([[[2]]]), + ], + "input_1", + None, + None, + tf.constant([[[True]]]), + ), + ( + [ + tf.constant([[[1, 2, 3]]]), + tf.constant([[[5]]]), + ], + "input_2", + None, + None, + tf.constant([[[False]]]), + ), + ( + [ + tf.constant([[[1, 2, 3]]]), + tf.constant([[[2], [9]]]), + ], + "input_3", + None, + "float64", + tf.constant([[[1.0], [0.0]]], dtype="float64"), + ), + ( + [ + tf.constant( + [ + [[1.5, 2.5, 3.5]], + [[4.5, 5.5, 6.5]], + ] + ), + tf.constant( + [ + [[2.5]], + [[7.5]], + ] + ), + ], + "input_4", + None, + None, + tf.constant([[[True]], [[False]]]), + ), + ( + [ + tf.constant([["1", "2", "3"]]), + tf.constant([["2"]]), + ], + "input_5", + "int64", + None, + tf.constant([[True]]), + ), + ( + # Integer array and integer scalar value, boolean output. + [ + tf.constant([[[10, 20, 30]]], dtype="int64"), + tf.constant([[[20]]], dtype="int64"), + ], + "input_6", + None, + None, + tf.constant([[[True]]]), + ), + ( + # Integer array and integer scalar value, cast to int output. + [ + tf.constant([[[10, 20, 30]]], dtype="int32"), + tf.constant([[[40]]], dtype="int32"), + ], + "input_7", + None, + "int32", + tf.constant([[[0]]], dtype="int32"), + ), + ], + ) + def test_array_contains( + self, + input_tensors, + input_name, + input_dtype, + output_dtype, + expected_output, + ): + # when + layer = ArrayContainsLayer( + name=input_name, + input_dtype=input_dtype, + output_dtype=output_dtype, + ) + output_tensor = layer(input_tensors) + # then + assert layer.name == input_name, "Layer name is not set properly" + assert ( + output_tensor.dtype == expected_output.dtype + ), "Output tensor dtype is not the same as expected tensor dtype" + assert ( + output_tensor.shape == expected_output.shape + ), "Output tensor shape is not the same as expected tensor shape" + + tf.debugging.assert_equal( + tf.cast(output_tensor, "float64"), tf.cast(expected_output, "float64") + ) + + @pytest.mark.parametrize( + "input_tensors, axis, keepdims, expected_output", + [ + ( + # Default: search over last axis, keep the collapsed dimension. + [ + tf.constant([[[1, 2, 3]]]), + tf.constant([[[2]]]), + ], + -1, + True, + tf.constant([[[True]]]), + ), + ( + # keepdims=False drops the collapsed axis. + [ + tf.constant([[[1, 2, 3]]]), + tf.constant([[[2]]]), + ], + -1, + False, + tf.constant([[True]]), + ), + ( + # Search over a non-final axis. + [ + tf.constant([[[1], [2], [3]]]), + tf.constant([[[2]]]), + ], + 1, + True, + tf.constant([[[True]]]), + ), + ], + ) + def test_array_contains_axis_keepdims( + self, + input_tensors, + axis, + keepdims, + expected_output, + ): + # when + layer = ArrayContainsLayer(axis=axis, keepdims=keepdims) + output_tensor = layer(input_tensors) + # then + assert ( + output_tensor.shape == expected_output.shape + ), "Output tensor shape is not the same as expected tensor shape" + tf.debugging.assert_equal(output_tensor, expected_output) + + @pytest.mark.parametrize( + "input_tensor, value_constant, output_dtype, expected_output", + [ + ( + # Single tensor input with a constant value to search for. + tf.constant([[[1, 2, 3]]]), + 2.0, + None, + tf.constant([[[True]]]), + ), + ( + tf.constant([[[1, 2, 3]]]), + 5.0, + None, + tf.constant([[[False]]]), + ), + ( + # Constant search with cast output. + tf.constant([[[1, 2, 3]], [[4, 5, 6]]]), + 4.0, + "float64", + tf.constant([[[0.0]], [[1.0]]], dtype="float64"), + ), + ], + ) + def test_array_contains_value_constant( + self, + input_tensor, + value_constant, + output_dtype, + expected_output, + ): + # when + layer = ArrayContainsLayer( + value_constant=value_constant, + output_dtype=output_dtype, + ) + output_tensor = layer(input_tensor) + # then + assert ( + output_tensor.shape == expected_output.shape + ), "Output tensor shape is not the same as expected tensor shape" + tf.debugging.assert_equal( + tf.cast(output_tensor, "float64"), tf.cast(expected_output, "float64") + ) + + def test_array_contains_value_constant_multiple_inputs_raises_error(self): + # given + layer = ArrayContainsLayer(value_constant=2.0) + # then + with pytest.raises(ValueError): + layer( + [ + tf.constant([[[1, 2, 3]]]), + tf.constant([[[2]]]), + ] + ) + + @pytest.mark.parametrize( + "input_tensors", + [ + ( + [ + # Too many input tensors + tf.constant([[[1, 2, 3]]]), + tf.constant([[[2]]]), + tf.constant([[[3]]]), + ], + ), + ( + [ + # Not enough input tensors + tf.constant([[[1, 2, 3]]]), + ], + ), + ], + ) + def test_array_contains_raises_error(self, input_tensors): + # when + layer = ArrayContainsLayer() + # then + with pytest.raises(ValueError): + layer(input_tensors) diff --git a/tests/kamae/keras/test_jit_compatibility.py b/tests/kamae/keras/test_jit_compatibility.py index bd9c2354..2ec2d600 100644 --- a/tests/kamae/keras/test_jit_compatibility.py +++ b/tests/kamae/keras/test_jit_compatibility.py @@ -25,6 +25,7 @@ from kamae.keras.core.layers import ( AbsoluteValueLayer, ArrayConcatenateLayer, + ArrayContainsLayer, ArrayCropLayer, ArrayReduceMaxLayer, ArraySplitLayer, @@ -107,6 +108,11 @@ [tf.random.normal((32, 10, 100, 3)), tf.random.normal((32, 10, 100, 3))], {"axis": -2}, ), + ( + ArrayContainsLayer, + [tf.random.normal((32, 1, 10)), tf.random.normal((32, 5, 1))], + None, + ), (ArrayReduceMaxLayer, [tf.random.normal((32, 10))], {"default_value": 0.0}), (ArraySplitLayer, [tf.random.normal((32, 10, 100, 3))], {"axis": -2}), ( diff --git a/tests/kamae/keras/test_layer_serialisation.py b/tests/kamae/keras/test_layer_serialisation.py index e634e7b3..a4679777 100644 --- a/tests/kamae/keras/test_layer_serialisation.py +++ b/tests/kamae/keras/test_layer_serialisation.py @@ -32,6 +32,7 @@ from kamae.keras.core.layers import ( AbsoluteValueLayer, ArrayConcatenateLayer, + ArrayContainsLayer, ArrayCropLayer, ArrayReduceMaxLayer, ArraySplitLayer, @@ -118,6 +119,18 @@ {"axis": -2}, False, ), + ( + ArrayContainsLayer, + [tf.random.normal((32, 1, 10)), tf.random.normal((32, 5, 1))], + {"axis": -1, "keepdims": True}, + False, + ), + ( + ArrayContainsLayer, + [tf.random.normal((32, 1, 10))], + {"axis": -1, "keepdims": True, "value_constant": 2.0}, + False, + ), ( ArrayReduceMaxLayer, [tf.random.normal((32, 10))], diff --git a/tests/kamae/spark/transformers/test_array_contains.py b/tests/kamae/spark/transformers/test_array_contains.py new file mode 100644 index 00000000..dd5c791e --- /dev/null +++ b/tests/kamae/spark/transformers/test_array_contains.py @@ -0,0 +1,488 @@ +# 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. + +import numpy as np +import pytest +import tensorflow as tf +from pyspark.sql.types import ArrayType, DoubleType, FloatType, StructField, StructType + +from kamae.spark.transformers import ArrayContainsTransformer + + +class TestArrayContains: + @pytest.fixture(scope="class") + def example_dataframe_with_arrays(self, spark_session): + return spark_session.createDataFrame( + [ + ([1, 2, 3], 2, 4), + ([1, 2, 3], 5, 1), + ([4, 5, 6], 4, 9), + ], + ["array_col", "value_col", "other_value_col"], + ) + + @pytest.fixture(scope="class") + def array_contains_transform_array_value_expected(self, spark_session): + return spark_session.createDataFrame( + [ + ([1, 2, 3], 2, 4, True), + ([1, 2, 3], 5, 1, False), + ([4, 5, 6], 4, 9, True), + ], + ["array_col", "value_col", "other_value_col", "array_contains_value"], + ) + + @pytest.fixture(scope="class") + def array_contains_transform_array_other_value_expected(self, spark_session): + return spark_session.createDataFrame( + [ + ([1, 2, 3], 2, 4, False), + ([1, 2, 3], 5, 1, True), + ([4, 5, 6], 4, 9, False), + ], + [ + "array_col", + "value_col", + "other_value_col", + "array_contains_other_value", + ], + ) + + @pytest.fixture(scope="class") + def example_dataframe_with_nested_arrays(self, spark_session): + return spark_session.createDataFrame( + [ + ([[1, 2, 3], [4, 5, 6]], 4), + ([[1, 2, 3], [4, 5, 6]], 9), + ([[7, 8, 9], [1, 1, 1]], 1), + ], + ["nested_array_col", "value_col"], + ) + + @pytest.fixture(scope="class") + def array_contains_transform_nested_expected(self, spark_session): + return spark_session.createDataFrame( + [ + ([[1, 2, 3], [4, 5, 6]], 4, [False, True]), + ([[1, 2, 3], [4, 5, 6]], 9, [False, False]), + ([[7, 8, 9], [1, 1, 1]], 1, [False, True]), + ], + ["nested_array_col", "value_col", "nested_array_contains"], + ) + + @pytest.mark.parametrize( + "input_cols, output_col, expected_dataframe", + [ + ( + ["array_col", "value_col"], + "array_contains_value", + "array_contains_transform_array_value_expected", + ), + ( + ["array_col", "other_value_col"], + "array_contains_other_value", + "array_contains_transform_array_other_value_expected", + ), + ], + ) + def test_spark_array_contains_transform( + self, + example_dataframe_with_arrays, + input_cols, + output_col, + expected_dataframe, + request, + ): + # given + expected = request.getfixturevalue(expected_dataframe) + # when + transformer = ArrayContainsTransformer( + inputCols=input_cols, + outputCol=output_col, + ) + actual = transformer.transform(example_dataframe_with_arrays) + # then + diff = actual.exceptAll(expected) + assert diff.isEmpty(), "Expected and actual dataframes are not equal" + + def test_spark_array_contains_transform_nested_arrays( + self, + example_dataframe_with_nested_arrays, + array_contains_transform_nested_expected, + ): + # given + expected = array_contains_transform_nested_expected + # when + transformer = ArrayContainsTransformer( + inputCols=["nested_array_col", "value_col"], + outputCol="nested_array_contains", + ) + actual = transformer.transform(example_dataframe_with_nested_arrays) + # then + diff = actual.exceptAll(expected) + assert diff.isEmpty(), "Expected and actual dataframes are not equal" + + def test_spark_array_contains_transform_null_elements(self, spark_session): + # given + # array_contains returns null when the array holds a null and the value + # is not otherwise present (or when the array itself is null); coalesce to + # False keeps parity with the always-boolean Keras output. + input_df = spark_session.createDataFrame( + [ + ([1, 2, 3], 5), + ([1, None, 3], 5), + ([1, None, 3], 1), + (None, 5), + ], + ["array_col", "value_col"], + ) + expected = spark_session.createDataFrame( + [ + ([1, 2, 3], 5, False), + ([1, None, 3], 5, False), + ([1, None, 3], 1, True), + (None, 5, False), + ], + ["array_col", "value_col", "array_contains_value"], + ) + # when + transformer = ArrayContainsTransformer( + inputCols=["array_col", "value_col"], + outputCol="array_contains_value", + ) + actual = transformer.transform(input_df) + # then + diff = actual.exceptAll(expected) + assert diff.isEmpty(), "Expected and actual dataframes are not equal" + + def test_spark_array_contains_transform_output_dtype_cast( + self, spark_session, example_dataframe_with_arrays + ): + # given + expected = spark_session.createDataFrame( + [ + ([1, 2, 3], 2, 4, 1.0), + ([1, 2, 3], 5, 1, 0.0), + ([4, 5, 6], 4, 9, 1.0), + ], + ["array_col", "value_col", "other_value_col", "array_contains_value"], + ) + # when + transformer = ArrayContainsTransformer( + inputCols=["array_col", "value_col"], + outputCol="array_contains_value", + outputDtype="double", + ) + actual = transformer.transform(example_dataframe_with_arrays) + # then + assert isinstance( + actual.schema["array_contains_value"].dataType, DoubleType + ), "Output column should be cast to DoubleType" + diff = actual.exceptAll(expected) + assert diff.isEmpty(), "Expected and actual dataframes are not equal" + + def test_array_contains_transform_defaults(self): + # when + array_contains_transform = ArrayContainsTransformer() + # then + assert array_contains_transform.getLayerName() == array_contains_transform.uid + assert ( + array_contains_transform.getOutputCol() + == f"{array_contains_transform.uid}__output" + ) + + @pytest.mark.parametrize( + "input_cols", + [ + ["array_col"], + ["array_col", "value_col", "other_value_col"], + ], + ) + def test_array_contains_transform_wrong_number_of_inputs_raises_error( + self, input_cols + ): + # then + with pytest.raises(ValueError): + ArrayContainsTransformer(inputCols=input_cols) + + def test_array_contains_transform_non_array_input_raises_error( + self, example_dataframe_with_arrays + ): + # given + transformer = ArrayContainsTransformer( + inputCols=["value_col", "other_value_col"], + outputCol="array_contains_output", + ) + # then + with pytest.raises(ValueError): + transformer.transform(example_dataframe_with_arrays).collect() + + def test_spark_array_contains_transform_constant( + self, spark_session, example_dataframe_with_arrays + ): + # given + expected = spark_session.createDataFrame( + [ + ([1, 2, 3], 2, 4, True), + ([1, 2, 3], 5, 1, True), + ([4, 5, 6], 4, 9, False), + ], + ["array_col", "value_col", "other_value_col", "array_contains_const"], + ) + # when + transformer = ArrayContainsTransformer( + inputCol="array_col", + outputCol="array_contains_const", + mathFloatConstant=2.0, + ) + actual = transformer.transform(example_dataframe_with_arrays) + # then + diff = actual.exceptAll(expected) + assert diff.isEmpty(), "Expected and actual dataframes are not equal" + + def test_spark_array_contains_transform_constant_nested( + self, spark_session, example_dataframe_with_nested_arrays + ): + # given + expected = spark_session.createDataFrame( + [ + ([[1, 2, 3], [4, 5, 6]], 4, [False, True]), + ([[1, 2, 3], [4, 5, 6]], 9, [False, True]), + ([[7, 8, 9], [1, 1, 1]], 1, [False, False]), + ], + ["nested_array_col", "value_col", "nested_array_contains_const"], + ) + # when + transformer = ArrayContainsTransformer( + inputCol="nested_array_col", + outputCol="nested_array_contains_const", + mathFloatConstant=5.0, + ) + actual = transformer.transform(example_dataframe_with_nested_arrays) + # then + diff = actual.exceptAll(expected) + assert diff.isEmpty(), "Expected and actual dataframes are not equal" + + def test_array_contains_transform_no_value_raises_error( + self, example_dataframe_with_arrays + ): + # given: single inputCol but no constant provided + transformer = ArrayContainsTransformer( + inputCol="array_col", + outputCol="array_contains_output", + ) + # then + with pytest.raises(ValueError): + transformer.transform(example_dataframe_with_arrays).collect() + + @pytest.mark.parametrize( + "input_arrays, input_values, input_dtype, output_dtype", + [ + ( + [[1, 2, 3], [1, 2, 3], [4, 5, 6]], + [2, 5, 4], + None, + None, + ), + ( + [[10, 20, 30, 40], [5, 6, 7, 8], [0, 0, 0, 0]], + [30, 100, 0], + "bigint", + "double", + ), + ], + ) + def test_array_contains_transform_spark_tf_parity( + self, + spark_session, + input_arrays, + input_values, + input_dtype, + output_dtype, + ): + # given + transformer = ArrayContainsTransformer( + inputCols=["array_col", "value_col"], + outputCol="output", + inputDtype=input_dtype, + outputDtype=output_dtype, + ) + # when + spark_df = spark_session.createDataFrame( + zip(input_arrays, input_values), + ["array_col", "value_col"], + ) + spark_values = ( + transformer.transform(spark_df) + .select("output") + .rdd.map(lambda r: r[0]) + .collect() + ) + array_tensor = tf.constant(input_arrays) + value_tensor = tf.constant([[v] for v in input_values]) + tensorflow_values = ( + transformer.get_keras_layer()([array_tensor, value_tensor]) + .numpy() + .flatten() + .tolist() + ) + + # then + np.testing.assert_almost_equal( + spark_values, + tensorflow_values, + ) + + @pytest.mark.parametrize( + "input_arrays, constant, input_dtype, output_dtype", + [ + ( + [[1, 2, 3], [4, 5, 6], [7, 8, 9]], + 5.0, + None, + None, + ), + ( + [[10, 20, 30, 40], [5, 6, 7, 8], [0, 0, 0, 0]], + 0.0, + "bigint", + "double", + ), + ( + # Non-integer constant on an integer array must stay unmatched on + # both sides (Spark must not narrow 2.5 -> 2). + [[1, 2, 3], [2, 4, 6], [7, 8, 9]], + 2.5, + None, + None, + ), + ( + # Same as above but int64: Keras widens the array to float64, so + # Spark must also compare in double (no narrowing) and stay unmatched. + [[1, 2, 3], [2, 4, 6], [7, 8, 9]], + 2.5, + "bigint", + None, + ), + ], + ) + def test_array_contains_transform_constant_spark_tf_parity( + self, + spark_session, + input_arrays, + constant, + input_dtype, + output_dtype, + ): + # given + transformer = ArrayContainsTransformer( + inputCol="array_col", + outputCol="output", + inputDtype=input_dtype, + outputDtype=output_dtype, + mathFloatConstant=constant, + ) + # when + spark_df = spark_session.createDataFrame( + [(a,) for a in input_arrays], + ["array_col"], + ) + spark_values = ( + transformer.transform(spark_df) + .select("output") + .rdd.map(lambda r: r[0]) + .collect() + ) + array_tensor = tf.constant(input_arrays) + tensorflow_values = ( + transformer.get_keras_layer()(array_tensor).numpy().flatten().tolist() + ) + + # then + np.testing.assert_almost_equal( + spark_values, + tensorflow_values, + ) + + def test_array_contains_transform_constant_float32_spark_tf_parity( + self, spark_session + ): + # given + # 0.1 isn't exactly representable in float32. With inputDtype=None nothing + # aligns the array and constant, so Spark must narrow the literal to the + # array's float32 element type to match the Keras layer (which narrows the + # constant); otherwise Spark compares in double and diverges. + input_arrays = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + schema = StructType([StructField("array_col", ArrayType(FloatType()))]) + transformer = ArrayContainsTransformer( + inputCol="array_col", + outputCol="output", + mathFloatConstant=0.1, + ) + # when + spark_df = spark_session.createDataFrame( + [(a,) for a in input_arrays], + schema, + ) + spark_values = ( + transformer.transform(spark_df) + .select("output") + .rdd.map(lambda r: r[0]) + .collect() + ) + array_tensor = tf.constant(input_arrays, dtype=tf.float32) + tensorflow_values = ( + transformer.get_keras_layer()(array_tensor).numpy().flatten().tolist() + ) + + # then + np.testing.assert_almost_equal( + spark_values, + tensorflow_values, + ) + + def test_array_contains_transform_constant_float64_spark_tf_parity( + self, spark_session + ): + # given + # Double array: Keras narrows the constant to float64, so Spark's double + # literal already aligns (no-op cast) and both compare in double. + input_arrays = [[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]] + schema = StructType([StructField("array_col", ArrayType(DoubleType()))]) + transformer = ArrayContainsTransformer( + inputCol="array_col", + outputCol="output", + mathFloatConstant=2.2, + ) + # when + spark_df = spark_session.createDataFrame( + [(a,) for a in input_arrays], + schema, + ) + spark_values = ( + transformer.transform(spark_df) + .select("output") + .rdd.map(lambda r: r[0]) + .collect() + ) + array_tensor = tf.constant(input_arrays, dtype=tf.float64) + tensorflow_values = ( + transformer.get_keras_layer()(array_tensor).numpy().flatten().tolist() + ) + + # then + np.testing.assert_almost_equal( + spark_values, + tensorflow_values, + )