feat: array contains

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)            |

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",

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()

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

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(),
+        )