From 9fb0a8a329ebcb26a20d479a5920964aff1e7c7e Mon Sep 17 00:00:00 2001 From: Abhishek Das <90371699+dummy1cx@users.noreply.github.com> Date: Tue, 11 Aug 2026 23:48:35 +0100 Subject: [PATCH 1/2] feat: add ListSum transformer and layer --- README.md | 1 + src/kamae/keras/tensorflow/layers/__init__.py | 2 + src/kamae/keras/tensorflow/layers/list_sum.py | 190 +++++ src/kamae/spark/transformers/__init__.py | 1 + src/kamae/spark/transformers/list_sum.py | 206 +++++ .../keras/tensorflow/layers/test_list_sum.py | 632 +++++++++++++++ tests/kamae/keras/test_jit_compatibility.py | 2 + tests/kamae/keras/test_layer_serialisation.py | 14 + .../kamae/spark/transformers/test_list_sum.py | 741 ++++++++++++++++++ 9 files changed, 1789 insertions(+) create mode 100644 src/kamae/keras/tensorflow/layers/list_sum.py create mode 100644 src/kamae/spark/transformers/list_sum.py create mode 100644 tests/kamae/keras/tensorflow/layers/test_list_sum.py create mode 100644 tests/kamae/spark/transformers/test_list_sum.py 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..2e9eb5a9 --- /dev/null +++ b/src/kamae/keras/tensorflow/layers/list_sum.py @@ -0,0 +1,190 @@ +# 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 NaNs results with. 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)) + + # 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) + + 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..6419b16b --- /dev/null +++ b/src/kamae/spark/transformers/list_sum.py @@ -0,0 +1,206 @@ +# 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. + :nanFillValue: Value to fill NaNs results with. 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..afb57b3c --- /dev/null +++ b/tests/kamae/keras/tensorflow/layers/test_list_sum.py @@ -0,0 +1,632 @@ +# 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) + + 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..5dd10a46 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,7 @@ "min_filter_value": 0, }, ), + (ListSumLayer, [tf.random.normal((100, 10, 5))], None), ] 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..7abb48e7 --- /dev/null +++ b/tests/kamae/spark/transformers/test_list_sum.py @@ -0,0 +1,741 @@ +# 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", + ) From 914c8ec1f192b85ff32c855fa0582189096cbacb Mon Sep 17 00:00:00 2001 From: Abhishek Das <90371699+dummy1cx@users.noreply.github.com> Date: Thu, 27 Aug 2026 15:55:02 +0100 Subject: [PATCH 2/2] fix: apply nan_fill_value to empty ListSum results nan_fill_value was accepted, documented and serialised but never used in the Keras layer, so setting it had no effect. The Spark transformer does apply it, filling the null that F.sum returns when the min filter leaves a window empty, so the two backends disagreed for any non-default value. Fill empty results in the layer to match, and tidy the docstrings in both files: indentation, run-on paragraphs, over-length lines and a missing param keyword on nanFillValue. --- src/kamae/keras/tensorflow/layers/list_sum.py | 39 ++++- src/kamae/spark/transformers/list_sum.py | 37 +++-- .../keras/tensorflow/layers/test_list_sum.py | 95 ++++++++++++ tests/kamae/keras/test_jit_compatibility.py | 5 + .../kamae/spark/transformers/test_list_sum.py | 144 ++++++++++++++++++ 5 files changed, 300 insertions(+), 20 deletions(-) diff --git a/src/kamae/keras/tensorflow/layers/list_sum.py b/src/kamae/keras/tensorflow/layers/list_sum.py index 2e9eb5a9..f61e76e1 100644 --- a/src/kamae/keras/tensorflow/layers/list_sum.py +++ b/src/kamae/keras/tensorflow/layers/list_sum.py @@ -35,10 +35,12 @@ class ListSumLayer(BaseLayer): - 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. + 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, @@ -77,11 +79,12 @@ def __init__( :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 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 NaNs results with. Defaults to 0. + :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__( @@ -151,6 +154,7 @@ def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: 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: @@ -167,6 +171,33 @@ def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: 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]: diff --git a/src/kamae/spark/transformers/list_sum.py b/src/kamae/spark/transformers/list_sum.py index 6419b16b..e70249aa 100644 --- a/src/kamae/spark/transformers/list_sum.py +++ b/src/kamae/spark/transformers/list_sum.py @@ -54,18 +54,21 @@ class ListSumTransformer( - 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. + - 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. + 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. + - 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. @@ -73,22 +76,24 @@ class ListSumTransformer( :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. + - 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 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 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. - :nanFillValue: Value to fill NaNs results with. 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 diff --git a/tests/kamae/keras/tensorflow/layers/test_list_sum.py b/tests/kamae/keras/tensorflow/layers/test_list_sum.py index afb57b3c..51722f1c 100644 --- a/tests/kamae/keras/tensorflow/layers/test_list_sum.py +++ b/tests/kamae/keras/tensorflow/layers/test_list_sum.py @@ -613,6 +613,101 @@ def test_listwise_sum( ), "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) diff --git a/tests/kamae/keras/test_jit_compatibility.py b/tests/kamae/keras/test_jit_compatibility.py index 5dd10a46..1884c1b7 100644 --- a/tests/kamae/keras/test_jit_compatibility.py +++ b/tests/kamae/keras/test_jit_compatibility.py @@ -271,6 +271,11 @@ }, ), (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/spark/transformers/test_list_sum.py b/tests/kamae/spark/transformers/test_list_sum.py index 7abb48e7..2901e6c1 100644 --- a/tests/kamae/spark/transformers/test_list_sum.py +++ b/tests/kamae/spark/transformers/test_list_sum.py @@ -739,3 +739,147 @@ def test_list_sum_transform_spark_tf_parity( 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", + )