feat: add ListSum transformer and layer - #73
Open
dummy1cx wants to merge 2 commits into
Open
Conversation
jamesdshinner
requested changes
Aug 27, 2026
jamesdshinner
left a comment
Contributor
There was a problem hiding this comment.
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. |
Contributor
There was a problem hiding this comment.
Let's ensure consistent indentation here
Contributor
There was a problem hiding this comment.
/I think you have two separate paragraphs flowing together
| ListwiseStatisticsParams, | ||
| NanFillValueParams, | ||
| ): | ||
| """ |
Contributor
There was a problem hiding this comment.
Same point on indentation and formatting here
mandrecki
requested changes
Aug 27, 2026
| self.top_n = top_n | ||
| self.sort_order = sort_order | ||
| self.min_filter_value = min_filter_value | ||
| self.nan_fill_value = nan_fill_value |
Contributor
There was a problem hiding this comment.
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.
2 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds
ListSumTransformer(Spark) andListSumLayer(Keras), extending the listwiseaggregation family alongside
ListMax. The implementation mirrorsListMaxinstructure, parameters and naming:
topN+sortOrder— sum over the top N items only, to approximate productionstatistics where the full list is not available at serving time.
withSegment— sum within segments of a second column rather than across the wholequery. 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:
_callmethod has been implemented in the new layer.compatible_dtypesproperty is defined in the new layer.@tf.keras.utils.register_keras_serializable(package=kamae.__name__).name,input_dtype, andoutput_dtypeas arguments to the constructor and that this is passed to the super constructor.get_configmethod.__init__.pyfile in thelayersdirectory.Spark Transformer/Estimator Checklist
Verify that:
__init__andsetParamsmethods.ListwiseStatisticsParamsandNanFillValueParams.)compatible_dtypesproperty has been implemented to specify the input/output data types that my transformer/estimator supports.get_tf_layermethod.__init__.pyfile in thetransformersdirectory.Finally, please verify that:
Notes for reviewers
Three decisions I would like a second opinion on:
minFilterValuesemantics. Values below the threshold contribute0to the sumrather than being dropped from the segment. The masking uses
tf.wherewithtf.zeros_likerather than a NaN-based approach, deliberately avoiding float-only opsso that integer value columns work. Happy to change if you would prefer different
semantics here.
stringincompatible_dtypes. This is not to sum strings.BaseLayervalidates the dtype of every input tensor, and under
withSegment=Truethe secondinput is a segment identifier which is commonly a string. Without
stringin thelist, string segment keys are rejected.
Known parity gap, deliberately not addressed here. If a segment becomes empty
after
minFilterValuefiltering, the Spark side appliesnanFillValueviafillnawhile the Keras layer returns
0.0. Reaching this requires bothminFilterValueanda non-zero
nanFillValue, which no current configuration sets. Closing it properlyneeds
map_fn_w_axisto accept more than two tensors, so I have kept it out of thisPR 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,topNwithand without a filter,
topNgreater than the list size, segmentation, segmentation withmultiple features, segmentation with string segment IDs, and segmentation combined with
minFilterValue, plus theValueErrorpaths. The Spark tests include Spark/TensorFlowparity cases. The layer is also registered in the JIT-compatibility and serialisation
meta-test suites.