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..243502b5 --- /dev/null +++ b/src/kamae/keras/core/layers/array_contains.py @@ -0,0 +1,112 @@ +# 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 + +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 enforce_multiple_tensor_input + + +@keras.saving.register_keras_serializable(package=kamae.__name__) +class ArrayContainsLayer(BaseLayer): + """ + Computes whether a value is contained in an array along the last axis. + + Expects two inputs `(array, value)` that broadcast on every axis except the + last. The `array` tensor holds the dimension to search over (e.g. shape + `(B, 1, N)`), while the `value` tensor has size 1 on the last axis (e.g. + shape `(B, L, 1)`). The output is the broadcast shape with the last axis + collapsed to 1 (e.g. `(B, L, 1)`), containing `1.0` where the value is found + and `0.0` otherwise. Both inputs must share the same dtype; use `input_dtype` + to cast them to a common dtype. + """ + + supported_backends = ALL_BACKENDS + jit_compatible = True + + def __init__( + self, + name: Optional[str] = None, + input_dtype: Optional[str] = None, + output_dtype: Optional[str] = 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`. + """ + super().__init__( + name=name, input_dtype=input_dtype, output_dtype=output_dtype, **kwargs + ) + + @property + def compatible_dtypes(self) -> Optional[List[str]]: + """ + Returns the compatible dtypes of the layer. + + :returns: List of compatible dtype names. + """ + return [ + "int8", + "uint8", + "int16", + "uint16", + "int32", + "uint32", + "int64", + "uint64", + "float16", + "float32", + "float64", + ] + + @enforce_multiple_tensor_input + def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: + """ + Computes membership of `value` within `array` along the last axis. + + Decorated with `@enforce_multiple_tensor_input` to ensure that the input + is an iterable of tensors. Raises an error if a single tensor is passed. + + After decoration, we check the length of the inputs to ensure we have the + right number of input tensors. + + :param inputs: List of two tensors `(array, value)` to compute membership + over. + :returns: The tensor resulting from the membership operation, with `1.0` + where the value is found and `0.0` otherwise. + """ + if len(inputs) != 2: + raise ValueError(f"Expected 2 inputs, got {len(inputs)} inputs instead.") + + array, value = inputs + any_match = ops.any(ops.equal(array, value), axis=-1, keepdims=True) + return ops.cast(any_match, "float32") + + 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. + """ + return super().get_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..3fc2bfcf --- /dev/null +++ b/src/kamae/spark/transformers/array_contains.py @@ -0,0 +1,159 @@ +# 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 ( + ArrayType, + ByteType, + DataType, + DoubleType, + FloatType, + IntegerType, + LongType, + ShortType, +) + +from kamae.keras.core.backend import ALL_BACKENDS +from kamae.keras.core.layers import ArrayContainsLayer +from kamae.spark.params import MultiInputSingleOutputParams + +from .base import BaseTransformer + +_NUMERIC_TYPES = (ByteType, ShortType, IntegerType, LongType, FloatType, DoubleType) + + +class ArrayContainsTransformer( + BaseTransformer, + MultiInputSingleOutputParams, +): + """ + ArrayContainsLayer Spark Transformer for use in Spark pipelines. + + This transformer checks whether a scalar value is contained in an array. + + Input: Two columns `[arrayCol, valueCol]`, where `arrayCol` is an + `Array[Numeric]` and `valueCol` is a scalar `Numeric`. + Output: Scalar `Double` equal to `1.0` if the value is in the array, + else `0.0`. + """ + + supported_backends = ALL_BACKENDS + jit_compatible = True + + @keyword_only + def __init__( + self, + inputCols: Optional[List[str]] = None, + outputCol: Optional[str] = None, + inputDtype: Optional[str] = None, + outputDtype: Optional[str] = None, + layerName: Optional[str] = None, + ) -> None: + """ + Initializes an ArrayContainsTransformer transformer. + + :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. + :returns: None - class instantiated. + """ + super().__init__() + 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: + """ + Transforms the input dataset. Creates a new column with name `outputCol`, + equal to `1.0` if the value in `valueCol` is contained in the array in + `arrayCol`, else `0.0`. + + :param dataset: Pyspark dataframe to transform. + :returns: Transformed pyspark dataframe. + """ + arr_c, val_c = self.getInputCols() + arr_t = self.get_column_datatype(dataset, arr_c) + val_t = self.get_column_datatype(dataset, val_c) + + if not isinstance(arr_t, ArrayType): + raise TypeError(f"arrayCol '{arr_c}' must be an ArrayType, got {arr_t}") + + if not isinstance(elem_t := arr_t.elementType, _NUMERIC_TYPES): + raise TypeError(f"arrayCol '{arr_c}' element must be numeric, got {elem_t}") + + if not isinstance(val_t, _NUMERIC_TYPES): + raise TypeError(f"valueCol '{val_c}' must be numeric, got {val_t}") + + output_col = ( + F.when(F.array_contains(F.col(arr_c), F.col(val_c)), F.lit(1.0)) + .otherwise(F.lit(0.0)) + .cast(DoubleType()) + ) + 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(), + ) 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..ec580ee1 --- /dev/null +++ b/tests/kamae/keras/tensorflow/layers/test_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. + +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([[[1.0]]]), + ), + ( + [ + tf.constant([[[1, 2, 3]]]), + tf.constant([[[5]]]), + ], + "input_2", + None, + None, + tf.constant([[[0.0]]]), + ), + ( + [ + 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([[[1.0]], [[0.0]]]), + ), + ( + [ + tf.constant([["1", "2", "3"]]), + tf.constant([["2"]]), + ], + "input_5", + "int64", + None, + tf.constant([[1.0]]), + ), + ], + ) + 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_near(output_tensor, expected_output, atol=1e-6) + + @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..fa79d1e7 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,12 @@ {"axis": -2}, False, ), + ( + ArrayContainsLayer, + [tf.random.normal((32, 1, 10)), tf.random.normal((32, 5, 1))], + None, + 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..b705794e --- /dev/null +++ b/tests/kamae/spark/transformers/test_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. + +import numpy as np +import pytest +import tensorflow as tf + +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, 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"], + ) + + @pytest.fixture(scope="class") + def array_contains_transform_array_other_value_expected(self, spark_session): + return spark_session.createDataFrame( + [ + ([1, 2, 3], 2, 4, 0.0), + ([1, 2, 3], 5, 1, 1.0), + ([4, 5, 6], 4, 9, 0.0), + ], + [ + "array_col", + "value_col", + "other_value_col", + "array_contains_other_value", + ], + ) + + @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_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(TypeError): + 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, + )