feat: array contains by mantasu · Pull Request #67 · ExpediaGroup/kamae

mantasu · 2026-07-31T12:22:15Z

Description

Adds an ArrayContains transformer and paired Keras layer that check whether a scalar value is contained in an array feature, outputting 1.0 if present and 0.0 otherwise.

Keras Layer Checklist

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 here.
The new layer is imported in the init.py file in the layers directory.

Spark Transformer/Estimator Checklist

The new Spark Transformer extends BaseTransformer.
If the new transform needs a fit method, a Spark Estimator has been implemented that extends BaseEstimator.
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 a implemented parameter class by extending the Params class here.
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/estimators directory.

Readme Checklist

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

mantasu · 2026-07-31T13:01:33Z

Could I please be also given write access so I wouldn't need to create the fork? 🙏

georyetti

Some comments:

Keras layer can contain more params to make it more flexible. We can set these as fixed from Spark side
I think this can return booleans and the user can set output dtypes for auto casting to floats
Can we use the Spark array helper functions to support nested arrays

georyetti · 2026-08-20T08:58:36Z

+            raise ValueError(f"Expected 2 inputs, got {len(inputs)} inputs instead.")
+
+        array, value = inputs
+        any_match = ops.any(ops.equal(array, value), axis=-1, keepdims=True)


Can we make axis a param of this layer? Then its more useful for a user that just reuses keras layers directly? We can set axis=-1 from Spark side

Same for keepdims too pls

georyetti · 2026-08-20T09:00:21Z

+        output_col = (
+            F.when(F.array_contains(F.col(arr_c), F.col(val_c)), F.lit(1.0))
+            .otherwise(F.lit(0.0))
+            .cast(DoubleType())
+        )


Can we reuse the helper methods we have for array computation here? Then it would support nested arrays also in Spark.

georyetti · 2026-08-20T09:01:50Z

+            "int8",
+            "uint8",
+            "int16",
+            "uint16",
+            "int32",
+            "uint32",
+            "int64",
+            "uint64",


Does this work for int values? Or just int arrays? If so can we add a test please. All tests currently are just with float values

Yeah should work for both arrays and scalars, added a test

georyetti · 2026-08-20T09:03:31Z

+
+        array, value = inputs
+        any_match = ops.any(ops.equal(array, value), axis=-1, keepdims=True)
+        return ops.cast(any_match, "float32")


Is it expected that array contains returns a float32? I would expect a boolean. The user can set output_dtype="float32" to get this behaviour so I would say we should preserve booleans here

georyetti

Two comments:

Pedantic one about using helpers, you have used some more internal ones, we have an easier wrapper for this.
Do we want to allow a pattern where the scalar value is just a normal constant. Like we do with the Divide/Multiply transformers?

georyetti · 2026-09-09T13:21:02Z

+        nesting_level, elem_t = get_array_nesting_level_and_element_dtype(arr_t)
+        if not isinstance(elem_t, _NUMERIC_TYPES):
+            raise TypeError(f"arrayCol '{arr_c}' element must be numeric, got {elem_t}")
+
+        if not isinstance(val_t, _NUMERIC_TYPES):
+            raise TypeError(f"valueCol '{val_c}' must be numeric, got {val_t}")
+
+        # Apply array_contains at the innermost level
+        contains_func = nested_transform(
+            func=lambda x: F.array_contains(x, F.col(val_c)),
+            nest_level=nesting_level - 1,
+        )
+        output_col = contains_func(F.col(arr_c))


We have a helper for this that makes it a bit simpler, can we use this? single_input_single_output_array_transform
Also there is no need to check for numeric types as this is done by the BaseTransformer using compatible_dtypes

georyetti · 2026-09-09T13:22:22Z

+        inputCols: Optional[List[str]] = None,
+        outputCol: Optional[str] = None,
+        inputDtype: Optional[str] = None,
+        outputDtype: Optional[str] = None,
+        layerName: Optional[str] = None,
+    ) -> None:
+        """
+        Initializes an ArrayContainsTransformer transformer.
+
+        :param inputCols: Input column names, given as `[arrayCol, valueCol]`.


Do we want to allow a single input col and a constant? So we can check if array_contains(x, -1) for example?

Fair idea, generalised it

georyetti

LGTM

feat: array contains

5e9e377

mantasu requested a review from a team as a code owner July 31, 2026 12:22

mantasu requested review from ddonghi and jacobjwood July 31, 2026 12:22

mantasu added 2 commits July 31, 2026 13:41

fix: readability

bae9cdc

fi: undo formatting

7d3af7d

georyetti requested changes Aug 20, 2026

View reviewed changes

mantasu and others added 2 commits September 3, 2026 08:47

fix: apply comments

477d42c

Merge branch 'main' into feat/array-contains

8133b18

mantasu force-pushed the feat/array-contains branch from 5583b8b to 8133b18 Compare September 3, 2026 08:04

georyetti requested changes Sep 9, 2026

View reviewed changes

fix: apply comments

eaf66ea

georyetti approved these changes Sep 15, 2026

View reviewed changes

Uh oh!

Conversation

mantasu commented Jul 31, 2026

Description

Keras Layer Checklist

Spark Transformer/Estimator Checklist

Readme Checklist

Uh oh!

mantasu commented Jul 31, 2026

Uh oh!

georyetti left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

georyetti left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

georyetti left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants