[ArranBasra]

Description

Fitting a pipeline currently runs one or more Spark passes per estimator. Each pass re-reads the input and recomputes the upstream transforms (including any explode), and the standard scalers also persist their whole input frame, which forces every column to be computed and cached for a fit that only reads one of them. On wide data with many estimators the fit is dominated by this repetition.

This PR adds an opt-in fuseEstimatorFits option to KamaeSparkPipeline that fits every estimator of a dependency layer in one aggregation, one pass over the data per layer, with nothing persisted. Default fit results are unchanged.

Two commits, meant to be reviewed separately:

1. refactor: split estimator fits into an aggregation part and a build-from-row part (no change in fitted values)

Estimators whose fit is a single aggregation now expose it in two parts on BaseEstimator:

• fit_aggregations(dataset, prefix): the helper columns and aggregation expressions (may add columns, never filters, explodes or samples rows);
• fit_from_row(row, prefix): builds the fitted transformer from the result row (driver only);
• can_fit_by_aggregation(dataset): whether the estimator fits this way (default False).

The normal fit (_fit_by_aggregation) is built from these two parts, keeping the existing persist where estimators already had it. Converted: ConditionalStandardScale (standard and binary), StandardScale, MinMaxScale, Impute (scalar inputs), and StringIndex, OneHotEncode, SharedStringIndex, SharedOneHotEncode (alphabetical string orders, scalar inputs). Nested-array scaler inputs, array inputs, frequency orders and SingleFeatureArrayStandardScale keep their own fit.

The label estimators' normal fit is unchanged (collect_labels_array). When fitted fused, their labels are collected with F.collect_set in the shared aggregation and ordered on the driver exactly as collect_labels_array orders them. Label estimators with maxNumLabels set keep the distributed path even when fusing, since collect_set gathers every distinct label before the limit is applied.

2. feat: opt-in fused estimator fits (fuseEstimatorFits, default False)

• Estimators are grouped into dependency layers through every column they read (inputs plus *Col/*Cols params such as relevanceCol and maskCols).
• For each layer: part one of every estimator that fits by aggregation, a single agg(), then part two of each. Other estimators (or those with their own sampleFraction) are fitted normally within their layer.
• Safeguards: the pipeline is fitted normally (with a warning) if a layer would give an estimator a different version of a column than the normal fit; an estimator whose helper columns add a row generator (e.g. explode) is fitted on its own (with a warning); on failure, estimators are pointed back at their original input columns.
• Ignored when fitSampleFraction is set.

Results

On a 17-estimator pipeline (9 scalers, 8 label estimators, 1 dependency layer) over 25.6M rows on a 12-node cluster:

	Normal fit	fuseEstimatorFits=True
Fit time	2h 44m	~14 min
Spark passes for the estimators	17	1
Labels		identical (all 8 estimators, up to 206k labels)
Scaler moments		within 1e-14 relative (summation order)

Testing

• Default path: fitted values are bitwise identical to 3.2.0 across the parameter combinations of every converted estimator (238 configurations, including masks, relevance, skipZeros, nanFillValue, binary scaling, mask tokens, maxNumLabels, nested and multi-column inputs).
• New tests in tests/kamae/spark/pipeline/:
  • test_fused_fit.py: fused vs normal fit per estimator, label ordering vs collect_labels_array, fallbacks, nested-array inputs, the row-generator safeguard;
  • test_fused_fit_layers.py: multi-layer graphs (chain, diamond, nested pipeline, transforms after unrelated estimators), asserting one fused aggregation per layer;
  • test_fused_fit_regressions.py: auxiliary columns as dependencies, masks, empty or fully-masked data, copied estimators sharing a uid, special characters in uids, reset on failure, rewritten columns, and that the normal fit releases its persisted data.
• Full test suite passes (1,783 passed, 1 skipped). The refactor commit on its own passes the existing Spark tests unchanged.

Checklists

The new layer / transformer checklists do not apply: no new transformers, estimators or Keras layers are added.

[georyetti]

One nit picky comment on potential collision naming of columns

I really like this addition thanks @ArranBasra !!!

[georyetti]

Nit picky comment, can we make this name collision proof? Very rare but this could already exist in the dataframe

[ArranBasra]

Good point, now using a random tmp prefix like elsewhere in Kamae. If it collides with another column it tries a new random prefix and added collision test.

[mantasu]

Was reviewing this with claude, it caught a couple things:

1. Copied estimators can fit on the wrong cast column. A .copy() keeps the same tmp_column_suffix. In the shared fused frame, copies reading the same input column with different inputDtypes overwrite or reuse each other's cast column. Reproduced: a double copy got the float mean, with no warning. Fix: a fresh suffix per estimator when fusing.

2. Label estimators can run out of memory on high-cardinality columns, even with fusion off. With alphabetical orders they now use collect_set, which sends every distinct label to the driver before maxNumLabels is applied. Suggest keeping the old distributed distinct() + take(n) unless fusing.

[ArranBasra]

Thanks Mantas I have addressed both comments.

1. Each fused estimator now gets its own tmp column suffix, so copies always fit on the correct cast (test added).
2. The normal fit is back to the distributed collect_labels_array. collect_set is only used when fusing, and only when maxNumLabels isn't set.