From 5e9e377d015a4842dd1f09264c325e34f95af9c9 Mon Sep 17 00:00:00 2001 From: mantasu Date: Fri, 31 Jul 2026 13:11:54 +0100 Subject: [PATCH 1/7] feat: array contains --- README.md | 1 + src/kamae/keras/core/layers/__init__.py | 2 + src/kamae/keras/core/layers/array_contains.py | 113 +++++++++++ src/kamae/spark/transformers/__init__.py | 1 + .../spark/transformers/array_contains.py | 172 ++++++++++++++++ .../tensorflow/layers/test_array_contains.py | 137 +++++++++++++ tests/kamae/keras/test_jit_compatibility.py | 27 ++- tests/kamae/keras/test_layer_serialisation.py | 11 +- .../spark/transformers/test_array_contains.py | 188 ++++++++++++++++++ 9 files changed, 642 insertions(+), 10 deletions(-) create mode 100644 src/kamae/keras/core/layers/array_contains.py create mode 100644 src/kamae/spark/transformers/array_contains.py create mode 100644 tests/kamae/keras/tensorflow/layers/test_array_contains.py create mode 100644 tests/kamae/spark/transformers/test_array_contains.py 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..0bc9a790 --- /dev/null +++ b/src/kamae/keras/core/layers/array_contains.py @@ -0,0 +1,113 @@ +# 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, received {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..1f2e6d9d --- /dev/null +++ b/src/kamae/spark/transformers/array_contains.py @@ -0,0 +1,172 @@ +# 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 columns, 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. + """ + array_col_name, value_col_name = self.getInputCols() + + array_type = self.get_column_datatype(dataset, array_col_name) + value_type = self.get_column_datatype(dataset, value_col_name) + + if not isinstance(array_type, ArrayType): + raise TypeError( + f"arrayCol '{array_col_name}' must be an ArrayType, got {array_type}" + ) + + elem_type = array_type.elementType + if not isinstance(elem_type, _NUMERIC_TYPES): + raise TypeError( + f"arrayCol '{array_col_name}' element type must be numeric, " + f"got {elem_type}" + ) + + if not isinstance(value_type, _NUMERIC_TYPES): + raise TypeError( + f"valueCol '{value_col_name}' must be numeric, got {value_type}" + ) + + output_col = ( + F.when( + F.array_contains(F.col(array_col_name), F.col(value_col_name)), + 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..cfca480b 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}), ( @@ -530,15 +536,18 @@ def test_all_layers_define_jit_compatible_and_supported_backends(): all_layers = multi_backend_layers + tf_only_layers for layer_cls in all_layers: - assert ( - "jit_compatible" in layer_cls.__dict__ - ), f"{layer_cls.__name__} must define 'jit_compatible' directly (not inherit it)" - assert isinstance( - layer_cls.jit_compatible, bool - ), f"{layer_cls.__name__}.jit_compatible must be bool, got {type(layer_cls.jit_compatible)}" - assert ( - "supported_backends" in layer_cls.__dict__ - ), f"{layer_cls.__name__} must define 'supported_backends' directly (not inherit it)" + assert "jit_compatible" in layer_cls.__dict__, ( + f"{layer_cls.__name__} must define 'jit_compatible' directly (not" + " inherit it)" + ) + assert isinstance(layer_cls.jit_compatible, bool), ( + f"{layer_cls.__name__}.jit_compatible must be bool, got" + f" {type(layer_cls.jit_compatible)}" + ) + assert "supported_backends" in layer_cls.__dict__, ( + f"{layer_cls.__name__} must define 'supported_backends' directly (not" + " inherit it)" + ) assert isinstance( layer_cls.supported_backends, frozenset ), f"{layer_cls.__name__}.supported_backends must be frozenset" diff --git a/tests/kamae/keras/test_layer_serialisation.py b/tests/kamae/keras/test_layer_serialisation.py index e634e7b3..cb7e8dc0 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))], @@ -560,7 +567,9 @@ def test_layer_serialisation( # LambdaFunctionLayer cannot serialize/deserialize lambda functions that reference # external modules (like tf) - this is a fundamental limitation of Python lambda serialization pytest.skip( - reason="LambdaFunctionLayer with module references cannot serialize in Keras 3" + reason=( + "LambdaFunctionLayer with module references cannot serialize in Keras 3" + ) ) if kwargs is None: kwargs = {} 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, + ) From bae9cdcaa1fe296cf414a72d3f946633261fc6de Mon Sep 17 00:00:00 2001 From: mantasu Date: Fri, 31 Jul 2026 13:41:31 +0100 Subject: [PATCH 2/7] fix: readability --- src/kamae/keras/core/layers/array_contains.py | 5 +-- .../spark/transformers/array_contains.py | 37 ++++++------------- tests/kamae/keras/test_jit_compatibility.py | 21 +++++------ 3 files changed, 23 insertions(+), 40 deletions(-) diff --git a/src/kamae/keras/core/layers/array_contains.py b/src/kamae/keras/core/layers/array_contains.py index 0bc9a790..243502b5 100644 --- a/src/kamae/keras/core/layers/array_contains.py +++ b/src/kamae/keras/core/layers/array_contains.py @@ -96,9 +96,8 @@ def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: where the value is found and `0.0` otherwise. """ if len(inputs) != 2: - raise ValueError( - f"Expected 2 inputs, received {len(inputs)} inputs instead." - ) + 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") diff --git a/src/kamae/spark/transformers/array_contains.py b/src/kamae/spark/transformers/array_contains.py index 1f2e6d9d..3fc2bfcf 100644 --- a/src/kamae/spark/transformers/array_contains.py +++ b/src/kamae/spark/transformers/array_contains.py @@ -112,9 +112,8 @@ def setInputCols(self, value: List[str]) -> "ArrayContainsTransformer": :returns: Instance of class with input columns set. """ if len(value) != 2: - raise ValueError( - f"Expected 2 input columns, received {len(value)} instead." - ) + raise ValueError(f"Expected 2 input cols, received {len(value)} instead.") + return self._set(inputCols=value) def _transform(self, dataset: DataFrame) -> DataFrame: @@ -126,33 +125,21 @@ def _transform(self, dataset: DataFrame) -> DataFrame: :param dataset: Pyspark dataframe to transform. :returns: Transformed pyspark dataframe. """ - array_col_name, value_col_name = self.getInputCols() - - array_type = self.get_column_datatype(dataset, array_col_name) - value_type = self.get_column_datatype(dataset, value_col_name) + 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(array_type, ArrayType): - raise TypeError( - f"arrayCol '{array_col_name}' must be an ArrayType, got {array_type}" - ) + if not isinstance(arr_t, ArrayType): + raise TypeError(f"arrayCol '{arr_c}' must be an ArrayType, got {arr_t}") - elem_type = array_type.elementType - if not isinstance(elem_type, _NUMERIC_TYPES): - raise TypeError( - f"arrayCol '{array_col_name}' element type must be numeric, " - f"got {elem_type}" - ) + 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(value_type, _NUMERIC_TYPES): - raise TypeError( - f"valueCol '{value_col_name}' must be numeric, got {value_type}" - ) + 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(array_col_name), F.col(value_col_name)), - F.lit(1.0), - ) + F.when(F.array_contains(F.col(arr_c), F.col(val_c)), F.lit(1.0)) .otherwise(F.lit(0.0)) .cast(DoubleType()) ) diff --git a/tests/kamae/keras/test_jit_compatibility.py b/tests/kamae/keras/test_jit_compatibility.py index cfca480b..2ec2d600 100644 --- a/tests/kamae/keras/test_jit_compatibility.py +++ b/tests/kamae/keras/test_jit_compatibility.py @@ -536,18 +536,15 @@ def test_all_layers_define_jit_compatible_and_supported_backends(): all_layers = multi_backend_layers + tf_only_layers for layer_cls in all_layers: - assert "jit_compatible" in layer_cls.__dict__, ( - f"{layer_cls.__name__} must define 'jit_compatible' directly (not" - " inherit it)" - ) - assert isinstance(layer_cls.jit_compatible, bool), ( - f"{layer_cls.__name__}.jit_compatible must be bool, got" - f" {type(layer_cls.jit_compatible)}" - ) - assert "supported_backends" in layer_cls.__dict__, ( - f"{layer_cls.__name__} must define 'supported_backends' directly (not" - " inherit it)" - ) + assert ( + "jit_compatible" in layer_cls.__dict__ + ), f"{layer_cls.__name__} must define 'jit_compatible' directly (not inherit it)" + assert isinstance( + layer_cls.jit_compatible, bool + ), f"{layer_cls.__name__}.jit_compatible must be bool, got {type(layer_cls.jit_compatible)}" + assert ( + "supported_backends" in layer_cls.__dict__ + ), f"{layer_cls.__name__} must define 'supported_backends' directly (not inherit it)" assert isinstance( layer_cls.supported_backends, frozenset ), f"{layer_cls.__name__}.supported_backends must be frozenset" From 7d3af7db2540b9a399c7a621572fbb20d2e25e58 Mon Sep 17 00:00:00 2001 From: mantasu Date: Fri, 31 Jul 2026 13:44:59 +0100 Subject: [PATCH 3/7] fi: undo formatting --- tests/kamae/keras/test_layer_serialisation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/kamae/keras/test_layer_serialisation.py b/tests/kamae/keras/test_layer_serialisation.py index cb7e8dc0..fa79d1e7 100644 --- a/tests/kamae/keras/test_layer_serialisation.py +++ b/tests/kamae/keras/test_layer_serialisation.py @@ -567,9 +567,7 @@ def test_layer_serialisation( # LambdaFunctionLayer cannot serialize/deserialize lambda functions that reference # external modules (like tf) - this is a fundamental limitation of Python lambda serialization pytest.skip( - reason=( - "LambdaFunctionLayer with module references cannot serialize in Keras 3" - ) + reason="LambdaFunctionLayer with module references cannot serialize in Keras 3" ) if kwargs is None: kwargs = {} From 477d42c1db7bcb6ca49f5c2c0285ef1ab4d33cf6 Mon Sep 17 00:00:00 2001 From: mantasu Date: Thu, 3 Sep 2026 08:47:29 +0100 Subject: [PATCH 4/7] fix: apply comments --- src/kamae/keras/core/layers/array_contains.py | 36 ++++---- .../spark/transformers/array_contains.py | 33 ++++--- .../tensorflow/layers/test_array_contains.py | 85 +++++++++++++++++-- tests/kamae/keras/test_layer_serialisation.py | 2 +- .../spark/transformers/test_array_contains.py | 78 +++++++++++++++-- 5 files changed, 191 insertions(+), 43 deletions(-) diff --git a/src/kamae/keras/core/layers/array_contains.py b/src/kamae/keras/core/layers/array_contains.py index 243502b5..5da52c0b 100644 --- a/src/kamae/keras/core/layers/array_contains.py +++ b/src/kamae/keras/core/layers/array_contains.py @@ -26,15 +26,11 @@ @keras.saving.register_keras_serializable(package=kamae.__name__) class ArrayContainsLayer(BaseLayer): """ - Computes whether a value is contained in an array along the last axis. + Computes whether a value is contained in an array along a given 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. + search axis, and returns a boolean where `value` is found in `array`. Set + `output_dtype` to cast the result (e.g. to a float). """ supported_backends = ALL_BACKENDS @@ -45,6 +41,8 @@ def __init__( name: Optional[str] = None, input_dtype: Optional[str] = None, output_dtype: Optional[str] = None, + axis: int = -1, + keepdims: bool = True, **kwargs: Any, ) -> None: """ @@ -53,10 +51,16 @@ def __init__( :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 search for the value. Defaults to + `-1`. + :param keepdims: Whether to keep the searched axis as a size 1 dimension + in the output. Defaults to `True`. """ super().__init__( name=name, input_dtype=input_dtype, output_dtype=output_dtype, **kwargs ) + self.axis = axis + self.keepdims = keepdims @property def compatible_dtypes(self) -> Optional[List[str]]: @@ -82,25 +86,17 @@ def compatible_dtypes(self) -> Optional[List[str]]: @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. + Computes membership of `value` within `array` along the search axis. :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. + :returns: A boolean tensor, `True` where the value is found. """ 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") + return ops.any(ops.equal(array, value), axis=self.axis, keepdims=self.keepdims) def get_config(self) -> Dict[str, Any]: """ @@ -109,4 +105,6 @@ def get_config(self) -> Dict[str, Any]: :returns: Dictionary of the configuration of the layer. """ - return super().get_config() + config = super().get_config() + config.update({"axis": self.axis, "keepdims": self.keepdims}) + return config diff --git a/src/kamae/spark/transformers/array_contains.py b/src/kamae/spark/transformers/array_contains.py index 3fc2bfcf..81438ad4 100644 --- a/src/kamae/spark/transformers/array_contains.py +++ b/src/kamae/spark/transformers/array_contains.py @@ -36,6 +36,10 @@ from kamae.keras.core.backend import ALL_BACKENDS from kamae.keras.core.layers import ArrayContainsLayer from kamae.spark.params import MultiInputSingleOutputParams +from kamae.spark.utils import ( + get_array_nesting_level_and_element_dtype, + nested_transform, +) from .base import BaseTransformer @@ -51,10 +55,11 @@ class ArrayContainsTransformer( 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`. + Input: Two columns `[arrayCol, valueCol]`, where `arrayCol` is a + (possibly nested) `Array[Numeric]` and `valueCol` is a scalar `Numeric`. + Output: `Boolean` (or nested `Array[Boolean]` for nested inputs) equal to + `True` if the value is in the innermost array, else `False`. Set + `outputDtype` to cast the boolean result to another dtype. """ supported_backends = ALL_BACKENDS @@ -118,9 +123,9 @@ def setInputCols(self, value: List[str]) -> "ArrayContainsTransformer": 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`. + Transforms the input dataset. Creates a new column `outputCol` that is + `True` where `valueCol` is contained in the innermost array of `arrayCol`. + Nested arrays are supported, yielding a nested array of booleans. :param dataset: Pyspark dataframe to transform. :returns: Transformed pyspark dataframe. @@ -132,17 +137,19 @@ def _transform(self, dataset: DataFrame) -> DataFrame: 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): + nesting_level, elem_t = get_array_nesting_level_and_element_dtype(arr_t) + if not isinstance(elem_t, _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()) + # Apply array_contains at the innermost level + contains_func = nested_transform( + func=lambda x: F.array_contains(x, F.col(val_c)), + nest_level=nesting_level - 1, ) + output_col = contains_func(F.col(arr_c)) return dataset.withColumn(self.getOutputCol(), output_col) def get_keras_layer(self) -> keras.layers.Layer: @@ -156,4 +163,6 @@ def get_keras_layer(self) -> keras.layers.Layer: name=self.getLayerName(), input_dtype=self.getInputKerasDtype(), output_dtype=self.getOutputKerasDtype(), + axis=-1, + keepdims=True, ) diff --git a/tests/kamae/keras/tensorflow/layers/test_array_contains.py b/tests/kamae/keras/tensorflow/layers/test_array_contains.py index ec580ee1..edd43b1b 100644 --- a/tests/kamae/keras/tensorflow/layers/test_array_contains.py +++ b/tests/kamae/keras/tensorflow/layers/test_array_contains.py @@ -30,7 +30,7 @@ class TestArrayContains: "input_1", None, None, - tf.constant([[[1.0]]]), + tf.constant([[[True]]]), ), ( [ @@ -40,7 +40,7 @@ class TestArrayContains: "input_2", None, None, - tf.constant([[[0.0]]]), + tf.constant([[[False]]]), ), ( [ @@ -70,7 +70,7 @@ class TestArrayContains: "input_4", None, None, - tf.constant([[[1.0]], [[0.0]]]), + tf.constant([[[True]], [[False]]]), ), ( [ @@ -80,7 +80,29 @@ class TestArrayContains: "input_5", "int64", None, - tf.constant([[1.0]]), + 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"), ), ], ) @@ -108,7 +130,60 @@ def test_array_contains( 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) + 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_tensors", diff --git a/tests/kamae/keras/test_layer_serialisation.py b/tests/kamae/keras/test_layer_serialisation.py index fa79d1e7..d659d084 100644 --- a/tests/kamae/keras/test_layer_serialisation.py +++ b/tests/kamae/keras/test_layer_serialisation.py @@ -122,7 +122,7 @@ ( ArrayContainsLayer, [tf.random.normal((32, 1, 10)), tf.random.normal((32, 5, 1))], - None, + {"axis": -1, "keepdims": True}, False, ), ( diff --git a/tests/kamae/spark/transformers/test_array_contains.py b/tests/kamae/spark/transformers/test_array_contains.py index b705794e..db25d222 100644 --- a/tests/kamae/spark/transformers/test_array_contains.py +++ b/tests/kamae/spark/transformers/test_array_contains.py @@ -15,6 +15,7 @@ import numpy as np import pytest import tensorflow as tf +from pyspark.sql.types import DoubleType from kamae.spark.transformers import ArrayContainsTransformer @@ -35,9 +36,9 @@ def example_dataframe_with_arrays(self, spark_session): 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), + ([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"], ) @@ -46,9 +47,9 @@ def array_contains_transform_array_value_expected(self, spark_session): 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), + ([1, 2, 3], 2, 4, False), + ([1, 2, 3], 5, 1, True), + ([4, 5, 6], 4, 9, False), ], [ "array_col", @@ -58,6 +59,28 @@ def array_contains_transform_array_other_value_expected(self, spark_session): ], ) + @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", [ @@ -93,6 +116,49 @@ def test_spark_array_contains_transform( 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_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() From eaf66ea222c1d6eaabe307143ceb2eedbf0d0dd9 Mon Sep 17 00:00:00 2001 From: mantasu Date: Wed, 9 Sep 2026 15:37:00 +0100 Subject: [PATCH 5/7] fix: apply comments --- src/kamae/keras/core/layers/array_contains.py | 65 +++++++--- .../spark/transformers/array_contains.py | 81 ++++++------ .../tensorflow/layers/test_array_contains.py | 58 +++++++++ tests/kamae/keras/test_layer_serialisation.py | 6 + .../spark/transformers/test_array_contains.py | 115 +++++++++++++++++- 5 files changed, 270 insertions(+), 55 deletions(-) diff --git a/src/kamae/keras/core/layers/array_contains.py b/src/kamae/keras/core/layers/array_contains.py index 5da52c0b..a83ba1ed 100644 --- a/src/kamae/keras/core/layers/array_contains.py +++ b/src/kamae/keras/core/layers/array_contains.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Iterable, List, Optional +from typing import Any, Dict, Iterable, List, Optional, Union import keras from keras import KerasTensor, ops @@ -20,7 +20,7 @@ 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 +from kamae.keras.core.utils.input_utils import allow_single_or_multiple_tensor_input @keras.saving.register_keras_serializable(package=kamae.__name__) @@ -28,9 +28,14 @@ class ArrayContainsLayer(BaseLayer): """ Computes whether a value is contained in an array along a given axis. - Expects two inputs `(array, value)` that broadcast on every axis except the - search axis, and returns a boolean where `value` is found in `array`. Set - `output_dtype` to cast the result (e.g. to a float). + 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 @@ -43,6 +48,7 @@ def __init__( output_dtype: Optional[str] = None, axis: int = -1, keepdims: bool = True, + value_constant: Optional[float] = None, **kwargs: Any, ) -> None: """ @@ -51,16 +57,20 @@ def __init__( :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 search for the value. Defaults to + :param axis: The axis along which to check for the value. Defaults to `-1`. - :param keepdims: Whether to keep the searched axis as a size 1 dimension + :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]]: @@ -70,6 +80,10 @@ def compatible_dtypes(self) -> Optional[List[str]]: :returns: List of compatible dtype names. """ return [ + "bfloat16", + "float16", + "float32", + "float64", "int8", "uint8", "int16", @@ -78,24 +92,31 @@ def compatible_dtypes(self) -> Optional[List[str]]: "uint32", "int64", "uint64", - "float16", - "float32", - "float64", ] - @enforce_multiple_tensor_input - def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: + @allow_single_or_multiple_tensor_input + def _call( + self, inputs: Union[KerasTensor, Iterable[KerasTensor]], **kwargs: Any + ) -> KerasTensor: """ - Computes membership of `value` within `array` along the search axis. + Computes whether `value` is present in `array` along the given axis. - :param inputs: List of two tensors `(array, value)` to compute membership - over. + :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 len(inputs) != 2: - raise ValueError(f"Expected 2 inputs, got {len(inputs)} inputs instead.") + 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") - array, value = inputs return ops.any(ops.equal(array, value), axis=self.axis, keepdims=self.keepdims) def get_config(self) -> Dict[str, Any]: @@ -106,5 +127,11 @@ def get_config(self) -> Dict[str, Any]: :returns: Dictionary of the configuration of the layer. """ config = super().get_config() - config.update({"axis": self.axis, "keepdims": self.keepdims}) + config.update( + { + "axis": self.axis, + "keepdims": self.keepdims, + "value_constant": self.value_constant, + } + ) return config diff --git a/src/kamae/spark/transformers/array_contains.py b/src/kamae/spark/transformers/array_contains.py index 81438ad4..364937e1 100644 --- a/src/kamae/spark/transformers/array_contains.py +++ b/src/kamae/spark/transformers/array_contains.py @@ -23,7 +23,6 @@ from pyspark import keyword_only from pyspark.sql import DataFrame from pyspark.sql.types import ( - ArrayType, ByteType, DataType, DoubleType, @@ -35,31 +34,47 @@ from kamae.keras.core.backend import ALL_BACKENDS from kamae.keras.core.layers import ArrayContainsLayer -from kamae.spark.params import MultiInputSingleOutputParams -from kamae.spark.utils import ( - get_array_nesting_level_and_element_dtype, - nested_transform, +from kamae.spark.params import ( + MathFloatConstantParams, + MultiInputSingleOutputParams, + SingleInputSingleOutputParams, ) +from kamae.spark.utils import single_input_single_output_array_transform from .base import BaseTransformer -_NUMERIC_TYPES = (ByteType, ShortType, IntegerType, LongType, FloatType, DoubleType) - class ArrayContainsTransformer( BaseTransformer, + SingleInputSingleOutputParams, MultiInputSingleOutputParams, + MathFloatConstantParams, ): """ 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 a - (possibly nested) `Array[Numeric]` and `valueCol` is a scalar `Numeric`. - Output: `Boolean` (or nested `Array[Boolean]` for nested inputs) equal to - `True` if the value is in the innermost array, else `False`. Set - `outputDtype` to cast the boolean result to another dtype. + 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 @@ -68,15 +83,20 @@ class ArrayContainsTransformer( @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 @@ -85,9 +105,12 @@ def __init__( 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) @@ -123,33 +146,20 @@ def setInputCols(self, value: List[str]) -> "ArrayContainsTransformer": def _transform(self, dataset: DataFrame) -> DataFrame: """ - Transforms the input dataset. Creates a new column `outputCol` that is - `True` where `valueCol` is contained in the innermost array of `arrayCol`. - Nested arrays are supported, yielding a nested array of booleans. + 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. """ - 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}") - - nesting_level, elem_t = get_array_nesting_level_and_element_dtype(arr_t) - if not isinstance(elem_t, _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}") + array_col, value_col = self.get_multiple_input_cols("mathFloatConstant", 2) + df = dataset.select(array_col, value_col) - # Apply array_contains at the innermost level - contains_func = nested_transform( - func=lambda x: F.array_contains(x, F.col(val_c)), - nest_level=nesting_level - 1, + output_col = single_input_single_output_array_transform( + input_col=array_col, + input_col_datatype=self.get_column_datatype(df, df.columns[0]), + func=lambda x: F.array_contains(x, value_col), ) - output_col = contains_func(F.col(arr_c)) return dataset.withColumn(self.getOutputCol(), output_col) def get_keras_layer(self) -> keras.layers.Layer: @@ -165,4 +175,5 @@ def get_keras_layer(self) -> keras.layers.Layer: 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 index edd43b1b..1e7c8d50 100644 --- a/tests/kamae/keras/tensorflow/layers/test_array_contains.py +++ b/tests/kamae/keras/tensorflow/layers/test_array_contains.py @@ -185,6 +185,64 @@ def test_array_contains_axis_keepdims( ), "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", [ diff --git a/tests/kamae/keras/test_layer_serialisation.py b/tests/kamae/keras/test_layer_serialisation.py index d659d084..a4679777 100644 --- a/tests/kamae/keras/test_layer_serialisation.py +++ b/tests/kamae/keras/test_layer_serialisation.py @@ -125,6 +125,12 @@ {"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 index db25d222..3afc7c33 100644 --- a/tests/kamae/spark/transformers/test_array_contains.py +++ b/tests/kamae/spark/transformers/test_array_contains.py @@ -192,7 +192,65 @@ def test_array_contains_transform_non_array_input_raises_error( outputCol="array_contains_output", ) # then - with pytest.raises(TypeError): + 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( @@ -252,3 +310,58 @@ def test_array_contains_transform_spark_tf_parity( 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", + ), + ], + ) + 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, + ) From 13edfb3420c58e6b074637060018a0d736706845 Mon Sep 17 00:00:00 2001 From: mantasu Date: Mon, 28 Sep 2026 16:00:43 +0100 Subject: [PATCH 6/7] fix: apply comments --- .../spark/transformers/array_contains.py | 2 +- .../spark/transformers/test_array_contains.py | 33 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/src/kamae/spark/transformers/array_contains.py b/src/kamae/spark/transformers/array_contains.py index 364937e1..3adc253e 100644 --- a/src/kamae/spark/transformers/array_contains.py +++ b/src/kamae/spark/transformers/array_contains.py @@ -158,7 +158,7 @@ def _transform(self, dataset: DataFrame) -> DataFrame: output_col = single_input_single_output_array_transform( input_col=array_col, input_col_datatype=self.get_column_datatype(df, df.columns[0]), - func=lambda x: F.array_contains(x, value_col), + func=lambda x: F.coalesce(F.array_contains(x, value_col), F.lit(False)), ) return dataset.withColumn(self.getOutputCol(), output_col) diff --git a/tests/kamae/spark/transformers/test_array_contains.py b/tests/kamae/spark/transformers/test_array_contains.py index 3afc7c33..2054c353 100644 --- a/tests/kamae/spark/transformers/test_array_contains.py +++ b/tests/kamae/spark/transformers/test_array_contains.py @@ -133,6 +133,39 @@ def test_spark_array_contains_transform_nested_arrays( 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 ): From 6dc9372732ed85987f116360f9cf5d35a5db2e15 Mon Sep 17 00:00:00 2001 From: mantasu Date: Mon, 28 Sep 2026 21:11:42 +0100 Subject: [PATCH 7/7] fix: spark dtypes --- .../spark/transformers/array_contains.py | 13 ++- .../spark/transformers/test_array_contains.py | 90 ++++++++++++++++++- 2 files changed, 100 insertions(+), 3 deletions(-) diff --git a/src/kamae/spark/transformers/array_contains.py b/src/kamae/spark/transformers/array_contains.py index 3adc253e..6e818e84 100644 --- a/src/kamae/spark/transformers/array_contains.py +++ b/src/kamae/spark/transformers/array_contains.py @@ -27,6 +27,7 @@ DataType, DoubleType, FloatType, + FractionalType, IntegerType, LongType, ShortType, @@ -39,7 +40,10 @@ MultiInputSingleOutputParams, SingleInputSingleOutputParams, ) -from kamae.spark.utils import single_input_single_output_array_transform +from kamae.spark.utils import ( + get_element_type, + single_input_single_output_array_transform, +) from .base import BaseTransformer @@ -154,10 +158,15 @@ def _transform(self, dataset: DataFrame) -> 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=self.get_column_datatype(df, df.columns[0]), + 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) diff --git a/tests/kamae/spark/transformers/test_array_contains.py b/tests/kamae/spark/transformers/test_array_contains.py index 2054c353..dd5c791e 100644 --- a/tests/kamae/spark/transformers/test_array_contains.py +++ b/tests/kamae/spark/transformers/test_array_contains.py @@ -15,7 +15,7 @@ import numpy as np import pytest import tensorflow as tf -from pyspark.sql.types import DoubleType +from pyspark.sql.types import ArrayType, DoubleType, FloatType, StructField, StructType from kamae.spark.transformers import ArrayContainsTransformer @@ -359,6 +359,22 @@ def test_array_contains_transform_spark_tf_parity( "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( @@ -398,3 +414,75 @@ def test_array_contains_transform_constant_spark_tf_parity( 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, + )