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

Sbranikas · 2026-09-22T13:29:54Z

+        output_col = single_input_single_output_array_transform(
+            input_col=array_col,
+            input_col_datatype=self.get_column_datatype(df, df.columns[0]),
+            func=lambda x: F.array_contains(x, value_col),


Just wanted to doulbe-check, array_contains can return NULL too, not just True/False, while for the Keras side ops.any(ops.equal(...)) always gives a boolean. if the array contains even a single NULL, array contains will result NULL overall. Is this the desired behaviour, to propagate the NULL? If not we can maybe wrap it with an explcit check?

Good catch, wrapped in F.coalesce(..., F.lit(False))

Sbranikas · 2026-09-22T13:42:51Z

+        else:
+            raise ValueError(f"Expected 2 inputs, got {len(inputs)} inputs instead")
+
+        return ops.any(ops.equal(array, value), axis=self.axis, keepdims=self.keepdims)


Is the equality check here identical to what array_contains does on the Spark side? Just trying to understand whether there's any chance of a mismatch on exact equality.

I see above (line 111) the cast only runs for the constant branch, so I wasn't sure what happens when the value comes in as a separate column with a different dtype.

Both use exact equality and inputDtype casts the array and value to the same dtype

Thanks Mantas!

The inputDtype point holds when it's defined, but it defaults to None, and then nothing aligns the array and the constant, so each side reconciles float vs double its own way:

Spark widens the array to double, and Keras on the opposite direction narrows the constant to float32. It might be pedantic, since this only matters for constants that aren't exactly representable in float32, like 0.1: array_contains([0.1f], 0.1) → false in Spark (0.10000000149 ≠ 0.1), but true in Keras. (I checked the Spark side locally with claude.)

Could we cast the literal to the array's element type on the Spark side (or mirror Spark's widening in Keras), and add a parity test with inputDtype=None and a constant like 0.1 (that float32 is not enough to describe it) on a float array? The current parity tests only use integer arrays and integer constants (5.0 and 0.0), so the float path isn't covered.

Nice, thank you! Made it so that behaviour is identical to #L111

Sbranikas

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 previously approved these changes Sep 15, 2026

View reviewed changes

Sbranikas reviewed Sep 22, 2026

View reviewed changes

fix: apply comments

13edfb3

mantasu dismissed georyetti’s stale review via 13edfb3 September 28, 2026 15:28

fix: spark dtypes

6dc9372

Sbranikas approved these changes Sep 28, 2026

View reviewed changes

georyetti approved these changes Sep 30, 2026

View reviewed changes

mantasu merged commit abce236 into ExpediaGroup:main Oct 1, 2026
28 checks passed

mantasu deleted the feat/array-contains branch October 1, 2026 07:40

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!

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!

Sbranikas left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants