diff --git a/README.md b/README.md index 88ad70ba..0a1ceb82 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,7 @@ torch_transformers = kamae.get_compatible_transformers('torch') | ListMin | Computes the listwise min of a feature, optionally calculated only on the top items based on another given feature. | [Link](src/kamae/keras/tensorflow/layers/list_min.py) | TensorFlow-only | [Link](src/kamae/spark/transformers/list_min.py) | | ListRank | Computes the listwise rank (ordering) of a feature. | [Link](src/kamae/keras/tensorflow/layers/list_rank.py) | TensorFlow-only | [Link](src/kamae/spark/transformers/list_rank.py) | | ListStdDev | Computes the listwise standard deviation of a feature, optionally calculated only on the top items based on another given feature. | [Link](src/kamae/keras/tensorflow/layers/list_std_dev.py) | TensorFlow-only | [Link](src/kamae/spark/transformers/list_std_dev.py) | +| ListSum | Computes the listwise sum of a feature, optionally calculated only on the top items based on another given feature. | [Link](src/kamae/keras/tensorflow/layers/list_sum.py) | TensorFlow-only | [Link](src/kamae/spark/transformers/list_sum.py) | | Log | Applies the natural logarithm `log(alpha + x)` transform . | [Link](src/kamae/keras/core/layers/log.py) | Multi-backend | [Link](src/kamae/spark/transformers/log.py) | | LogicalAnd | Performs an and(x, y) operation on multiple boolean features. | [Link](src/kamae/keras/core/layers/logical_and.py) | Multi-backend | [Link](src/kamae/spark/transformers/logical_and.py) | | LogicalNot | Performs a not(x) operation on a single boolean feature. | [Link](src/kamae/keras/core/layers/logical_not.py) | Multi-backend | [Link](src/kamae/spark/transformers/logical_not.py) | diff --git a/src/kamae/keras/tensorflow/layers/__init__.py b/src/kamae/keras/tensorflow/layers/__init__.py index 81e81324..5e505352 100644 --- a/src/kamae/keras/tensorflow/layers/__init__.py +++ b/src/kamae/keras/tensorflow/layers/__init__.py @@ -37,6 +37,7 @@ from .list_min import ListMinLayer # noqa: F401 from .list_rank import ListRankLayer # noqa: F401 from .list_std_dev import ListStdDevLayer # noqa: F401 +from .list_sum import ListSumLayer # noqa: F401 from .min_hash_index import MinHashIndexLayer # noqa: F401 from .one_hot_encode import OneHotEncodeLayer # noqa: F401 from .ordinal_array_encode import OrdinalArrayEncodeLayer # noqa: F401 @@ -75,6 +76,7 @@ "ListMinLayer", "ListRankLayer", "ListStdDevLayer", + "ListSumLayer", "MinHashIndexLayer", "OneHotEncodeLayer", "OrdinalArrayEncodeLayer", diff --git a/src/kamae/keras/tensorflow/layers/list_sum.py b/src/kamae/keras/tensorflow/layers/list_sum.py new file mode 100644 index 00000000..f61e76e1 --- /dev/null +++ b/src/kamae/keras/tensorflow/layers/list_sum.py @@ -0,0 +1,221 @@ +# 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 +import tensorflow as tf +from keras import KerasTensor + +import kamae +from kamae.keras.core.backend import TENSORFLOW_ONLY +from kamae.keras.core.base import BaseLayer +from kamae.keras.core.utils.input_utils import allow_single_or_multiple_tensor_input +from kamae.keras.tensorflow.utils.list_utils import get_top_n, segmented_operation +from kamae.keras.tensorflow.utils.transform_utils import map_fn_w_axis + + +@tf.keras.utils.register_keras_serializable(package=kamae.__name__) +class ListSumLayer(BaseLayer): + """ + Calculate the sum across the axis dimension. + - If one tensor is passed, the transformer calculates the sum of the tensor + based on all the items in the given axis dimension. + - If inputCols is set, + - If with_segment = True: the layer calculates the sum of the first tensor + segmented by values of the second tensor. + + Example: calculate the sum price of hotels within star ratings + + - If with_segment = False: the layer calculates the sum of the first tensor + based on second tensor's topN items in the same given axis dimension. + + By using the topN items to calculate the statistics, we can better approximate + the real statistics in production. It is suggested to use a large enough topN to + get a good approximation of the statistics, and an important feature to sort on, + such as item's past production. + + Example: calculate the sum price in the same query, based only on the top N + items sorted by descending production. + """ + + supported_backends = TENSORFLOW_ONLY + jit_compatible = True + + def __init__( + self, + name: Optional[str] = None, + input_dtype: Optional[str] = None, + output_dtype: Optional[str] = None, + top_n: Optional[int] = None, + sort_order: str = "asc", + with_segment: bool = False, + min_filter_value: Optional[float] = None, + nan_fill_value: float = 0.0, + axis: int = 1, + **kwargs: Any, + ) -> None: + """ + Initializes the Listwise Sum layer. + + WARNING: The code is fully tested for axis=1 only. Further testing is needed. + + WARNING: The code can be affected by the value of the padding items. Always + make sure to filter out the padding items value with min_filter_value. + + :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 top_n: The number of top items to consider when calculating the sum. + :param sort_order: The order to sort the second tensor by. Defaults to `asc`. + :param with_segment: Whether the second tensor should be used for + segmentation (True) or sorting (False). Defaults to False. + :param min_filter_value: The minimum filter value to ignore values during + calculation. Defaults to None (no filter). + :param nan_fill_value: The value to fill empty results with, i.e. when the + min filter leaves no values to sum. Defaults to 0. + :param axis: The axis to calculate the statistics across. Defaults to 1. + """ + super().__init__( + name=name, input_dtype=input_dtype, output_dtype=output_dtype, **kwargs + ) + self.top_n = top_n + self.sort_order = sort_order + self.min_filter_value = min_filter_value + self.nan_fill_value = nan_fill_value + self.axis = axis + self.with_segment = with_segment + + @property + def compatible_dtypes(self) -> Optional[List[str]]: + """ + Returns the compatible dtypes of the layer. + + :returns: The compatible dtypes of the layer. + """ + return [ + "bfloat16", + "float16", + "float32", + "float64", + "int8", + "int16", + "int32", + "int64", + "string", + ] + + @allow_single_or_multiple_tensor_input + def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: + """ + Calculate the listwise sum, optionally sorting and + filtering based on the second input tensor, or segmenting + based on the second input tensor. Behaviour is set by with_segment. + + :param inputs: The iterable tensor for the feature. + :returns: The new tensor result column. + """ + val_tensor = inputs[0] + output_shape = tf.shape(val_tensor) + + # Define use of second input + if len(inputs) == 2: + if self.with_segment: + segment_tensor = inputs[1] + else: + sort_tensor = inputs[1] + if self.top_n is None: + raise ValueError("topN must be specified when using a sort column.") + val_tensor = get_top_n( + val_tensor=val_tensor, + axis=self.axis, + sort_tensor=sort_tensor, + sort_order=self.sort_order, + top_n=self.top_n, + ) + else: + if self.with_segment: + raise ValueError("with_segment set to True, expected two inputs.") + + # Values excluded by the min filter contribute 0 to the sum. + # Kept int/string-safe (no float-only ops), mirroring ListMaxLayer, so + # integer value columns and string segment columns both work. + if self.min_filter_value is not None: + mask = tf.greater_equal(val_tensor, self.min_filter_value) + val_tensor = tf.where(mask, val_tensor, tf.zeros_like(val_tensor)) + kept = tf.cast(mask, tf.int32) + + # Apply segmented calculation + if self.with_segment: + listwise_sum = map_fn_w_axis( + elems=[val_tensor, segment_tensor], + fn=lambda x: segmented_operation(x, tf.math.unsorted_segment_sum), + axis=self.axis, + fn_output_signature=tf.TensorSpec( + shape=val_tensor.shape[self.axis :], dtype=val_tensor.dtype + ), + ) + listwise_sum = tf.ensure_shape(listwise_sum, val_tensor.shape) + else: + listwise_sum = tf.reduce_sum(val_tensor, axis=self.axis, keepdims=True) + listwise_sum = tf.broadcast_to(listwise_sum, output_shape) + + if self.min_filter_value is not None: + # Summing zero surviving values gives 0, which is indistinguishable from a + # genuine zero sum. Spark yields null there and fills it with nanFillValue, + # so the same substitution is needed here to keep the two in parity. + if self.with_segment: + any_kept = map_fn_w_axis( + elems=[kept, segment_tensor], + fn=lambda x: segmented_operation(x, tf.math.unsorted_segment_max), + axis=self.axis, + fn_output_signature=tf.TensorSpec( + shape=kept.shape[self.axis :], dtype=kept.dtype + ), + ) + any_kept = tf.ensure_shape(any_kept, kept.shape) + else: + any_kept = tf.reduce_max(kept, axis=self.axis, keepdims=True) + any_kept = tf.broadcast_to(any_kept, output_shape) + + # nan_fill_value is a Python float, which tf.constant cannot convert + # directly to an integer dtype. Narrowing via numpy first handles the + # integer dtypes while preserving full precision for the float ones. + fill_val = tf.constant( + listwise_sum.dtype.as_numpy_dtype(self.nan_fill_value), + dtype=listwise_sum.dtype, + ) + listwise_sum = tf.where(any_kept > 0, listwise_sum, fill_val) + + return listwise_sum + + def get_config(self) -> Dict[str, Any]: + """ + Gets the configuration of the layer. + Used for saving and loading from a model. + + :returns: Dictionary of the configuration of the layer. + """ + config = super().get_config() + config.update( + { + "top_n": self.top_n, + "sort_order": self.sort_order, + "min_filter_value": self.min_filter_value, + "nan_fill_value": self.nan_fill_value, + "axis": self.axis, + "with_segment": self.with_segment, + } + ) + return config diff --git a/src/kamae/spark/transformers/__init__.py b/src/kamae/spark/transformers/__init__.py index de563832..15ad00d5 100644 --- a/src/kamae/spark/transformers/__init__.py +++ b/src/kamae/spark/transformers/__init__.py @@ -51,6 +51,7 @@ from .list_min import ListMinTransformer # noqa: F401 from .list_rank import ListRankTransformer # noqa: F401 from .list_std_dev import ListStdDevTransformer # noqa: F401 +from .list_sum import ListSumTransformer # noqa: F401 from .log import LogTransformer # noqa: F401 from .logical_and import LogicalAndTransformer # noqa: F401 from .logical_not import LogicalNotTransformer # noqa: F401 diff --git a/src/kamae/spark/transformers/list_sum.py b/src/kamae/spark/transformers/list_sum.py new file mode 100644 index 00000000..e70249aa --- /dev/null +++ b/src/kamae/spark/transformers/list_sum.py @@ -0,0 +1,211 @@ +# 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 List, Optional + +import pyspark.sql.functions as F +import tensorflow as tf +from pyspark import keyword_only +from pyspark.sql import DataFrame +from pyspark.sql.types import ( + ByteType, + DataType, + DoubleType, + FloatType, + IntegerType, + LongType, + ShortType, + StringType, +) + +from kamae.keras.core.backend import TENSORFLOW_ONLY +from kamae.keras.tensorflow.layers import ListSumLayer +from kamae.spark.params import ( + ListwiseStatisticsParams, + MultiInputSingleOutputParams, + NanFillValueParams, + SingleInputSingleOutputParams, +) +from kamae.spark.utils import check_and_apply_listwise_op + +from .base import BaseTransformer + + +class ListSumTransformer( + BaseTransformer, + SingleInputSingleOutputParams, + MultiInputSingleOutputParams, + ListwiseStatisticsParams, + NanFillValueParams, +): + """ + Calculate the listwise sum across the query id column. + - If inputCol is set, the transformer calculates the sum of the input column + based on all the items with the same query id column value. + - If inputCols is set, behaviour depends on the value of withSegment: + - If withSegment = True: the transformer calculates the sum of the first + column with the same query id column value, segmented by values of the + second column. + + Example: calculate the sum price of hotels within star ratings, in the + same query. + + - If withSegment = False: the transformer calculates the sum of the first + column with the same query id column value, based on second column's topN + items. When using the second input as sorting column, topN must be provided. + + By using the topN items to calculate the statistics, we can better + approximate the real statistics in production. A large enough topN should + be used, to obtain a good approximation of the statistics, and an important + feature to sort on, such as item's production. + + Example: calculate the sum price in the same query, based on the top N + items sorted by descending production. + + :param inputCol: Value column, on which to calculate the sum. + :param inputCols: Input column names. + - The first is the value column, on which to calculate the sum. + - The second is the sort or segment column. The role of the second input is + governed by the value of withSegment as described above. + :param outputCol: Name of output col. + :param inputDtype: Data Type of input. + :param outputDtype: Data Type of output. + :param layerName: The name of the transformer, which typically + should be the name of the produced feature. + :param queryIdCol: Name of column to aggregate upon. It is required. + :param topN: Filter for limiting the items to calculate the statistics. + Not used when withSegment = True. + :param sortOrder: Option of 'asc' or 'desc' which defines order + for listwise operation. Default is 'asc'. Not used when withSegment = True. + :param withSegment: Whether to use the second input column to partition the + statistic calculation. Defaults to False. + :param minFilterValue: Minimum value to remove padded values + defaults to >= 0. + :param nanFillValue: Value to fill empty results with, i.e. when the min filter + leaves no values to sum. Defaults to 0. + """ + + jit_compatible = True + + supported_backends = TENSORFLOW_ONLY + + @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, + queryIdCol: Optional[str] = None, + topN: Optional[int] = None, + sortOrder: str = "asc", + withSegment: bool = False, + minFilterValue: Optional[float] = None, + nanFillValue: float = 0.0, + ) -> None: + super().__init__() + self._setDefault( + topN=None, + sortOrder="asc", + minFilterValue=None, + nanFillValue=0, + withSegment=False, + ) + 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(), + StringType(), + ] + + def _transform(self, dataset: DataFrame) -> DataFrame: + """ + Calculate the listwise sum, optionally sorting and + filtering based on the second input column. + :param dataset: The dataframe with signals and features. + :returns: The dataframe dataset with the new feature. + """ + if not self.isDefined("queryIdCol"): + raise ValueError("queryIdCol must be set on listwise transformers.") + + # Define the columns to use for the calculation + if self.isDefined("inputCols"): + with_segment = self.getWithSegment() + if with_segment: + val_col_name = self.getInputCols()[0] + segment_col_name = self.getInputCols()[1] + sort_col_name = None + else: + val_col_name = self.getInputCols()[0] + sort_col_name = self.getInputCols()[1] + segment_col_name = None + else: + val_col_name = self.getInputCol() + sort_col_name = None + segment_col_name = None + + dataset = dataset.withColumn( + self.getOutputCol(), + check_and_apply_listwise_op( + dataset, + F.sum, + self.getQueryIdCol(), + val_col_name, + sort_col_name, + segment_col_name, + self.getSortOrder(), + self.getTopN(), + self.getMinFilterValue(), + ), + ) + + # Replace Nulls/Nans + dataset = dataset.fillna({self.getOutputCol(): self.getNanFillValue()}) + + return dataset + + def get_keras_layer(self) -> tf.keras.layers.Layer: + """ + Gets the Keras layer for the listwise-sum transformer. + + :returns: Keras layer with name equal to the layerName parameter that + performs a summing operation. + """ + return ListSumLayer( + name=self.getLayerName(), + input_dtype=self.getInputKerasDtype(), + output_dtype=self.getOutputKerasDtype(), + top_n=self.getTopN(), + sort_order=self.getSortOrder(), + with_segment=self.getWithSegment(), + min_filter_value=self.getMinFilterValue(), + nan_fill_value=self.getNanFillValue(), + axis=1, + ) diff --git a/tests/kamae/keras/tensorflow/layers/test_list_sum.py b/tests/kamae/keras/tensorflow/layers/test_list_sum.py new file mode 100644 index 00000000..51722f1c --- /dev/null +++ b/tests/kamae/keras/tensorflow/layers/test_list_sum.py @@ -0,0 +1,727 @@ +# 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.tensorflow.layers import ListSumLayer + + +class TestListSum: + @pytest.mark.parametrize( + "inputs, min_filter_value, top_n, with_segment, sort_order, input_dtype, output_dtype, expected_output", + [ + # Base case + ( + [ + # values + tf.constant( + [ + [ + [1.0], + [1.0], + [9.0], + [4.0], + [6.0], + [2.0], + [0.0], + [0.0], + ], + [ + [5.0], + [1.0], + [9.0], + [4.0], + [6.0], + [8.0], + [0.0], + [0.0], + ], + ], + dtype=tf.float32, + ), + ], + None, + None, + False, + "asc", + "float64", + "float32", + tf.constant( + [ + [ + [23.0], + [23.0], + [23.0], + [23.0], + [23.0], + [23.0], + [23.0], + [23.0], + ], + [ + [33.0], + [33.0], + [33.0], + [33.0], + [33.0], + [33.0], + [33.0], + [33.0], + ], + ], + dtype=tf.float32, + ), + ), + # With min_filter_value. The excluded values are zeros, which contribute + # nothing to a sum, so the result matches the base case. + ( + [ + tf.constant( + [ + [ + [1.0], + [1.0], + [9.0], + [4.0], + [6.0], + [2.0], + [0.0], + [0.0], + ], + [ + [5.0], + [1.0], + [9.0], + [4.0], + [6.0], + [8.0], + [0.0], + [0.0], + ], + ], + dtype=tf.float32, + ), + ], + 1, + None, + False, + "asc", + "float64", + "float32", + tf.constant( + [ + [ + [23.0], + [23.0], + [23.0], + [23.0], + [23.0], + [23.0], + [23.0], + [23.0], + ], + [ + [33.0], + [33.0], + [33.0], + [33.0], + [33.0], + [33.0], + [33.0], + [33.0], + ], + ], + dtype=tf.float32, + ), + ), + # With top_n + ( + [ + # values + tf.constant( + [ + [ + [1.0], + [1.0], + [9.0], + [4.0], + [6.0], + [2.0], + [0.0], + [0.0], + ], + [ + [5.0], + [1.0], + [9.0], + [4.0], + [6.0], + [8.0], + [0.0], + [0.0], + ], + ], + dtype=tf.float32, + ), + # sort + tf.constant( + [ + [ + [1.0], + [2.0], + [3.0], + [4.0], + [5.0], + [6.0], + [7.0], + [8.0], + ], + [ + [8.0], + [7.0], + [6.0], + [5.0], + [4.0], + [3.0], + [2.0], + [1.0], + ], + ], + dtype=tf.float32, + ), + ], + None, + 5, + False, + "asc", + "float64", + "float32", + # Top 5 ascending picks values [1, 1, 9, 4, 6] and [0, 0, 8, 6, 4]. + tf.constant( + [ + [ + [21.0], + [21.0], + [21.0], + [21.0], + [21.0], + [21.0], + [21.0], + [21.0], + ], + [ + [18.0], + [18.0], + [18.0], + [18.0], + [18.0], + [18.0], + [18.0], + [18.0], + ], + ], + dtype=tf.float32, + ), + ), + # With top_n and filter + ( + [ + # values + tf.constant( + [ + [ + [1.0], + [1.0], + [9.0], + [4.0], + [6.0], + [2.0], + [0.0], + [0.0], + ], + [ + [5.0], + [1.0], + [9.0], + [4.0], + [6.0], + [8.0], + [0.0], + [0.0], + ], + ], + dtype=tf.float32, + ), + # sort + tf.constant( + [ + [ + [1.0], + [2.0], + [3.0], + [4.0], + [5.0], + [6.0], + [7.0], + [8.0], + ], + [ + [8.0], + [7.0], + [6.0], + [5.0], + [4.0], + [3.0], + [2.0], + [1.0], + ], + ], + dtype=tf.float32, + ), + ], + 1, + 5, + False, + "asc", + "float64", + "float32", + tf.constant( + [ + [ + [21.0], + [21.0], + [21.0], + [21.0], + [21.0], + [21.0], + [21.0], + [21.0], + ], + [ + [18.0], + [18.0], + [18.0], + [18.0], + [18.0], + [18.0], + [18.0], + [18.0], + ], + ], + dtype=tf.float32, + ), + ), + # With top_n > list size + ( + [ + # values + tf.constant( + [ + [ + [1.0], + [1.0], + [9.0], + ], + [ + [5.0], + [1.0], + [9.0], + ], + ], + dtype=tf.float32, + ), + # sort + tf.constant( + [ + [ + [1.0], + [2.0], + [3.0], + ], + [ + [8.0], + [7.0], + [6.0], + ], + ], + dtype=tf.float32, + ), + ], + 1, + 5, + False, + "asc", + "float64", + "float32", + tf.constant( + [ + [ + [11.0], + [11.0], + [11.0], + ], + [ + [15.0], + [15.0], + [15.0], + ], + ], + dtype=tf.float32, + ), + ), + # With segmentation + ( + [ + # values + tf.constant( + [ + [ + [1.0], + [1.0], + [9.0], + ], + [ + [5.0], + [1.0], + [9.0], + ], + ], + dtype=tf.float32, + ), + # segment + tf.constant( + [ + [ + [1.0], + [2.0], + [2.0], + ], + [ + [1.0], + [2.0], + [2.0], + ], + ], + dtype=tf.float32, + ), + ], + None, + None, + True, + "asc", + "float64", + "float32", + tf.constant( + [ + [ + [1.0], + [10.0], + [10.0], + ], + [ + [5.0], + [10.0], + [10.0], + ], + ], + dtype=tf.float32, + ), + ), + # With segmentation and multiple features + ( + [ + # values + tf.constant( + [[[1.0, 10.0], [2.0, 20.0], [3.0, 30.0]]], dtype=tf.float32 + ), + # segment + tf.constant( + [[[1.0, 1.0], [2.0, 2.0], [2.0, 2.0]]], dtype=tf.float32 + ), + ], + None, + None, + True, + "asc", + "float64", + "float32", + tf.constant( + [[[1.0, 10.0], [5.0, 50.0], [5.0, 50.0]]], dtype=tf.float32 + ), + ), + # With segmentation ID as string + ( + [ + # values + tf.constant( + [ + [ + [1.0], + [1.0], + [9.0], + ], + [ + [5.0], + [1.0], + [9.0], + ], + ], + dtype=tf.float32, + ), + # segment + tf.constant( + [ + [ + ["1.0"], + ["2.0"], + ["2.0"], + ], + [ + ["1.0"], + ["2.0"], + ["2.0"], + ], + ], + dtype=tf.string, + ), + ], + None, + None, + True, + "asc", + "float64", + "float32", + tf.constant( + [ + [ + [1.0], + [10.0], + [10.0], + ], + [ + [5.0], + [10.0], + [10.0], + ], + ], + dtype=tf.float32, + ), + ), + # With segmentation and min_filter_val + ( + [ + # values + tf.constant( + [ + [ + [1.0], + [1.0], + [9.0], + ], + [ + [5.0], + [1.0], + [9.0], + ], + ], + dtype=tf.float32, + ), + # segment + tf.constant( + [ + [ + [1.0], + [2.0], + [2.0], + ], + [ + [1.0], + [2.0], + [2.0], + ], + ], + dtype=tf.float32, + ), + ], + 2.0, + None, + True, + "asc", + "float64", + "float32", + tf.constant( + [ + [ + [0.0], + [9.0], + [9.0], + ], + [ + [5.0], + [9.0], + [9.0], + ], + ], + dtype=tf.float32, + ), + ), + ], + ) + def test_listwise_sum( + self, + inputs, + min_filter_value, + top_n, + with_segment, + sort_order, + input_dtype, + output_dtype, + expected_output, + ): + # when + name = "listwise_sum_test" + layer = ListSumLayer( + name=name, + min_filter_value=min_filter_value, + input_dtype=input_dtype, + output_dtype=output_dtype, + sort_order=sort_order, + top_n=top_n, + with_segment=with_segment, + ) + inputs = inputs if len(inputs) > 1 else inputs[0] + output_tensor = layer(inputs) + # then + assert layer.name == name, "Layer name is not set properly" + assert ( + output_tensor.dtype == expected_output.dtype + ), "Output tensor dtype is not the same as expected tensor dtype" + assert ( + output_tensor.shape == expected_output.shape + ), "Output tensor shape is not the same as expected tensor shape" + tf.debugging.assert_equal(output_tensor, expected_output) + + @pytest.mark.parametrize( + "inputs, min_filter_value, nan_fill_value, with_segment, expected_output", + [ + # Nothing survives the filter, so the second list has nothing to sum + ( + [ + tf.constant( + [ + [[1.0], [2.0], [3.0]], + [[-999.0], [-999.0], [-999.0]], + ], + dtype=tf.float32, + ), + ], + 0.0, + -1.0, + False, + tf.constant( + [ + [[6.0], [6.0], [6.0]], + [[-1.0], [-1.0], [-1.0]], + ], + dtype=tf.float32, + ), + ), + # The default fill value leaves the empty sum at zero + ( + [ + tf.constant( + [ + [[1.0], [2.0], [3.0]], + [[-999.0], [-999.0], [-999.0]], + ], + dtype=tf.float32, + ), + ], + 0.0, + 0.0, + False, + tf.constant( + [ + [[6.0], [6.0], [6.0]], + [[0.0], [0.0], [0.0]], + ], + dtype=tf.float32, + ), + ), + # A list that genuinely sums to zero keeps its zero + ( + [ + tf.constant([[[0.0], [0.0], [0.0]]], dtype=tf.float32), + ], + 0.0, + -1.0, + False, + tf.constant([[[0.0], [0.0], [0.0]]], dtype=tf.float32), + ), + # The filter empties one segment but not the other + ( + [ + # values + tf.constant([[[1.0], [4.0], [-999.0]]], dtype=tf.float32), + # segment + tf.constant([[[1.0], [1.0], [2.0]]], dtype=tf.float32), + ], + 0.0, + -1.0, + True, + tf.constant([[[5.0], [5.0], [-1.0]]], dtype=tf.float32), + ), + ], + ) + def test_listwise_sum_nan_fill_value( + self, + inputs, + min_filter_value, + nan_fill_value, + with_segment, + expected_output, + ): + # given + layer = ListSumLayer( + name="listwise_sum_nan_fill_value", + min_filter_value=min_filter_value, + nan_fill_value=nan_fill_value, + with_segment=with_segment, + ) + # when + output_tensor = layer(inputs if len(inputs) > 1 else inputs[0]) + # 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) + + def test_listwise_sum_raises_without_top_n_when_sorting(self): + # given + layer = ListSumLayer(name="listwise_sum_no_top_n", with_segment=False) + inputs = [ + tf.constant([[[1.0], [2.0], [3.0]]]), + tf.constant([[[1.0], [2.0], [3.0]]]), + ] + # when / then + with pytest.raises(ValueError, match="topN must be specified"): + layer(inputs) + + def test_listwise_sum_raises_with_segment_and_single_input(self): + # given + layer = ListSumLayer(name="listwise_sum_single_input", with_segment=True) + # when / then + with pytest.raises(ValueError, match="expected two inputs"): + layer(tf.constant([[[1.0], [2.0], [3.0]]])) diff --git a/tests/kamae/keras/test_jit_compatibility.py b/tests/kamae/keras/test_jit_compatibility.py index bd9c2354..1884c1b7 100644 --- a/tests/kamae/keras/test_jit_compatibility.py +++ b/tests/kamae/keras/test_jit_compatibility.py @@ -78,6 +78,7 @@ ListMinLayer, ListRankLayer, ListStdDevLayer, + ListSumLayer, MinHashIndexLayer, OneHotEncodeLayer, OrdinalArrayEncodeLayer, @@ -269,6 +270,12 @@ "min_filter_value": 0, }, ), + (ListSumLayer, [tf.random.normal((100, 10, 5))], None), + ( + ListSumLayer, + [tf.random.normal((100, 10, 5))], + {"min_filter_value": 0.0, "nan_fill_value": -1.0}, + ), ] diff --git a/tests/kamae/keras/test_layer_serialisation.py b/tests/kamae/keras/test_layer_serialisation.py index e634e7b3..160c471f 100644 --- a/tests/kamae/keras/test_layer_serialisation.py +++ b/tests/kamae/keras/test_layer_serialisation.py @@ -85,6 +85,7 @@ ListMinLayer, ListRankLayer, ListStdDevLayer, + ListSumLayer, MinHashIndexLayer, OneHotEncodeLayer, OrdinalArrayEncodeLayer, @@ -251,6 +252,7 @@ (ListMaxLayer, [tf.random.normal((100, 10, 5))], None, False), (ListMeanLayer, [tf.random.normal((100, 10, 5))], None, False), (ListMinLayer, [tf.random.normal((100, 10, 5))], None, False), + (ListSumLayer, [tf.random.normal((100, 10, 5))], None, False), ( IfStatementLayer, [tf.random.normal((100, 10, 5)), tf.random.normal((100, 10, 5))], @@ -425,6 +427,18 @@ }, False, ), + ( + ListSumLayer, + [tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])], + { + "axis": 1, + "top_n": 5, + "sort_order": "descending", + "nan_fill_value": 0, + "min_filter_value": 0, + }, + False, + ), ( StringAffixLayer, [tf.constant("a", shape=(100, 10, 1))], diff --git a/tests/kamae/spark/transformers/test_list_sum.py b/tests/kamae/spark/transformers/test_list_sum.py new file mode 100644 index 00000000..2901e6c1 --- /dev/null +++ b/tests/kamae/spark/transformers/test_list_sum.py @@ -0,0 +1,885 @@ +# 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 pyspark.sql.functions as F +import pytest +import tensorflow as tf + +from kamae.spark.transformers import ListSumTransformer + +from ..test_helpers import tensor_to_python_type + + +class TestListSum: + @pytest.fixture(scope="class") + def listwise_transform_df_no_filter(self, spark_session): + return spark_session.createDataFrame( + [ + (1, 2, 14.0), + (1, 2, 14.0), + (1, 2, 14.0), + (1, 8, 14.0), + (2, 10, 30.0), + (2, 20, 30.0), + (3, None, 5.0), # should be ignored + (3, 5, 5.0), + ], + [ + "search_id", + "value_col", + "expected", + ], + ) + + @pytest.fixture(scope="class") + def listwise_transform_df_min_value(self, spark_session): + return spark_session.createDataFrame( + [ + (1, -999, 12.0), # should be ignored + (1, 2, 12.0), + (1, 2, 12.0), + (1, 8, 12.0), + (2, -999, 20.0), # should be ignored + (2, 20, 20.0), + (3, None, 5.0), # should be ignored + (3, 5, 5.0), + ], + [ + "search_id", + "value_col", + "expected", + ], + ) + + @pytest.fixture(scope="class") + def listwise_transform_df_sort_desc(self, spark_session): + return spark_session.createDataFrame( + [ + (1, 1, 1, 12.0), # should be ignored in top3 desc + (1, 2, 2, 12.0), + (1, 2, 3, 12.0), + (1, 8, 4, 12.0), + ], + [ + "search_id", + "value_col", + "sort_col", + "expected", + ], + ) + + @pytest.fixture(scope="class") + def listwise_transform_df_sort_asc(self, spark_session): + return spark_session.createDataFrame( + [ + (1, 5, 1, 9.0), + (1, 2, 2, 9.0), + (1, 2, 3, 9.0), + (1, 8, 4, 9.0), # should be ignored in top3 asc + ], + [ + "search_id", + "value_col", + "sort_col", + "expected", + ], + ) + + @pytest.fixture(scope="class") + def listwise_transform_df_segment(self, spark_session): + return spark_session.createDataFrame( + [ + (1, 5, 1, 13.0), + (1, 2, 2, 4.0), + (1, 2, 2, 4.0), + (1, 8, 1, 13.0), + ], + [ + "search_id", + "value_col", + "segment_col", + "expected", + ], + ) + + @pytest.mark.parametrize( + "input_dataframe, value_col, min_filter_value, output_col, input_dtype, output_dtype", + [ + ( + "listwise_transform_df_no_filter", + "value_col", + None, + "expected", + "float", + "float", + ), + ( + "listwise_transform_df_min_value", + "value_col", + 0.0, + "expected", + "float", + "float", + ), + ], + ) + def test_spark_sum_transform( + self, + input_dataframe, + value_col, + min_filter_value, + output_col, + input_dtype, + output_dtype, + request, + ): + # given + input_dataframe = request.getfixturevalue(input_dataframe) + # when + transformer = ListSumTransformer( + inputCol=value_col, + outputCol=output_col, + inputDtype=input_dtype, + outputDtype=output_dtype, + queryIdCol="search_id", + minFilterValue=min_filter_value, + ) + actual = transformer.transform(input_dataframe.drop("expected")) + # then + expected = input_dataframe.select( + F.col("expected").cast(output_dtype).alias("expected") + ) + diff = actual.select("expected").exceptAll(expected) + assert diff.isEmpty(), "Expected and actual dataframes are not equal" + + @pytest.mark.parametrize( + "input_dataframe, value_col, sort_col, top_n, sort_order, output_col, input_dtype, output_dtype", + [ + ( + "listwise_transform_df_sort_desc", + "value_col", + "sort_col", + 3, + "desc", + "expected", + "float", + "float", + ), + ( + "listwise_transform_df_sort_asc", + "value_col", + "sort_col", + 3, + "asc", + "expected", + "float", + "float", + ), + ], + ) + def test_spark_sum_transform_with_sort( + self, + input_dataframe, + value_col, + sort_col, + top_n, + sort_order, + output_col, + input_dtype, + output_dtype, + request, + ): + # given + input_dataframe = request.getfixturevalue(input_dataframe) + # when + transformer = ListSumTransformer( + inputCols=[value_col, sort_col], + outputCol=output_col, + inputDtype=input_dtype, + outputDtype=output_dtype, + queryIdCol="search_id", + topN=top_n, + sortOrder=sort_order, + ) + actual = transformer.transform(input_dataframe.drop("expected")) + # then + expected = input_dataframe.select( + F.col("expected").cast(output_dtype).alias("expected") + ) + diff = actual.select("expected").exceptAll(expected) + assert diff.isEmpty(), "Expected and actual dataframes are not equal" + + @pytest.mark.parametrize( + "input_dataframe, value_col, segment_col, output_col, input_dtype, output_dtype", + [ + ( + "listwise_transform_df_segment", + "value_col", + "segment_col", + "expected", + "float", + "float", + ), + ], + ) + def test_spark_sum_transform_with_segmentation( + self, + input_dataframe, + value_col, + segment_col, + output_col, + input_dtype, + output_dtype, + request, + ): + # given + input_dataframe = request.getfixturevalue(input_dataframe) + # when + transformer = ListSumTransformer( + inputCols=[value_col, segment_col], + outputCol=output_col, + inputDtype=input_dtype, + outputDtype=output_dtype, + queryIdCol="search_id", + withSegment=True, + ) + actual = transformer.transform(input_dataframe.drop("expected")) + # then + expected = input_dataframe.select( + F.col("expected").cast(output_dtype).alias("expected") + ) + diff = actual.select("expected").exceptAll(expected) + assert diff.isEmpty(), "Expected and actual dataframes are not equal" + + def test_spark_sum_transform_raises_without_query_id_col( + self, listwise_transform_df_no_filter + ): + # given + transformer = ListSumTransformer( + inputCol="value_col", + outputCol="output", + ) + # when / then + with pytest.raises(ValueError): + transformer.transform(listwise_transform_df_no_filter.drop("expected")) + + @pytest.mark.parametrize( + "list_size, qid_tensor, input_tensors, min_filter_value, with_segment, top_n, input_dtype, output_dtype", + [ + # Base case + ( + 8, + tf.constant( + [ + [1], + [1], + [1], + [1], + [1], + [1], + [1], + [1], + [2], + [2], + [2], + [2], + [2], + [2], + [2], + [2], + ], + dtype=tf.float32, + ), + [ + # values + tf.constant( + [ + [1.0], + [1.0], + [9.0], + [4.0], + [6.0], + [2.0], + [0.0], + [0.0], + [5.0], + [1.0], + [9.0], + [4.0], + [6.0], + [8.0], + [0.0], + [0.0], + ], + dtype=tf.float32, + ), + ], + None, + False, + None, + "double", + "float", + ), + # With min_filter_value + ( + 8, + tf.constant( + [ + [1], + [1], + [1], + [1], + [1], + [1], + [1], + [1], + [2], + [2], + [2], + [2], + [2], + [2], + [2], + [2], + ], + dtype=tf.float32, + ), + [ + # values + tf.constant( + [ + [1.0], + [1.0], + [9.0], + [4.0], + [6.0], + [2.0], + [0.0], + [0.0], + [5.0], + [1.0], + [9.0], + [4.0], + [6.0], + [8.0], + [0.0], + [0.0], + ], + dtype=tf.float32, + ), + ], + 1, + False, + None, + "double", + "float", + ), + # With top_n + ( + 8, + tf.constant( + [ + [1], + [1], + [1], + [1], + [1], + [1], + [1], + [1], + [2], + [2], + [2], + [2], + [2], + [2], + [2], + [2], + ], + dtype=tf.float32, + ), + [ + # values + tf.constant( + [ + [1.0], + [1.0], + [9.0], + [4.0], + [6.0], + [2.0], + [0.0], + [0.0], + [5.0], + [1.0], + [9.0], + [4.0], + [6.0], + [8.0], + [0.0], + [0.0], + ], + dtype=tf.float32, + ), + # sort + tf.constant( + [ + [1.0], + [2.0], + [3.0], + [4.0], + [5.0], + [6.0], + [7.0], + [8.0], + [8.0], + [7.0], + [6.0], + [5.0], + [4.0], + [3.0], + [2.0], + [1.0], + ], + dtype=tf.float32, + ), + ], + None, + False, + 5, + "double", + "float", + ), + # With top_n and filter + ( + 8, + tf.constant( + [ + [1], + [1], + [1], + [1], + [1], + [1], + [1], + [1], + [2], + [2], + [2], + [2], + [2], + [2], + [2], + [2], + ], + dtype=tf.float32, + ), + [ + # values + tf.constant( + [ + [1.0], + [1.0], + [9.0], + [4.0], + [6.0], + [2.0], + [0.0], + [0.0], + [5.0], + [1.0], + [9.0], + [4.0], + [6.0], + [8.0], + [0.0], + [0.0], + ], + dtype=tf.float32, + ), + # sort + tf.constant( + [ + [1.0], + [2.0], + [3.0], + [4.0], + [5.0], + [6.0], + [7.0], + [8.0], + [8.0], + [7.0], + [6.0], + [5.0], + [4.0], + [3.0], + [2.0], + [1.0], + ], + dtype=tf.float32, + ), + ], + 1, + False, + 5, + "double", + "float", + ), + # With top_n > list size + ( + 3, + tf.constant( + [ + [1], + [1], + [1], + [2], + [2], + [2], + ], + dtype=tf.float32, + ), + [ + # values + tf.constant( + [ + [1.0], + [1.0], + [9.0], + [5.0], + [1.0], + [9.0], + ], + dtype=tf.float32, + ), + # sort + tf.constant( + [ + [1.0], + [2.0], + [3.0], + [8.0], + [7.0], + [6.0], + ], + dtype=tf.float32, + ), + ], + 1, + False, + 5, + "double", + "float", + ), + # With segmentation + ( + 3, + tf.constant( + [ + [1.0], + [1.0], + [1.0], + [2.0], + [2.0], + [2.0], + ], + dtype=tf.float32, + ), + [ + # values + tf.constant( + [ + [1.0], + [1.0], + [4.0], + [5.0], + [5.0], + [20.0], + ], + dtype=tf.float32, + ), + # segment + tf.constant( + [ + [1.0], + [1.0], + [2.0], + [1.0], + [1.0], + [2.0], + ], + dtype=tf.float32, + ), + ], + 0, + True, + None, + "double", + "float", + ), + # With segmentation & filter + ( + 3, + tf.constant( + [ + [1], + [1], + [1], + [2], + [2], + [2], + ], + dtype=tf.float32, + ), + [ + # values + tf.constant( + [ + [1.0], + [1.0], + [4.0], + [5.0], + [5.0], + [20.0], + ], + dtype=tf.float32, + ), + # segment + tf.constant( + [ + [1.0], + [1.0], + [2.0], + [1.0], + [1.0], + [2.0], + ], + dtype=tf.float32, + ), + ], + 2.0, + True, + None, + "double", + "float", + ), + ], + ) + def test_list_sum_transform_spark_tf_parity( + self, + spark_session, + list_size, + qid_tensor, + input_tensors, + min_filter_value, + with_segment, + top_n, + input_dtype, + output_dtype, + ): + col_names = [f"input{i}" for i in range(len(input_tensors))] + # given + transformer = ListSumTransformer( + inputCol=col_names[0] if len(col_names) == 1 else None, + inputCols=col_names if len(col_names) > 1 else None, + outputCol="output", + inputDtype=input_dtype, + outputDtype=output_dtype, + queryIdCol="search_id", + minFilterValue=min_filter_value, + topN=top_n, + withSegment=with_segment, + sortOrder="asc", + ) + # when + qid_inputs_tensors = [qid_tensor] + input_tensors + qid_col_names = ["search_id"] + col_names + spark_df = spark_session.createDataFrame( + [ + tuple([tensor_to_python_type(ti) for ti in t]) + for t in zip(*qid_inputs_tensors) + ], + qid_col_names, + ) + + spark_values = ( + transformer.transform(spark_df) + .select("output") + .rdd.map(lambda r: r[0]) + .collect() + ) + + # reshape the input tensors to match the expected shape based on list size + input_tensors = [tf.reshape(t, (-1, list_size, 1)) for t in input_tensors] + tensorflow_values = np.reshape( + [ + np.squeeze(v) + for v in transformer.get_keras_layer()(input_tensors).numpy().tolist() + ], + -1, + ) + + # then + np.testing.assert_almost_equal( + spark_values, + tensorflow_values, + decimal=6, + err_msg="Spark and Tensorflow transform outputs are not equal", + ) + + @pytest.mark.parametrize( + "list_size, qid_tensor, input_tensors, min_filter_value, nan_fill_value, with_segment", + [ + # The filter empties the whole of the second query + ( + 3, + tf.constant( + [ + [1.0], + [1.0], + [1.0], + [2.0], + [2.0], + [2.0], + ], + dtype=tf.float32, + ), + [ + # values + tf.constant( + [ + [1.0], + [2.0], + [3.0], + [-999.0], + [-999.0], + [-999.0], + ], + dtype=tf.float32, + ), + ], + 0.0, + -1.0, + False, + ), + # The filter empties one segment of the first query + ( + 3, + tf.constant( + [ + [1.0], + [1.0], + [1.0], + [2.0], + [2.0], + [2.0], + ], + dtype=tf.float32, + ), + [ + # values + tf.constant( + [ + [1.0], + [4.0], + [-999.0], + [5.0], + [5.0], + [20.0], + ], + dtype=tf.float32, + ), + # segment + tf.constant( + [ + [1.0], + [1.0], + [2.0], + [1.0], + [1.0], + [2.0], + ], + dtype=tf.float32, + ), + ], + 0.0, + -1.0, + True, + ), + ], + ) + def test_list_sum_transform_spark_tf_parity_with_nan_fill_value( + self, + spark_session, + list_size, + qid_tensor, + input_tensors, + min_filter_value, + nan_fill_value, + with_segment, + ): + col_names = [f"input{i}" for i in range(len(input_tensors))] + # given + transformer = ListSumTransformer( + inputCol=col_names[0] if len(col_names) == 1 else None, + inputCols=col_names if len(col_names) > 1 else None, + outputCol="output", + inputDtype="double", + outputDtype="float", + queryIdCol="search_id", + minFilterValue=min_filter_value, + nanFillValue=nan_fill_value, + withSegment=with_segment, + sortOrder="asc", + ) + # when + qid_inputs_tensors = [qid_tensor] + input_tensors + qid_col_names = ["search_id"] + col_names + spark_df = spark_session.createDataFrame( + [ + tuple([tensor_to_python_type(ti) for ti in t]) + for t in zip(*qid_inputs_tensors) + ], + qid_col_names, + ) + + spark_values = ( + transformer.transform(spark_df) + .select("output") + .rdd.map(lambda r: r[0]) + .collect() + ) + + # reshape the input tensors to match the expected shape based on list size + input_tensors = [tf.reshape(t, (-1, list_size, 1)) for t in input_tensors] + tensorflow_values = np.reshape( + [ + np.squeeze(v) + for v in transformer.get_keras_layer()(input_tensors).numpy().tolist() + ], + -1, + ) + + # then + assert ( + nan_fill_value in spark_values + ), "Expected the emptied list to fall back to nanFillValue" + np.testing.assert_almost_equal( + spark_values, + tensorflow_values, + decimal=6, + err_msg="Spark and Tensorflow transform outputs are not equal", + )