feat: opt-in pipeline fit optimisations, native bucketize, and sampled fitting support with caching by ConorWorthington · Pull Request #69 · ExpediaGroup/kamae

ConorWorthington · 2026-08-05T13:56:38Z

Description

Provide a short description of the PR changes.

The below checklists come from the docs page on adding new transformers here

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 here.
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.
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.

Finally, please verify that:

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

Plans are too slow to materialise - this is our attempt to speed it up

Wrap the moments aggregation in StandardScale, SingleFeatureArrayStandardScale and ConditionalStandardScale estimators in a guarded persist/unpersist so the array-size probe and the aggregation reuse a materialised result instead of re-scanning the upstream lineage twice. Repair the incomplete persist edit in ConditionalStandardScale._fit. Add checkpointInterval / pruneInputColumns coverage to the pipeline tests and a checkpoint directory to the spark_session fixture. Surface estimator fit errors as RuntimeError chained from the original exception. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Merge latest version

georyetti

Mostly minor changes to pipeline logic. Also bucketize changes I assume should not be here. Lastly you need to run a uv lock.

georyetti · 2026-08-12T09:19:03Z

        """
-        super().__init__(stages=stages)
+        kwargs = self._input_kwargs
+        super().__init__()


The new super call drops the stages=stages, was this intended?

georyetti · 2026-08-12T09:19:31Z

+        kwargs = self._input_kwargs
+        super().__init__()
+        self._setDefault(
+            checkpointInterval=0,


Default to None instead of 0?

georyetti · 2026-08-12T09:25:26Z

+            for param in stage.params:
+                if not (param.name.endswith("Col") or param.name.endswith("Cols")):
+                    continue
+                if not stage.isDefined(param):
+                    continue
+                value = stage.getOrDefault(param)
+                if isinstance(value, str):
+                    required_input_columns.add(value)
+                elif isinstance(value, (list, tuple)):
+                    required_input_columns.update(
+                        item for item in value if isinstance(item, str)
+                    )


This feels more complex than I think it needs to be. Two things:

Why do we need to iterate through the stage.params at all if we have already added the inputs using get_layer_inputs_outputs

Even if we do need, we can just check for presence of inputCol or inputCols as a stage always has strictly one of these defined. if stage.hasParam("inputCol") and stage.isDefined("inputCol"): value = stage.getInputCol() and after check input cols. But this is really what get_layer_inputs_outputs does.

You're right that inputCol/inputCols are already covered by get_layer_inputs_outputs but the sweep isn't for those.

It's there for the other column params a stage reads that aren't its input col: maskCols/relevanceCol on ConditionalStandardScaleEstimator, and queryIdCol on the listwise transformers. Those get read at fit time, so if we don't keep them, pruneInputColumns=True drops them and the fit blows up with "column not found". There are regression tests covering exactly that.

If the suffix-matching feels too hacky, I'm happy to swap it for a small get_fit_input_columns() hook on the handful of stages that need it instead. Let me know what you'd prefer.

Ok makes sense.

georyetti · 2026-08-12T09:26:56Z

+        """
+        required_input_columns = self.collect_required_input_columns(stages)
+        columns_to_keep = [c for c in dataset.columns if c in required_input_columns]
+        if columns_to_keep and len(columns_to_keep) < len(dataset.columns):


Pedantic but do we need that second condition? If we are inside this function then we are pruning, and columns_to_keep is always a subset. So I would just check it's not empty and otherwise select

georyetti · 2026-08-12T09:28:00Z

+        fit. 0 (or None) disables checkpointing.
+        :returns: KamaeSparkPipeline object with checkpointInterval set.
+        """
+        return self._set(checkpointInterval=value)


We should error on checkpoint interval being negative here. Personally I think we should error on 0 too and treat None as the no checkpoint behaviour

georyetti · 2026-08-12T09:29:50Z

+        :returns: KamaeSparkPipeline object with params set.
        """
        kwargs = self._input_kwargs
        return self._set(**kwargs)


This set params does not use the setter methods at all. So any validation in them will not be respected when the user passes arguments to the init as opposed to using the setter method. If you check how I defined the setParams for the estimator and transformer you can see I use the setter method.

georyetti · 2026-08-12T11:11:43Z

-                return None
-            # We add 1 because we want to reserve the 0 index for mask/padding.
-            return bisect_right(splits, value) + 1
+        def bucketize(value: Column) -> Column:


Same to assume that all this bucket logic changes to this transformer are not meant to be in this PR?

I was looking for inefficiencies in the library... I was going to patch more but only bucketize was really impacted in a way I could speed up simply. We don't really use it but figured it was a good improvement so may as well include.

Hmmm ok if you want to keep it I will have more comments, will add them then you can decide.

georyetti · 2026-08-12T11:12:50Z

 dependencies = [
    "pyspark>=3.4.0,<4.0.0",
    "pandas>=1.3.4,<3.0.0",
+    "pyarrow>=4.0.0",


Adding new dependencies needs a uv lock pls

Lock updated

- Restore stages=stages in __init__ super call - checkpointInterval defaults to None; reject non-positive via setter - Route setParams through setter methods so validation runs - Drop redundant length check in prune_unused_input_columns - Regenerate uv.lock to include pyarrow (required by pandas_udf) Retains the aux-column sweep in collect_required_input_columns: it is load-bearing for pruning correctness (maskCols/relevanceCol/queryIdCol are not returned by get_layer_inputs_outputs) and defended by regression tests. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

georyetti

Effectively fine, just revert the type hint for bucketize back to tf layer as its still tf only

georyetti · 2026-08-17T13:12:37Z

+            for param in stage.params:
+                if not (param.name.endswith("Col") or param.name.endswith("Cols")):
+                    continue
+                if not stage.isDefined(param):
+                    continue
+                value = stage.getOrDefault(param)
+                if isinstance(value, str):
+                    required_input_columns.add(value)
+                elif isinstance(value, (list, tuple)):
+                    required_input_columns.update(
+                        item for item in value if isinstance(item, str)
+                    )


Ok makes sense.

georyetti · 2026-08-17T13:14:06Z

-                return None
-            # We add 1 because we want to reserve the 0 index for mask/padding.
-            return bisect_right(splits, value) + 1
+        def bucketize(value: Column) -> Column:


Hmmm ok if you want to keep it I will have more comments, will add them then you can decide.

georyetti · 2026-08-17T13:16:38Z

        )

-    def get_keras_layer(self) -> tf.keras.layers.Layer:
+    def get_keras_layer(self) -> keras.layers.Layer:


Why are we making this type hint as keras layer? The layer still uses tf only functions so we should revert this back

…s.Layer BucketizeLayer is TensorFlow-only, so the tf-specific return type is accurate. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…peline Adds a default-off boolean pipeline param that, at the first estimator-fit boundary, projects the working frame to the columns still read downstream and persists (MEMORY_AND_DISK) that narrow frame once, reused by all subsequent estimators. This collapses repeated full scans of a wide input across independent sibling estimators into a single populating scan plus in-RAM reuse. When both cacheIntermediateData and cacheEstimatorInput are enabled, cacheEstimatorInput takes precedence (with a warning) since it is a strictly narrower cache and the intermediate cache would evict it. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

georyetti

Sorry some minor comments, I don't think its particularly blocking so can approve if you have good reason for it.

georyetti · 2026-08-20T08:52:49Z

        if cached_dataset is not None:
            cached_dataset.unpersist()
+        if estimator_input_cache is not None:
+            estimator_input_cache.unpersist()


Got a little confused here as we have 3 dataframes:

dataset

cached_dataset

new_cached

And we set the first 2 to the 3rd, but then only unpersist the 2nd...?

Can we simplify this to just persist dataset?

georyetti · 2026-08-20T08:53:48Z

+                        estimator_input_cache = dataset.select(*keep_columns).persist(
+                            StorageLevel.MEMORY_AND_DISK
+                        )
+                        dataset = estimator_input_cache


Similar here, what is the use of estimator_input_cache here if we just set to dataset and then have to remember to unpersist()? Can we not just reuse the dataset variable and unpersist() just that at the end?

…line Adds a default-None float param (0, 1] that draws a single sample of the input up-front, persists (MEMORY_AND_DISK) and materialises it once, and fits every estimator from that shared sample with each estimator's own sampleFraction temporarily disabled and restored afterwards. This collapses the N independent per-estimator Bernoulli scans of a wide source (which spilled and GC-thrashed at scale) into a single populating scan plus in-RAM reuse. fitSampleSeed makes the sample reproducible. fitSampleFraction is incompatible with cacheIntermediateData/cacheEstimatorInput (which persist frames it is designed to avoid), so enabling it warns and disables them. It only computes correct statistics for sample-robust estimators (mean/std/quantiles); a runtime warning documents that vocabulary builders, min/max scalers and distinct counts need exact statistics. Also addresses review feedback on the cache bookkeeping: since the two cache strategies are mutually exclusive, collapse the separate cached_dataset / estimator_input_cache / new_cached handles into a single persisted_frame kept distinct from dataset (which model.transform reassigns), with one unpersist. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

ConorWorthington · 2026-08-20T17:00:35Z

Pushed a new commit (2e0a0cd) with two things:

1. New opt-in param: fitSampleFraction (+ fitSampleSeed)

This is a fit-time optimisation for the case that motivated the caching work: a pipeline of several independent estimators (each ArrayConcatenate -> ConditionalStandardScaleEstimator sampling internally at ~0.005) where every estimator does its own full Bernoulli scan of the wide source. Persisting the intermediate/estimator-input frame doesn't help there because the frame is too big for the storage-memory budget — it spills (~1.7 TiB) and the executors GC-thrash.

Insight: the fit only needs a sample. So when fitSampleFraction is set to a float in (0, 1], _fit:

draws a single dataset.sample(fraction, seed=fitSampleSeed) up-front, persist(MEMORY_AND_DISK)s it and forces one count() so it materialises exactly once;
runs the whole fit loop against that shared sample;
temporarily forces each estimator's own sampleFraction to None (restored in a finally) so we don't sample-a-sample;
unpersists the sample in the finally.

The returned model is unchanged — transformers still apply to the full dataset at transform time; only fit-time stats come from the sample.

Caveats (also in the param docstring + a runtime warning):

It's incompatible with cacheIntermediateData/cacheEstimatorInput (they persist the frames this option avoids), so enabling it warns and disables them.
It fits all estimators from one shared sample, so it's only correct for sample-robust statistics (mean/std/quantiles, e.g. ConditionalStandardScale). Vocabulary builders (StringIndexer/OneHot), min/max scalers and distinct counts need exact/global stats and will be inaccurate — hence the runtime warning.

Tests cover: None = unchanged; sampled vs full mean/stddev within a loose tolerance on seeded data; an accumulator assertion that the source is scanned once (not once-per-estimator); and the mutual-exclusion warning.

Note: this is correctness-tested, but I have not benchmarked the actual speedup at prod scale — the single-scan behaviour is proven, the wall-clock win is expected but unmeasured.

2. Addressed your two latest comments on the cache bookkeeping

Since cacheIntermediateData and cacheEstimatorInput are mutually exclusive, I collapsed the cached_dataset / estimator_input_cache / new_cached handles into a single persisted_frame with one unpersist(). I kept it as a separate handle (rather than just unpersisting dataset at the end) because dataset gets reassigned by model.transform(...) between boundaries, so by loop-end it no longer points at the persisted frame — there's a docstring note explaining that.

Replace the all-estimators sampling behaviour of fitSampleFraction with a per-estimator opt-in boolean (useFitSample) on SampleFractionParams. Only estimators with useFitSample=True fit on the shared pipeline sample; all others fit on the full input, so vocabulary builders and min/max scalers that need exact/global statistics stay correct. Warn on the two conflicting configurations: an estimator that sets both useFitSample=True and its own sampleFraction (shared sample wins), and useFitSample=True with no pipeline fitSampleFraction (no-op). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

… load The scalar pandas_udf path let Arrow deliver Spark NULLs as NaN/pd.NA for numeric series, bypassing the `is None` null/OOV guards in element funcs. Restore Python None before mapping, guarded by hasnans so the null-free fast path (the common case) keeps its speedup. Also persist the pipeline-level fit params (checkpointInterval, cache/prune flags, fitSampleFraction/Seed) in the pipeline writer's metadata and restore them in the reader, so non-default values survive a save/load round-trip instead of silently resetting to defaults. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

georyetti

Just one minor nitpick, we have kamae classes for reading and writing metadata/params, could we use them. We added these due to slowdowns in how metadata get written on databricks

georyetti · 2026-08-24T12:20:04Z

