Repository navigation
feat: array contains #67
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,137 @@ | ||
| # Copyright [2024] Expedia, Inc. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| from typing import Any, Dict, Iterable, List, Optional, Union | ||
|
|
||
| import keras | ||
| from keras import KerasTensor, ops | ||
|
|
||
| import kamae | ||
| from kamae.keras.core.backend import ALL_BACKENDS | ||
| from kamae.keras.core.base import BaseLayer | ||
| from kamae.keras.core.utils.input_utils import allow_single_or_multiple_tensor_input | ||
|
|
||
|
|
||
| @keras.saving.register_keras_serializable(package=kamae.__name__) | ||
| class ArrayContainsLayer(BaseLayer): | ||
| """ | ||
| Computes whether a value is contained in an array along a given axis. | ||
|
|
||
| With `value_constant` set, takes a single `array` input and checks whether | ||
| that constant is present. Otherwise takes two inputs `(array, value)`. | ||
| Returns a boolean; set `output_dtype` to cast the result. | ||
|
|
||
| Example: | ||
| >>> layer = ArrayContainsLayer(value_constant=2) | ||
| >>> layer([[1, 2, 3], [4, 5, 6]]) | ||
| [[True], [False]] | ||
| """ | ||
|
|
||
| supported_backends = ALL_BACKENDS | ||
| jit_compatible = True | ||
|
|
||
| def __init__( | ||
| self, | ||
| name: Optional[str] = None, | ||
| input_dtype: Optional[str] = None, | ||
| output_dtype: Optional[str] = None, | ||
| axis: int = -1, | ||
| keepdims: bool = True, | ||
| value_constant: Optional[float] = None, | ||
| **kwargs: Any, | ||
| ) -> None: | ||
| """ | ||
| Initializes the ArrayContainsLayer layer. | ||
|
|
||
| :param name: Name of the layer, defaults to `None`. | ||
| :param input_dtype: The dtype to cast the input to. Defaults to `None`. | ||
| :param output_dtype: The dtype to cast the output to. Defaults to `None`. | ||
| :param axis: The axis along which to check for the value. Defaults to | ||
| `-1`. | ||
| :param keepdims: Whether to keep the reduced axis as a size 1 dimension | ||
| in the output. Defaults to `True`. | ||
| :param value_constant: Optional constant to check for. If set, the layer | ||
| takes a single `array` input. If `None`, two inputs `(array, value)` are | ||
| required. Defaults to `None`. | ||
| """ | ||
| super().__init__( | ||
| name=name, input_dtype=input_dtype, output_dtype=output_dtype, **kwargs | ||
| ) | ||
| self.axis = axis | ||
| self.keepdims = keepdims | ||
| self.value_constant = value_constant | ||
|
|
||
| @property | ||
| def compatible_dtypes(self) -> Optional[List[str]]: | ||
| """ | ||
| Returns the compatible dtypes of the layer. | ||
|
|
||
| :returns: List of compatible dtype names. | ||
| """ | ||
| return [ | ||
| "bfloat16", | ||
| "float16", | ||
| "float32", | ||
| "float64", | ||
| "int8", | ||
| "uint8", | ||
| "int16", | ||
| "uint16", | ||
| "int32", | ||
| "uint32", | ||
| "int64", | ||
| "uint64", | ||
| ] | ||
|
|
||
| @allow_single_or_multiple_tensor_input | ||
| def _call( | ||
| self, inputs: Union[KerasTensor, Iterable[KerasTensor]], **kwargs: Any | ||
| ) -> KerasTensor: | ||
| """ | ||
| Computes whether `value` is present in `array` along the given axis. | ||
|
|
||
| :param inputs: Single `array` tensor (when `value_constant` is set) or a | ||
| list of two tensors `(array, value)`. | ||
| :returns: A boolean tensor, `True` where the value is found. | ||
| """ | ||
| if self.value_constant is not None and len(inputs) == 1: | ||
| # Constant value given upon initialization | ||
| array, value = inputs[0], self.value_constant | ||
| array, value = self._force_cast_to_compatible_numeric_type(array, value) | ||
| elif self.value_constant is None and len(inputs) == 2: | ||
| # Dynamic value given | ||
| array, value = inputs | ||
| elif self.value_constant is not None: | ||
| raise ValueError("Expected 1 input when `value_constant` is set") | ||
| else: | ||
| raise ValueError(f"Expected 2 inputs, got {len(inputs)} inputs instead") | ||
|
|
||
| return ops.any(ops.equal(array, value), axis=self.axis, keepdims=self.keepdims) | ||
|
There was a problem hiding this comment. Is the equality check here identical to what array_contains does on the Spark side? Just trying to understand whether there's any chance of a mismatch on exact equality. I see above (line 111) the cast only runs for the constant branch, so I wasn't sure what happens when the value comes in as a separate column with a different dtype.
Member
Author
There was a problem hiding this comment. Both use exact equality and There was a problem hiding this comment. Thanks Mantas! The inputDtype point holds when it's defined, but it defaults to None, and then nothing aligns the array and the constant, so each side reconciles float vs double its own way: Spark widens the array to double, and Keras on the opposite direction narrows the constant to float32. It might be pedantic, since this only matters for constants that aren't exactly representable in float32, like 0.1: array_contains([0.1f], 0.1) → false in Spark (0.10000000149 ≠ 0.1), but true in Keras. (I checked the Spark side locally with claude.) Could we cast the literal to the array's element type on the Spark side (or mirror Spark's widening in Keras), and add a parity test with inputDtype=None and a constant like 0.1 (that float32 is not enough to describe it) on a float array? The current parity tests only use integer arrays and integer constants (5.0 and 0.0), so the float path isn't covered.
Member
Author
There was a problem hiding this comment. Nice, thank you! Made it so that behaviour is identical to #L111 |
||
|
|
||
| def get_config(self) -> Dict[str, Any]: | ||
| """ | ||
| Gets the configuration of the ArrayContains layer. | ||
| Used for saving and loading from a model. | ||
|
|
||
| :returns: Dictionary of the configuration of the layer. | ||
| """ | ||
| config = super().get_config() | ||
| config.update( | ||
| { | ||
| "axis": self.axis, | ||
| "keepdims": self.keepdims, | ||
| "value_constant": self.value_constant, | ||
| } | ||
| ) | ||
| return config | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,188 @@ | ||
| # Copyright [2024] Expedia, Inc. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| # pylint: disable=unused-argument | ||
| # pylint: disable=invalid-name | ||
| # pylint: disable=too-many-ancestors | ||
| # pylint: disable=no-member | ||
| from typing import List, Optional | ||
|
|
||
| import keras | ||
| import pyspark.sql.functions as F | ||
| from pyspark import keyword_only | ||
| from pyspark.sql import DataFrame | ||
| from pyspark.sql.types import ( | ||
| ByteType, | ||
| DataType, | ||
| DoubleType, | ||
| FloatType, | ||
| FractionalType, | ||
| IntegerType, | ||
| LongType, | ||
| ShortType, | ||
| ) | ||
|
|
||
| from kamae.keras.core.backend import ALL_BACKENDS | ||
| from kamae.keras.core.layers import ArrayContainsLayer | ||
| from kamae.spark.params import ( | ||
| MathFloatConstantParams, | ||
| MultiInputSingleOutputParams, | ||
| SingleInputSingleOutputParams, | ||
| ) | ||
| from kamae.spark.utils import ( | ||
| get_element_type, | ||
| single_input_single_output_array_transform, | ||
| ) | ||
|
|
||
| from .base import BaseTransformer | ||
|
|
||
|
|
||
| class ArrayContainsTransformer( | ||
| BaseTransformer, | ||
| SingleInputSingleOutputParams, | ||
| MultiInputSingleOutputParams, | ||
| MathFloatConstantParams, | ||
| ): | ||
| """ | ||
| ArrayContainsLayer Spark Transformer for use in Spark pipelines. | ||
|
|
||
| Checks whether a scalar value is present in a (possibly nested) numeric | ||
| array. The value is either a second input column (`inputCols`) or the | ||
| `mathFloatConstant` (`inputCol`). Outputs a boolean; set `outputDtype` to | ||
| cast the result. | ||
|
|
||
| Example: | ||
|
|
||
| >>> df.show() | ||
| +---------+ | ||
| | a| | ||
| +---------+ | ||
| |[1, 2, 3]| | ||
| |[4, 5, 6]| | ||
| +---------+ | ||
| >>> t = ArrayContainsTransformer(inputCol="a", outputCol="b", mathFloatConstant=2) | ||
| >>> t.transform(df).show() | ||
| +---------+-----+ | ||
| | a| b| | ||
| +---------+-----+ | ||
| |[1, 2, 3]| true| | ||
| |[4, 5, 6]|false| | ||
| +---------+-----+ | ||
| """ | ||
|
|
||
| supported_backends = ALL_BACKENDS | ||
| jit_compatible = True | ||
|
|
||
| @keyword_only | ||
| def __init__( | ||
| self, | ||
| inputCol: Optional[str] = None, | ||
| inputCols: Optional[List[str]] = None, | ||
| outputCol: Optional[str] = None, | ||
| inputDtype: Optional[str] = None, | ||
| outputDtype: Optional[str] = None, | ||
| layerName: Optional[str] = None, | ||
| mathFloatConstant: Optional[float] = None, | ||
| ) -> None: | ||
| """ | ||
| Initializes an ArrayContainsTransformer transformer. | ||
|
|
||
| :param inputCol: Input array column name. Only used if inputCols is not | ||
| specified. If specified, we check whether `mathFloatConstant` is contained | ||
| in this array column. | ||
| :param inputCols: Input column names, given as `[arrayCol, valueCol]`. | ||
|
Comment on lines
+91
to
+104
Contributor
There was a problem hiding this comment. Do we want to allow a single input col and a constant? So we can check if
Member
Author
There was a problem hiding this comment. Fair idea, generalised it |
||
| :param outputCol: Output column name. | ||
| :param inputDtype: Input data type to cast input column(s) to before | ||
| transforming. | ||
| :param outputDtype: Output data type to cast the output column to after | ||
| transforming. | ||
| :param layerName: Name of the layer. Used as the name of the Keras layer | ||
| in the keras model. If not set, we use the uid of the Spark transformer. | ||
| :param mathFloatConstant: Optional constant value to check for. Used with | ||
| `inputCol`. If not provided, then `inputCols` is required. | ||
| :returns: None - class instantiated. | ||
| """ | ||
| super().__init__() | ||
| self._setDefault(mathFloatConstant=None) | ||
| kwargs = self._input_kwargs | ||
| self.setParams(**kwargs) | ||
|
|
||
| @property | ||
| def compatible_dtypes(self) -> Optional[List[DataType]]: | ||
| """ | ||
| List of compatible data types for the layer. | ||
| If the computation can be performed on any data type, return None. | ||
|
|
||
| :returns: List of compatible data types for the layer. | ||
| """ | ||
| return [ | ||
| FloatType(), | ||
| DoubleType(), | ||
| ByteType(), | ||
| ShortType(), | ||
| IntegerType(), | ||
| LongType(), | ||
| ] | ||
|
|
||
| def setInputCols(self, value: List[str]) -> "ArrayContainsTransformer": | ||
| """ | ||
| Sets the input columns, ensuring exactly two are provided: | ||
| `[arrayCol, valueCol]`. | ||
|
|
||
| :param value: List of two input column names. | ||
| :returns: Instance of class with input columns set. | ||
| """ | ||
| if len(value) != 2: | ||
| raise ValueError(f"Expected 2 input cols, received {len(value)} instead.") | ||
|
|
||
| return self._set(inputCols=value) | ||
|
|
||
| def _transform(self, dataset: DataFrame) -> DataFrame: | ||
| """ | ||
| Adds `outputCol`, `True` where the value is present in the innermost | ||
| array. The value is a second input column or `mathFloatConstant`. | ||
|
|
||
| :param dataset: Pyspark dataframe to transform. | ||
| :returns: Transformed pyspark dataframe. | ||
| """ | ||
| array_col, value_col = self.get_multiple_input_cols("mathFloatConstant", 2) | ||
| df = dataset.select(array_col, value_col) | ||
| c_dt = self.get_column_datatype(df, df.columns[0]) | ||
| e_dt = get_element_type(c_dt) | ||
|
|
||
| if self.getMathFloatConstant() is not None and isinstance(e_dt, FractionalType): | ||
| value_col = value_col.cast(e_dt) | ||
|
|
||
| output_col = single_input_single_output_array_transform( | ||
| input_col=array_col, | ||
| input_col_datatype=c_dt, | ||
| func=lambda x: F.coalesce(F.array_contains(x, value_col), F.lit(False)), | ||
| ) | ||
| return dataset.withColumn(self.getOutputCol(), output_col) | ||
|
|
||
| def get_keras_layer(self) -> keras.layers.Layer: | ||
| """ | ||
| Gets the Keras layer for the array contains transformer. | ||
|
|
||
| :returns: Keras layer with name equal to the layerName parameter that | ||
| performs the array contains operation. | ||
| """ | ||
| return ArrayContainsLayer( | ||
| name=self.getLayerName(), | ||
| input_dtype=self.getInputKerasDtype(), | ||
| output_dtype=self.getOutputKerasDtype(), | ||
| axis=-1, | ||
| keepdims=True, | ||
| value_constant=self.getMathFloatConstant(), | ||
| ) | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does this work for int values? Or just int arrays? If so can we add a test please. All tests currently are just with float values
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah should work for both arrays and scalars, added a test