From ffd79f801053bd4efd9157db0ea62146a269a29e Mon Sep 17 00:00:00 2001 From: Abhishek Das <90371699+dummy1cx@users.noreply.github.com> Date: Wed, 12 Aug 2026 12:39:29 +0100 Subject: [PATCH 1/3] feat: support integer dtypes in ListMax --- src/kamae/keras/tensorflow/layers/list_max.py | 8 +- src/kamae/spark/transformers/list_max.py | 15 +- .../keras/tensorflow/layers/test_list_max.py | 180 ++++++++++++++++++ .../kamae/spark/transformers/test_list_max.py | 146 ++++++++++++++ 4 files changed, 347 insertions(+), 2 deletions(-) diff --git a/src/kamae/keras/tensorflow/layers/list_max.py b/src/kamae/keras/tensorflow/layers/list_max.py index 12e91a52..ecfacbd3 100644 --- a/src/kamae/keras/tensorflow/layers/list_max.py +++ b/src/kamae/keras/tensorflow/layers/list_max.py @@ -108,6 +108,10 @@ def compatible_dtypes(self) -> Optional[List[str]]: "float16", "float32", "float64", + "int8", + "int16", + "int32", + "int64", "string", ] @@ -167,7 +171,9 @@ def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: listwise_max = tf.broadcast_to(listwise_max, output_shape) if self.min_filter_value is not None: - fill_val = tf.constant(self.nan_fill_value, dtype=listwise_max.dtype) + # tf.cast rather than tf.constant(..., dtype=...): nan_fill_value is a + # Python float, which cannot be converted directly to an integer dtype. + fill_val = tf.cast(self.nan_fill_value, listwise_max.dtype) listwise_max = tf.where(listwise_max != neg_inf, listwise_max, fill_val) return listwise_max diff --git a/src/kamae/spark/transformers/list_max.py b/src/kamae/spark/transformers/list_max.py index 2fdc2834..f2f43f2a 100644 --- a/src/kamae/spark/transformers/list_max.py +++ b/src/kamae/spark/transformers/list_max.py @@ -18,7 +18,16 @@ import tensorflow as tf from pyspark import keyword_only from pyspark.sql import DataFrame -from pyspark.sql.types import DataType, DoubleType, FloatType, StringType +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 ListMaxLayer @@ -124,6 +133,10 @@ def compatible_dtypes(self) -> Optional[List[DataType]]: return [ FloatType(), DoubleType(), + ByteType(), + ShortType(), + IntegerType(), + LongType(), StringType(), ] diff --git a/tests/kamae/keras/tensorflow/layers/test_list_max.py b/tests/kamae/keras/tensorflow/layers/test_list_max.py index 7ffa8db1..6fb48295 100644 --- a/tests/kamae/keras/tensorflow/layers/test_list_max.py +++ b/tests/kamae/keras/tensorflow/layers/test_list_max.py @@ -576,6 +576,186 @@ class TestListMax: dtype=tf.float32, ), ), + # Integer values with segmentation + ( + [ + # values + tf.constant( + [ + [ + [1], + [5], + [3], + ], + [ + [7], + [2], + [9], + ], + ], + dtype=tf.int64, + ), + # segment + tf.constant( + [ + [ + [1], + [1], + [2], + ], + [ + [1], + [1], + [2], + ], + ], + dtype=tf.int64, + ), + ], + None, + None, + True, + "asc", + "int64", + "int64", + tf.constant( + [ + [ + [5], + [5], + [3], + ], + [ + [7], + [7], + [9], + ], + ], + dtype=tf.int64, + ), + ), + # Integer values with segmentation and min_filter_value, where the + # filter empties a segment entirely and nan_fill_value is applied. + ( + [ + # values + tf.constant( + [ + [ + [1], + [1], + [9], + ], + [ + [5], + [1], + [9], + ], + ], + dtype=tf.int64, + ), + # segment + tf.constant( + [ + [ + [1], + [2], + [2], + ], + [ + [1], + [2], + [2], + ], + ], + dtype=tf.int64, + ), + ], + 2, + None, + True, + "asc", + "int64", + "int64", + tf.constant( + [ + [ + [0], + [9], + [9], + ], + [ + [5], + [9], + [9], + ], + ], + dtype=tf.int64, + ), + ), + # Remaining integer widths, exercising segmentation together with + # min_filter_value emptying a segment, so that nan_fill_value is + # applied for every integer dtype the layer accepts. + *[ + ( + [ + # values + tf.constant( + [ + [ + [1], + [1], + [9], + ], + [ + [5], + [1], + [9], + ], + ], + dtype=int_dtype, + ), + # segment + tf.constant( + [ + [ + [1], + [2], + [2], + ], + [ + [1], + [2], + [2], + ], + ], + dtype=int_dtype, + ), + ], + 2, + None, + True, + "asc", + int_dtype.name, + int_dtype.name, + tf.constant( + [ + [ + [0], + [9], + [9], + ], + [ + [5], + [9], + [9], + ], + ], + dtype=int_dtype, + ), + ) + for int_dtype in [tf.int8, tf.int16, tf.int32] + ], ], ) def test_listwise_max( diff --git a/tests/kamae/spark/transformers/test_list_max.py b/tests/kamae/spark/transformers/test_list_max.py index 9f927a72..98007969 100644 --- a/tests/kamae/spark/transformers/test_list_max.py +++ b/tests/kamae/spark/transformers/test_list_max.py @@ -664,6 +664,152 @@ def test_spark_mean_transform_with_segmentation( "double", "float", ), + # Integer values with segmentation + ( + 3, + tf.constant( + [ + [1], + [1], + [1], + [2], + [2], + [2], + ], + dtype=tf.int64, + ), + [ + # values + tf.constant( + [ + [1], + [1], + [4], + [5], + [5], + [20], + ], + dtype=tf.int64, + ), + # segment + tf.constant( + [ + [1], + [1], + [2], + [1], + [1], + [2], + ], + dtype=tf.int64, + ), + ], + None, + True, + None, + "bigint", + "bigint", + ), + # Integer values with segmentation & filter + ( + 3, + tf.constant( + [ + [1], + [1], + [1], + [2], + [2], + [2], + ], + dtype=tf.int64, + ), + [ + # values + tf.constant( + [ + [1], + [1], + [4], + [5], + [5], + [20], + ], + dtype=tf.int64, + ), + # segment + tf.constant( + [ + [1], + [1], + [2], + [1], + [1], + [2], + ], + dtype=tf.int64, + ), + ], + 2, + True, + None, + "bigint", + "bigint", + ), + # Remaining integer widths, with segmentation & filter, to check + # Spark/Keras parity for every integer dtype the transformer accepts. + *[ + ( + 3, + tf.constant( + [ + [1], + [1], + [1], + [2], + [2], + [2], + ], + dtype=tf_dtype, + ), + [ + # values + tf.constant( + [ + [1], + [1], + [4], + [5], + [5], + [20], + ], + dtype=tf_dtype, + ), + # segment + tf.constant( + [ + [1], + [1], + [2], + [1], + [1], + [2], + ], + dtype=tf_dtype, + ), + ], + 2, + True, + None, + spark_dtype, + spark_dtype, + ) + for tf_dtype, spark_dtype in [ + (tf.int8, "tinyint"), + (tf.int16, "smallint"), + (tf.int32, "int"), + ] + ], ], ) def test_list_max_transform_spark_tf_parity( From cc4dfd8d3226c10fe028ed676788e9658504dbdc Mon Sep 17 00:00:00 2001 From: Abhishek Das <90371699+dummy1cx@users.noreply.github.com> Date: Wed, 12 Aug 2026 12:59:49 +0100 Subject: [PATCH 2/3] fix: preserve float precision when building ListMax fill value --- src/kamae/keras/tensorflow/layers/list_max.py | 10 ++++-- .../keras/tensorflow/layers/test_list_max.py | 34 +++++++++++++++++++ 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/src/kamae/keras/tensorflow/layers/list_max.py b/src/kamae/keras/tensorflow/layers/list_max.py index ecfacbd3..be213a45 100644 --- a/src/kamae/keras/tensorflow/layers/list_max.py +++ b/src/kamae/keras/tensorflow/layers/list_max.py @@ -171,9 +171,13 @@ def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: listwise_max = tf.broadcast_to(listwise_max, output_shape) if self.min_filter_value is not None: - # tf.cast rather than tf.constant(..., dtype=...): nan_fill_value is a - # Python float, which cannot be converted directly to an integer dtype. - fill_val = tf.cast(self.nan_fill_value, listwise_max.dtype) + # 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_max.dtype.as_numpy_dtype(self.nan_fill_value), + dtype=listwise_max.dtype, + ) listwise_max = tf.where(listwise_max != neg_inf, listwise_max, fill_val) return listwise_max diff --git a/tests/kamae/keras/tensorflow/layers/test_list_max.py b/tests/kamae/keras/tensorflow/layers/test_list_max.py index 6fb48295..5c4c9a47 100644 --- a/tests/kamae/keras/tensorflow/layers/test_list_max.py +++ b/tests/kamae/keras/tensorflow/layers/test_list_max.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import numpy as np import pytest import tensorflow as tf @@ -791,3 +792,36 @@ def test_listwise_max( 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( + "dtype, nan_fill_value, expected_fill", + [ + (tf.float64, 0.1, 0.1), + (tf.float64, 123.456, 123.456), + (tf.float32, 0.1, np.float32(0.1)), + (tf.int64, 7.0, 7), + (tf.int32, 7.0, 7), + ], + ) + def test_listwise_max_nan_fill_value_dtype( + self, dtype, nan_fill_value, expected_fill + ): + """The fill value applied to a segment emptied by min_filter_value must + keep full precision on float dtypes and be usable on integer dtypes.""" + # given, a segment whose values are all removed by the filter + values = tf.constant([[[1], [1], [9]]], dtype=dtype) + segments = tf.constant([[[1], [2], [2]]], dtype=dtype) + layer = ListMaxLayer( + name="listwise_max_nan_fill_test", + min_filter_value=5, + with_segment=True, + nan_fill_value=nan_fill_value, + input_dtype=dtype.name, + output_dtype=dtype.name, + ) + # when + output_tensor = layer([values, segments]) + # then + assert ( + output_tensor.numpy().flatten()[0] == expected_fill + ), "Emptied segment was not filled with the exact nan_fill_value" From 6770121c5da54e2fd25d299958d24e4c78517405 Mon Sep 17 00:00:00 2001 From: Abhishek Das <90371699+dummy1cx@users.noreply.github.com> Date: Thu, 27 Aug 2026 21:47:23 +0100 Subject: [PATCH 3/3] fix: detect emptied ListMax segments from the filter mask The layer decided whether a filter had emptied a segment by testing the reduction result against dtype.min, the neutral element it substitutes for filtered entries. That reads as a sentinel, but on the narrow integer dtypes it is ordinary data: a segment whose genuine maximum is -128 on int8 was indistinguishable from one the filter had emptied, and the real value was overwritten with nanFillValue. Track which entries survive the filter in a mask and reduce that alongside the values, so emptiness is answered directly instead of inferred from the result. This mirrors the approach taken in ListSum. Comparing an integer tensor against the float min_filter_value also raised a TypeError, so integer thresholds are now narrowed explicitly. Rounding up leaves >= unchanged on integers, and a threshold beyond the bounds of the dtype is resolved without a cast, which would otherwise wrap around and silently invert the comparison. The shared min_filter_mask helper holds this logic so ListSum can adopt it. Co-authored-by: Cursor --- src/kamae/keras/tensorflow/layers/list_max.py | 37 +++++++-- src/kamae/keras/tensorflow/utils/__init__.py | 7 +- .../keras/tensorflow/utils/list_utils.py | 26 ++++++ .../keras/tensorflow/layers/test_list_max.py | 81 +++++++++++++++++++ .../kamae/keras/tensorflow/test_list_utils.py | 51 +++++++++++- .../kamae/spark/transformers/test_list_max.py | 58 +++++++++++++ 6 files changed, 251 insertions(+), 9 deletions(-) diff --git a/src/kamae/keras/tensorflow/layers/list_max.py b/src/kamae/keras/tensorflow/layers/list_max.py index be213a45..161021b9 100644 --- a/src/kamae/keras/tensorflow/layers/list_max.py +++ b/src/kamae/keras/tensorflow/layers/list_max.py @@ -22,7 +22,11 @@ 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.list_utils import ( + get_top_n, + min_filter_mask, + segmented_operation, +) from kamae.keras.tensorflow.utils.transform_utils import map_fn_w_axis @@ -149,11 +153,12 @@ def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: # Apply the mask to filter out elements less than or equal to the threshold if self.min_filter_value is not None: - mask = tf.greater_equal(val_tensor, self.min_filter_value) - neg_inf = val_tensor.dtype.min - val_tensor = tf.where(mask, val_tensor, neg_inf) - else: - val_tensor = val_tensor + mask = min_filter_mask(val_tensor, self.min_filter_value) + # The dtype minimum is only a neutral element for the reduction here: a + # real value equal to it still wins the max, so substituting it for the + # filtered entries cannot change the result. + val_tensor = tf.where(mask, val_tensor, val_tensor.dtype.min) + kept = tf.cast(mask, tf.int32) # Apply segmented calculation if self.with_segment: @@ -171,6 +176,24 @@ def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: listwise_max = tf.broadcast_to(listwise_max, output_shape) if self.min_filter_value is not None: + # Whether anything survived the filter is read from the mask rather than + # by testing the result against the sentinel. dtype.min is a legitimate + # value for the narrow integer dtypes (-128 for int8), so a sentinel test + # would silently overwrite real data with nan_fill_value. + 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. @@ -178,7 +201,7 @@ def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: listwise_max.dtype.as_numpy_dtype(self.nan_fill_value), dtype=listwise_max.dtype, ) - listwise_max = tf.where(listwise_max != neg_inf, listwise_max, fill_val) + listwise_max = tf.where(any_kept > 0, listwise_max, fill_val) return listwise_max diff --git a/src/kamae/keras/tensorflow/utils/__init__.py b/src/kamae/keras/tensorflow/utils/__init__.py index ca2949f6..880a6141 100644 --- a/src/kamae/keras/tensorflow/utils/__init__.py +++ b/src/kamae/keras/tensorflow/utils/__init__.py @@ -37,5 +37,10 @@ datetime_year, unix_timestamp_to_datetime, ) -from .list_utils import get_top_n, listify_tensors, segmented_operation # noqa: F401 +from .list_utils import ( # noqa: F401 + get_top_n, + listify_tensors, + min_filter_mask, + segmented_operation, +) from .transform_utils import map_fn_w_axis # noqa: F401 diff --git a/src/kamae/keras/tensorflow/utils/list_utils.py b/src/kamae/keras/tensorflow/utils/list_utils.py index 264fd028..9f5f48da 100644 --- a/src/kamae/keras/tensorflow/utils/list_utils.py +++ b/src/kamae/keras/tensorflow/utils/list_utils.py @@ -11,6 +11,7 @@ # 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 math from typing import Any, Callable, List, Union import numpy as np @@ -70,6 +71,31 @@ def get_top_n( ) +def min_filter_mask(val_tensor: Tensor, min_filter_value: float) -> Tensor: + """ + Build the mask of the entries that meet a minimum value threshold. + + TensorFlow will not compare an integer tensor against a float threshold, so for + those the threshold is rounded up, which is equivalent under >= on integers. + A threshold beyond the bounds of the dtype cannot be narrowed at all, since the + cast would wrap around silently, so those two cases are answered directly. + Floats need neither adjustment because they saturate to +/-inf rather than wrap. + + :param val_tensor: Value tensor to filter. + :param min_filter_value: Minimum value an entry must meet to be kept. + :returns: Boolean tensor that is True wherever the entry is kept. + """ + if not val_tensor.dtype.is_integer: + return tf.greater_equal(val_tensor, min_filter_value) + + threshold = math.ceil(min_filter_value) + if threshold <= val_tensor.dtype.min: + return tf.ones_like(val_tensor, dtype=tf.bool) + if threshold > val_tensor.dtype.max: + return tf.zeros_like(val_tensor, dtype=tf.bool) + return tf.greater_equal(val_tensor, tf.cast(threshold, val_tensor.dtype)) + + def listify_tensors(x: Union[tf.Tensor, np.ndarray, List[Any]]) -> List[Any]: """ Converts any tensors or numpy arrays to lists for config serialization. diff --git a/tests/kamae/keras/tensorflow/layers/test_list_max.py b/tests/kamae/keras/tensorflow/layers/test_list_max.py index 5c4c9a47..6cab99a2 100644 --- a/tests/kamae/keras/tensorflow/layers/test_list_max.py +++ b/tests/kamae/keras/tensorflow/layers/test_list_max.py @@ -825,3 +825,84 @@ def test_listwise_max_nan_fill_value_dtype( assert ( output_tensor.numpy().flatten()[0] == expected_fill ), "Emptied segment was not filled with the exact nan_fill_value" + + @pytest.mark.parametrize("with_segment", [False, True]) + @pytest.mark.parametrize("dtype", [tf.int8, tf.int16, tf.int32, tf.int64]) + def test_listwise_max_dtype_minimum_is_a_valid_value(self, dtype, with_segment): + """The dtype minimum is ordinary data on the narrow integer types, -128 for + int8, so a list whose genuine maximum is that value must be returned as-is + rather than mistaken for a list emptied by the filter.""" + # given, a list whose real maximum is the dtype minimum, kept by the filter + dtype_min = dtype.min + values = tf.constant([[[dtype_min], [dtype_min], [dtype_min]]], dtype=dtype) + layer = ListMaxLayer( + name="listwise_max_dtype_min_test", + min_filter_value=dtype_min, + nan_fill_value=0.0, + with_segment=with_segment, + input_dtype=dtype.name, + output_dtype=dtype.name, + ) + inputs = values + if with_segment: + inputs = [values, tf.constant([[[1], [1], [2]]], dtype=dtype)] + # when + output_tensor = layer(inputs) + # then + tf.debugging.assert_equal( + output_tensor, + tf.constant([[[dtype_min], [dtype_min], [dtype_min]]], dtype=dtype), + ) + + @pytest.mark.parametrize( + "min_filter_value, expected", + [ + (-200.0, 3), + (-129.0, 3), + (-128.0, 3), + (127.0, 0), + (128.0, 0), + (200.0, 0), + ], + ) + def test_listwise_max_min_filter_value_outside_integer_dtype_range( + self, min_filter_value, expected + ): + """Narrowing a threshold that falls outside the value dtype wraps around, so + the out-of-range cases are decided directly: below the dtype minimum keeps + every value, above the dtype maximum keeps none of them.""" + # given, int8 values of 1, 2 and 3 against thresholds beyond the int8 bounds + values = tf.constant([[[1], [2], [3]]], dtype=tf.int8) + layer = ListMaxLayer( + name="listwise_max_out_of_range_filter_test", + min_filter_value=min_filter_value, + nan_fill_value=0.0, + input_dtype="int8", + output_dtype="int8", + ) + # when + output_tensor = layer(values) + # then + tf.debugging.assert_equal( + output_tensor, tf.constant([[[expected]] * 3], dtype=tf.int8) + ) + + @pytest.mark.parametrize("dtype", [tf.int8, tf.int16, tf.int32, tf.int64]) + def test_listwise_max_float_min_filter_value_on_integer_input(self, dtype): + """min_filter_value is declared as a float, so passing one against an + integer input column must work rather than fail the dtype comparison.""" + # given + values = tf.constant([[[1], [2], [3]]], dtype=dtype) + layer = ListMaxLayer( + name="listwise_max_float_filter_test", + min_filter_value=0.0, + nan_fill_value=-1.0, + input_dtype=dtype.name, + output_dtype=dtype.name, + ) + # when + output_tensor = layer(values) + # then + tf.debugging.assert_equal( + output_tensor, tf.constant([[[3], [3], [3]]], dtype=dtype) + ) diff --git a/tests/kamae/keras/tensorflow/test_list_utils.py b/tests/kamae/keras/tensorflow/test_list_utils.py index 71a4a06e..8636e1f1 100644 --- a/tests/kamae/keras/tensorflow/test_list_utils.py +++ b/tests/kamae/keras/tensorflow/test_list_utils.py @@ -15,7 +15,7 @@ import pytest import tensorflow as tf -from kamae.keras.tensorflow.utils import get_top_n +from kamae.keras.tensorflow.utils import get_top_n, min_filter_mask class TestGetTopN: @@ -69,3 +69,52 @@ def test_get_top_with_batch_1( ): output = get_top_n(val_tensor, axis, sort_tensor, top_n, sort_order) tf.debugging.assert_equal(output, expected_output) + + +class TestMinFilterMask: + @pytest.mark.parametrize( + "dtype, min_filter_value, expected", + [ + # Integer values cannot be compared against a float threshold directly, + # and rounding up leaves >= unchanged. + (tf.int32, 0.0, [False, True, True]), + (tf.int32, 0.5, [False, False, True]), + (tf.int32, 1.0, [False, False, True]), + (tf.int32, -0.5, [False, True, True]), + # Floats keep the fractional threshold as given. + (tf.float32, 0.5, [False, False, True]), + (tf.float64, -0.5, [False, True, True]), + ], + ) + def test_min_filter_mask_threshold(self, dtype, min_filter_value, expected): + val_tensor = tf.constant([-1, 0, 1], dtype=dtype) + output = min_filter_mask(val_tensor, min_filter_value) + tf.debugging.assert_equal(output, tf.constant(expected, dtype=tf.bool)) + + @pytest.mark.parametrize( + "min_filter_value, expected", + [ + # Narrowing these to int8 would wrap around, so they are answered without + # a cast: below the dtype minimum keeps everything, above the dtype + # maximum keeps nothing. + (-200.0, [True, True, True]), + (-129.0, [True, True, True]), + (-128.0, [True, True, True]), + (127.0, [False, False, False]), + (128.0, [False, False, False]), + (200.0, [False, False, False]), + ], + ) + def test_min_filter_mask_threshold_outside_dtype_range( + self, min_filter_value, expected + ): + val_tensor = tf.constant([1, 2, 3], dtype=tf.int8) + output = min_filter_mask(val_tensor, min_filter_value) + tf.debugging.assert_equal(output, tf.constant(expected, dtype=tf.bool)) + + def test_min_filter_mask_keeps_dtype_minimum(self): + """The dtype minimum is a real value, not a sentinel, so a threshold at the + bound must keep it.""" + val_tensor = tf.constant([-128, -127], dtype=tf.int8) + output = min_filter_mask(val_tensor, -128.0) + tf.debugging.assert_equal(output, tf.constant([True, True], dtype=tf.bool)) diff --git a/tests/kamae/spark/transformers/test_list_max.py b/tests/kamae/spark/transformers/test_list_max.py index 98007969..812ada45 100644 --- a/tests/kamae/spark/transformers/test_list_max.py +++ b/tests/kamae/spark/transformers/test_list_max.py @@ -880,3 +880,61 @@ def test_list_max_transform_spark_tf_parity( decimal=6, err_msg="Spark and Tensorflow transform outputs are not equal", ) + + @pytest.mark.parametrize( + "spark_dtype, tf_dtype", + [ + ("tinyint", tf.int8), + ("smallint", tf.int16), + ("int", tf.int32), + ("bigint", tf.int64), + ], + ) + def test_list_max_transform_spark_tf_parity_at_dtype_minimum( + self, spark_session, spark_dtype, tf_dtype + ): + """Spark only returns null once a filter empties a window, so a query whose + genuine maximum is the dtype minimum keeps that value. Tensorflow has to + agree rather than read the reduction result back as an emptied query.""" + # given, two queries of two items filtered at the dtype minimum, so nothing + # is removed. The first has the dtype minimum as its real maximum, the + # second keeps an ordinary value as a control. + dtype_min = tf_dtype.min + list_size = 2 + query_ids = [1, 1, 2, 2] + values = [dtype_min, dtype_min, dtype_min, 5] + transformer = ListMaxTransformer( + inputCol="input0", + outputCol="output", + inputDtype=spark_dtype, + outputDtype=spark_dtype, + queryIdCol="search_id", + minFilterValue=dtype_min, + nanFillValue=0.0, + ) + spark_df = spark_session.createDataFrame( + list(zip(query_ids, values)), ["search_id", "input0"] + ) + # when + spark_values = ( + transformer.transform(spark_df) + .select("output") + .rdd.map(lambda r: r[0]) + .collect() + ) + input_tensor = tf.reshape( + tf.constant(values, dtype=tf_dtype), (-1, list_size, 1) + ) + tensorflow_values = np.reshape( + transformer.get_keras_layer()(input_tensor).numpy(), -1 + ).tolist() + # then + assert spark_values == [dtype_min, dtype_min, 5, 5], ( + "The dtype minimum should be kept as a real value, and nanFillValue " + "should not appear when nothing is filtered" + ) + np.testing.assert_equal( + spark_values, + tensorflow_values, + err_msg="Spark and Tensorflow transform outputs are not equal", + )