+            for p in self.instance.params
+            if p.name != "stages" and self.instance.isSet(p)
+        }
+        DefaultParamsWriter.saveMetadata(


Can we use the kamae classes we have here: KamaeDefaultParamsWriter

georyetti · 2026-08-24T12:22:14Z

@@ -754,7 +831,15 @@ def load(self, path: str) -> KamaeSparkPipeline:
        """
        metadata = DefaultParamsReader.loadMetadata(path, self.sc)


Sorry spotted we don't use the reader here too. I know its not your change but could you use here too? KamaeDefaultParamsReader

Swap DefaultParamsReader/DefaultParamsWriter for the Kamae variants so pipeline metadata read/write uses the Databricks fast-path workaround the rest of kamae already relies on, keeping the write and read sides consistent. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Sbranikas

Just a minor question on the unpersist().

Sbranikas · 2026-08-25T13:52:36Z

+                    if sample_dataset is not None:
+                        sample_dataset = model.transform(sample_dataset)
+        if persisted_frame is not None:
+            persisted_frame.unpersist()


Just a small question, but right now only sampled_dataset gets cleaned up on failure (via the outer _fit's finally) — persisted_frame doesn't, right?

ConorWorthington and others added 24 commits June 29, 2026 16:31

Added localCheckPointInterval

96bc515

Plans are too slow to materialise - this is our attempt to speed it up

Update pipeline.py

b581ff1

Update single_feature_array_standard_scale.py

32502cc

Update standard_scale.py

4e367d3

Update conditional_standard_scale.py

2d2f75a

Update base.py

f16c0ef

Update bucketize.py

8089c4f

Update bucketize.py

9453119

Update conditional_standard_scale.py

06043b5

Update single_feature_array_standard_scale.py

b1096ce

Update standard_scale.py

70cd178

Update bucketize.py

c3db6b7

Update base.py

847f181

Update pipeline.py

6d5964a

Update pipeline.py

ff3e9f7

Update standard_scale.py

f8cff4d

Update single_feature_array_standard_scale.py

5bb8cd9

Update conditional_standard_scale.py

4c470d4

Update transform_utils.py

a8a8750

Update pyproject.toml

71cf01a

Update pipeline.py

6c4e18f

Update pipeline.py

771ab9c

Merge pull request #1 from ConorWorthington/pipeline-fit-optimisations

a2985e1

Merge latest version

ConorWorthington requested a review from a team as a code owner August 5, 2026 13:56

ConorWorthington requested review from ddonghi and mandrecki August 5, 2026 13:56

ConorWorthington changed the title ~~feat: opt-in pipeline fit optimisations, native bucketize, and sampled fitting~~ feat: opt-in pipeline fit optimisations, native bucketize, and sampled fitting support with caching Aug 5, 2026

ConorWorthington added 2 commits August 5, 2026 16:56

Patch black

94dea39

Patch error

4f5cea2

ConorWorthington added 4 commits August 5, 2026 16:57

Patch test case

3878d2f

Remove complexity, remove comments

1918090

Patch for lint

96fc184

Patch for linter

6af4ba4

georyetti requested changes Aug 12, 2026

View reviewed changes

georyetti requested changes Aug 17, 2026

View reviewed changes

refactor: annotate bucketize get_keras_layer return as tf.keras.layer…

60fa70b

…s.Layer BucketizeLayer is TensorFlow-only, so the tf-specific return type is accurate. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

georyetti previously approved these changes Aug 18, 2026

View reviewed changes

ConorWorthington dismissed georyetti’s stale review via 36859af August 18, 2026 17:09

georyetti requested changes Aug 20, 2026

View reviewed changes

cworthington and others added 2 commits August 21, 2026 15:32

georyetti requested changes Aug 24, 2026

View reviewed changes

georyetti approved these changes Aug 24, 2026

View reviewed changes

Sbranikas approved these changes Aug 25, 2026

View reviewed changes

Merge branch 'main' into main

2bf1a82

georyetti merged commit d581c6c into ExpediaGroup:main Aug 25, 2026
28 checks passed

		@@ -754,7 +831,15 @@ def load(self, path: str) -> KamaeSparkPipeline:
		"""
		metadata = DefaultParamsReader.loadMetadata(path, self.sc)

Uh oh!

Conversation

ConorWorthington commented Aug 5, 2026

Description

Keras Layer Checklist

Spark Transformer/Estimator Checklist

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!

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!

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!

ConorWorthington commented Aug 20, 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!

Sbranikas left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants