Skip to content

feat: add ListSum transformer and layer - #73

Open
dummy1cx wants to merge 2 commits into
ExpediaGroup:mainfrom
dummy1cx:kamae-listsum
Open

feat: add ListSum transformer and layer#73
dummy1cx wants to merge 2 commits into
ExpediaGroup:mainfrom
dummy1cx:kamae-listsum

Conversation

@dummy1cx

Copy link
Copy Markdown

Description

Adds ListSumTransformer (Spark) and ListSumLayer (Keras), extending the listwise
aggregation family alongside ListMax. The implementation mirrors ListMax in
structure, parameters and naming:

  • topN + sortOrder — sum over the top N items only, to approximate production
    statistics where the full list is not available at serving time.
  • withSegment — sum within segments of a second column rather than across the whole
    query. For example, the summed price of hotels within each star rating in a query.
  • minFilterValue — exclude padded values before aggregating.

Motivation: multi-level ranking models need listwise sums aggregated per segment, which
currently has no equivalent in the library.

Keras Layer Checklist

Verify that:

  • The new Keras layer extends BaseLayer
  • The _call method has been implemented in the new layer.
  • The compatible_dtypes property is defined in the new layer.
  • The new layer is decorated with @tf.keras.utils.register_keras_serializable(package=kamae.__name__).
  • The new layer takes a name, input_dtype, and output_dtype as arguments to the constructor and that this is passed to the super constructor.
  • The Keras layer is serializable. I have implemented the get_config method.
  • There are unit tests of the new layer.
  • There is a specific test of layer serialisation added.
  • The new layer is imported in the __init__.py file in the layers directory.

Spark Transformer/Estimator Checklist

Verify that:

  • The new Spark Transformer extends BaseTransformer.
  • If the new transform needs a fit method, a Spark Estimator has been implemented that extends BaseEstimator. (N/A — no fit required.)
  • The instructions in the above docs page have been followed for the __init__ and setParams methods.
  • The transformer uses one of the input/output mixin classes from base.py.
  • If the new transformer requires more parameters that would need to be serialised to the Spark ML pipeline, there is an implemented parameter class. (Reuses the existing ListwiseStatisticsParams and NanFillValueParams.)
  • The compatible_dtypes property has been implemented to specify the input/output data types that my transformer/estimator supports.
  • A Keras subclassed layer is returned in the transformer's get_tf_layer method.
  • There are unit tests of the new transform. In particular, there are parity tests between the Spark and Keras implementations.
  • The new transformer/estimator is imported in the __init__.py file in the transformers directory.

Finally, please verify that:

  • There is a new entry (alphabetical order) in the README table describing the new layer/transformer

Notes for reviewers

Three decisions I would like a second opinion on:

  1. minFilterValue semantics. Values below the threshold contribute 0 to the sum
    rather than being dropped from the segment. The masking uses tf.where with
    tf.zeros_like rather than a NaN-based approach, deliberately avoiding float-only ops
    so that integer value columns work. Happy to change if you would prefer different
    semantics here.

  2. string in compatible_dtypes. This is not to sum strings. BaseLayer
    validates the dtype of every input tensor, and under withSegment=True the second
    input is a segment identifier which is commonly a string. Without string in the
    list, string segment keys are rejected.

  3. Known parity gap, deliberately not addressed here. If a segment becomes empty
    after minFilterValue filtering, the Spark side applies nanFillValue via fillna
    while the Keras layer returns 0.0. Reaching this requires both minFilterValue and
    a non-zero nanFillValue, which no current configuration sets. Closing it properly
    needs map_fn_w_axis to accept more than two tensors, so I have kept it out of this
    PR to keep the change reviewable. Happy to fold it in here instead if you would rather
    not merge a known gap.

Testing

Unit tests for both implementations cover the base case, minFilterValue, topN with
and without a filter, topN greater than the list size, segmentation, segmentation with
multiple features, segmentation with string segment IDs, and segmentation combined with
minFilterValue, plus the ValueError paths. The Spark tests include Spark/TensorFlow
parity cases. The layer is also registered in the JIT-compatibility and serialisation
meta-test suites.

@dummy1cx
dummy1cx requested a review from a team as a code owner August 11, 2026 22:51
@dummy1cx
dummy1cx requested review from ddonghi and jacobjwood August 11, 2026 22:51
@jamesdshinner
jamesdshinner self-requested a review August 12, 2026 10:16
georyetti
georyetti previously approved these changes Aug 25, 2026

@georyetti georyetti left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@jamesdshinner jamesdshinner left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Couple of pedantic style comments only

Comment on lines +32 to +48
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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's ensure consistent indentation here

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

/I think you have two separate paragraphs flowing together

ListwiseStatisticsParams,
NanFillValueParams,
):
"""

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same point on indentation and formatting here

self.top_n = top_n
self.sort_order = sort_order
self.min_filter_value = min_filter_value
self.nan_fill_value = nan_fill_value

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this used anywhere? let's add a test in parity check that with a non-default nan_fill_value

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants