From 96bc5153234c0d9590e933c656410326b98ccbac Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Mon, 29 Jun 2026 16:31:59 +0100 Subject: [PATCH 01/36] Added localCheckPointInterval Plans are too slow to materialise - this is our attempt to speed it up --- src/kamae/spark/pipeline/pipeline.py | 93 +++++++++++++++++++++++++--- 1 file changed, 85 insertions(+), 8 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index 1da0d7ea..1a72a515 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -17,7 +17,7 @@ import networkx as nx from pyspark import keyword_only from pyspark.ml import Pipeline -from pyspark.ml.param import Params +from pyspark.ml.param import Param, Params, TypeConverters from pyspark.ml.pipeline import PipelineReader, PipelineSharedReadWrite, PipelineWriter from pyspark.ml.util import DefaultParamsReader, MLWriter from pyspark.sql import DataFrame @@ -38,17 +38,49 @@ class KamaeSparkPipeline(Pipeline): KamaeSparkPipeline is a subclass of pyspark.ml.Pipeline that is used to chain together BaseTransformers. It maintains the same functionality as pyspark.ml.Pipeline e.g. serialisation. + + The `localCheckpointInterval` param optionally bounds the depth of the Spark + logical plan built up while fitting a multi-estimator pipeline. When set to a + positive integer it triggers an ephemeral `DataFrame.localCheckpoint(eager=True)` + every `localCheckpointInterval` stages (evaluated at estimator-fit action + boundaries), physically truncating the accumulated lineage. This is a + depth-bounding / reliability feature: it guards against deep-plan failures such + as "plan too large", 64KB codegen, and CodeCache-full errors, and avoids + re-executing the full upstream lineage on every estimator fit. Its throughput + impact is data-dependent and NOT guaranteed positive (localCheckpoint persists + the full, wide intermediate DataFrame to executor local disk with no column + pruning), so benchmark before relying on it for speed. The default of 0 disables + checkpointing entirely, leaving fit behaviour byte-for-byte unchanged. """ + localCheckpointInterval = Param( + Params._dummy(), + "localCheckpointInterval", + "Number of stages between ephemeral localCheckpoint(eager=True) calls during " + "fit, used to bound logical-plan depth. 0 (the default) disables " + "checkpointing and leaves fit behaviour exactly unchanged.", + typeConverter=TypeConverters.toInt, + ) + @keyword_only - def __init__(self, *, stages: Optional[List["KamaePipelineStage"]] = None) -> None: + def __init__( + self, + *, + stages: Optional[List["KamaePipelineStage"]] = None, + localCheckpointInterval: int = 0, + ) -> None: """ Initialises the KamaeSparkPipeline object. :param stages: List of LayerTransformers to chain together. + :param localCheckpointInterval: Number of stages between ephemeral + localCheckpoint(eager=True) calls during fit. 0 (default) disables it. :returns: None - class instantiated. """ - super().__init__(stages=stages) + kwargs = self._input_kwargs + super().__init__() + self._setDefault(localCheckpointInterval=0) + self.setParams(**kwargs) def setStages(self, value: List["KamaePipelineStage"]) -> "KamaeSparkPipeline": """ @@ -67,15 +99,38 @@ def getStages(self) -> List["KamaePipelineStage"]: """ return self.getOrDefault("stages") + def setLocalCheckpointInterval(self, value: int) -> "KamaeSparkPipeline": + """ + Sets the `localCheckpointInterval` parameter. + + :param value: Number of stages between ephemeral localCheckpoint calls during + fit. 0 (or None) disables checkpointing. + :returns: KamaeSparkPipeline object with localCheckpointInterval set. + """ + return self._set(localCheckpointInterval=value) + + def getLocalCheckpointInterval(self) -> int: + """ + Gets the value of the `localCheckpointInterval` parameter. + + :returns: The localCheckpointInterval value. + """ + return self.getOrDefault(self.localCheckpointInterval) + @keyword_only def setParams( - self, *, stages: Optional["KamaePipelineStage"] = None + self, + *, + stages: Optional["KamaePipelineStage"] = None, + localCheckpointInterval: int = 0, ) -> "KamaeSparkPipeline": """ Sets the keyword arguments of the pipeline. :param stages: List of pipeline stages. - :returns: KamaeSparkPipeline object with stages set. + :param localCheckpointInterval: Number of stages between ephemeral + localCheckpoint(eager=True) calls during fit. 0 (default) disables it. + :returns: KamaeSparkPipeline object with params set. """ kwargs = self._input_kwargs return self._set(**kwargs) @@ -139,6 +194,14 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": Calls the super fit method of the pyspark.ml.Pipeline class and then constructs a KamaeSparkPipelineModel uses the stages from the fit pipeline. + If `localCheckpointInterval` is a positive integer, the working DataFrame is + ephemerally checkpointed via `localCheckpoint(eager=True)` roughly every + `localCheckpointInterval` stages (at estimator-fit action boundaries) to bound + logical-plan depth. localCheckpoint(eager=True) preserves the data exactly and + only truncates lineage, so fitted results are numerically identical to the + default (interval=0) behaviour. The default of 0 (or None) disables + checkpointing entirely. + :param dataset: PySpark DataFrame to fit the pipeline to. :returns: KamaeSparkPipelineModel object. """ @@ -162,15 +225,29 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": estimator_parent_stages = self.collect_estimator_parents( expanded_pipeline_stages ) + # Optional, opt-in plan-depth bounding. 0 (or None) keeps behaviour unchanged. + local_checkpoint_interval = self.getLocalCheckpointInterval() + checkpoint_enabled = ( + local_checkpoint_interval is not None and local_checkpoint_interval > 0 + ) + last_checkpoint_index = 0 # Fit each stage, appending the transformer to the list of transformers # If the stage is a parent of an estimator, transform the dataset. transformers: List[BaseTransformer] = [] - for stage in expanded_pipeline_stages: + for index, stage in enumerate(expanded_pipeline_stages): if isinstance(stage, BaseTransformer): transformers.append(stage) if stage in estimator_parent_stages: dataset = stage.transform(dataset) else: + # Truncate the accumulated lineage just before the fit action so the + # plan is physically bounded. eager=True forces materialisation now. + if ( + checkpoint_enabled + and index - last_checkpoint_index >= local_checkpoint_interval + ): + dataset = dataset.localCheckpoint(eager=True) + last_checkpoint_index = index model = stage.fit(dataset) transformers.append(model) if stage in estimator_parent_stages: @@ -215,7 +292,7 @@ class KamaeSparkPipelineReader(PipelineReader): Util class for reading a pipeline from a persistent storage path. """ - def __init__(self, cls: Type[KamaeSparkPipeline]) -> None: + def __init__(self, cls: Type[KamaeSparkPipeline]): super().__init__(cls=cls) def load(self, path: str) -> KamaeSparkPipeline: @@ -235,5 +312,5 @@ class KamaeSparkPipelineWriter(PipelineWriter): Util class for writing a pipeline to a persistent storage path. """ - def __init__(self, instance: KamaeSparkPipeline) -> None: + def __init__(self, instance: KamaeSparkPipeline): super().__init__(instance=instance) From b581ff1be944091b9523fd4dd5b7af94365258e0 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:36:23 +0100 Subject: [PATCH 02/36] Update pipeline.py --- src/kamae/spark/pipeline/pipeline.py | 167 ++++++++++++++++++++------- 1 file changed, 127 insertions(+), 40 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index 1a72a515..e1d972ea 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -21,6 +21,7 @@ from pyspark.ml.pipeline import PipelineReader, PipelineSharedReadWrite, PipelineWriter from pyspark.ml.util import DefaultParamsReader, MLWriter from pyspark.sql import DataFrame +from pyspark.storagelevel import StorageLevel from kamae.graph import PipelineGraph from kamae.spark.estimators import BaseEstimator @@ -39,47 +40,80 @@ class KamaeSparkPipeline(Pipeline): together BaseTransformers. It maintains the same functionality as pyspark.ml.Pipeline e.g. serialisation. - The `localCheckpointInterval` param optionally bounds the depth of the Spark + The `checkpointInterval` param optionally bounds the depth of the Spark logical plan built up while fitting a multi-estimator pipeline. When set to a - positive integer it triggers an ephemeral `DataFrame.localCheckpoint(eager=True)` - every `localCheckpointInterval` stages (evaluated at estimator-fit action + positive integer it triggers a reliable `DataFrame.checkpoint(eager=True)` + every `checkpointInterval` stages (evaluated at estimator-fit action boundaries), physically truncating the accumulated lineage. This is a depth-bounding / reliability feature: it guards against deep-plan failures such as "plan too large", 64KB codegen, and CodeCache-full errors, and avoids - re-executing the full upstream lineage on every estimator fit. Its throughput - impact is data-dependent and NOT guaranteed positive (localCheckpoint persists - the full, wide intermediate DataFrame to executor local disk with no column - pruning), so benchmark before relying on it for speed. The default of 0 disables - checkpointing entirely, leaving fit behaviour byte-for-byte unchanged. + re-executing the full upstream lineage on every estimator fit. + + Reliable checkpointing writes the intermediate DataFrame to the checkpoint + directory configured via `spark.sparkContext.setCheckpointDir()`, which + must point at fault-tolerant storage (DFS/cloud storage). Unlike local + checkpointing it survives executor loss (e.g. autoscaling, spot reclaim, OOM), + at the cost of writing to remote storage rather than executor-local disk. A + checkpoint directory MUST be set before fitting with a positive interval. Its + throughput impact is data-dependent and NOT guaranteed positive (the full, wide + intermediate DataFrame is persisted with no column pruning), so benchmark before + relying on it for speed. The default of 0 disables checkpointing entirely, + leaving fit behaviour byte-for-byte unchanged. + + The `cacheIntermediateData` param optionally persists the working DataFrame + (MEMORY_AND_DISK) at each estimator-fit boundary so that the estimator's fit + action - and any subsequent transforms - reuse a materialised result instead of + re-executing (and re-reading from source) the full upstream lineage on every + estimator. Only one intermediate frame is held at a time: each new persist + unpersists the one it supersedes, and the final frame is released before + returning. Unlike `checkpointInterval` it does not truncate the logical plan or + require a checkpoint directory; it is purely a re-scan-avoidance optimisation. + It preserves data exactly, so fitted results are identical to the default. The + default of False leaves fit behaviour unchanged. """ - localCheckpointInterval = Param( + checkpointInterval = Param( Params._dummy(), - "localCheckpointInterval", - "Number of stages between ephemeral localCheckpoint(eager=True) calls during " - "fit, used to bound logical-plan depth. 0 (the default) disables " + "checkpointInterval", + "Number of stages between reliable checkpoint(eager=True) calls during " + "fit, used to bound logical-plan depth. Requires a checkpoint directory set " + "via spark.sparkContext.setCheckpointDir. 0 (the default) disables " "checkpointing and leaves fit behaviour exactly unchanged.", typeConverter=TypeConverters.toInt, ) + cacheIntermediateData = Param( + Params._dummy(), + "cacheIntermediateData", + "If True, persist the working DataFrame (MEMORY_AND_DISK) at each " + "estimator-fit boundary so estimator fits reuse a materialised result " + "rather than re-scanning the upstream lineage from source. False (the " + "default) leaves fit behaviour exactly unchanged.", + typeConverter=TypeConverters.toBoolean, + ) + @keyword_only def __init__( self, *, stages: Optional[List["KamaePipelineStage"]] = None, - localCheckpointInterval: int = 0, + checkpointInterval: int = 0, + cacheIntermediateData: bool = False, ) -> None: """ Initialises the KamaeSparkPipeline object. :param stages: List of LayerTransformers to chain together. - :param localCheckpointInterval: Number of stages between ephemeral - localCheckpoint(eager=True) calls during fit. 0 (default) disables it. + :param checkpointInterval: Number of stages between reliable + checkpoint(eager=True) calls during fit. 0 (default) disables it. + :param cacheIntermediateData: If True, persist the working DataFrame at + each estimator-fit boundary to avoid re-scanning the upstream lineage. + False (default) disables it. :returns: None - class instantiated. """ kwargs = self._input_kwargs super().__init__() - self._setDefault(localCheckpointInterval=0) + self._setDefault(checkpointInterval=0, cacheIntermediateData=False) self.setParams(**kwargs) def setStages(self, value: List["KamaePipelineStage"]) -> "KamaeSparkPipeline": @@ -99,37 +133,58 @@ def getStages(self) -> List["KamaePipelineStage"]: """ return self.getOrDefault("stages") - def setLocalCheckpointInterval(self, value: int) -> "KamaeSparkPipeline": + def setCheckpointInterval(self, value: int) -> "KamaeSparkPipeline": """ - Sets the `localCheckpointInterval` parameter. + Sets the `checkpointInterval` parameter. - :param value: Number of stages between ephemeral localCheckpoint calls during + :param value: Number of stages between reliable checkpoint calls during fit. 0 (or None) disables checkpointing. - :returns: KamaeSparkPipeline object with localCheckpointInterval set. + :returns: KamaeSparkPipeline object with checkpointInterval set. + """ + return self._set(checkpointInterval=value) + + def getCheckpointInterval(self) -> int: + """ + Gets the value of the `checkpointInterval` parameter. + + :returns: The checkpointInterval value. """ - return self._set(localCheckpointInterval=value) + return self.getOrDefault(self.checkpointInterval) - def getLocalCheckpointInterval(self) -> int: + def setCacheIntermediateData(self, value: bool) -> "KamaeSparkPipeline": """ - Gets the value of the `localCheckpointInterval` parameter. + Sets the `cacheIntermediateData` parameter. - :returns: The localCheckpointInterval value. + :param value: Whether to persist the working DataFrame at each + estimator-fit boundary during fit. + :returns: KamaeSparkPipeline object with cacheIntermediateData set. """ - return self.getOrDefault(self.localCheckpointInterval) + return self._set(cacheIntermediateData=value) + + def getCacheIntermediateData(self) -> bool: + """ + Gets the value of the `cacheIntermediateData` parameter. + + :returns: The cacheIntermediateData value. + """ + return self.getOrDefault(self.cacheIntermediateData) @keyword_only def setParams( self, *, stages: Optional["KamaePipelineStage"] = None, - localCheckpointInterval: int = 0, + checkpointInterval: int = 0, + cacheIntermediateData: bool = False, ) -> "KamaeSparkPipeline": """ Sets the keyword arguments of the pipeline. :param stages: List of pipeline stages. - :param localCheckpointInterval: Number of stages between ephemeral - localCheckpoint(eager=True) calls during fit. 0 (default) disables it. + :param checkpointInterval: Number of stages between reliable + checkpoint(eager=True) calls during fit. 0 (default) disables it. + :param cacheIntermediateData: If True, persist the working DataFrame at + each estimator-fit boundary. False (default) disables it. :returns: KamaeSparkPipeline object with params set. """ kwargs = self._input_kwargs @@ -194,16 +249,25 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": Calls the super fit method of the pyspark.ml.Pipeline class and then constructs a KamaeSparkPipelineModel uses the stages from the fit pipeline. - If `localCheckpointInterval` is a positive integer, the working DataFrame is - ephemerally checkpointed via `localCheckpoint(eager=True)` roughly every - `localCheckpointInterval` stages (at estimator-fit action boundaries) to bound - logical-plan depth. localCheckpoint(eager=True) preserves the data exactly and + If `checkpointInterval` is a positive integer, the working DataFrame is + reliably checkpointed via `checkpoint(eager=True)` roughly every + `checkpointInterval` stages (at estimator-fit action boundaries) to bound + logical-plan depth. checkpoint(eager=True) preserves the data exactly and only truncates lineage, so fitted results are numerically identical to the - default (interval=0) behaviour. The default of 0 (or None) disables - checkpointing entirely. + default (interval=0) behaviour. A checkpoint directory must be configured via + `spark.sparkContext.setCheckpointDir` before fitting with a positive interval. + The default of 0 (or None) disables checkpointing entirely. + + If `cacheIntermediateData` is True, the working DataFrame is persisted + (MEMORY_AND_DISK) at each estimator-fit boundary so the fit action and any + subsequent transform reuse a materialised result rather than re-scanning the + upstream lineage. Persistence preserves data exactly, so fitted results are + identical to the default (False) behaviour. The default of False disables it. :param dataset: PySpark DataFrame to fit the pipeline to. :returns: KamaeSparkPipelineModel object. + :raises ValueError: If checkpointing is enabled but no checkpoint directory + has been set on the SparkContext. """ expanded_pipeline_stages = self.expand_pipeline_stages() @@ -226,11 +290,23 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": expanded_pipeline_stages ) # Optional, opt-in plan-depth bounding. 0 (or None) keeps behaviour unchanged. - local_checkpoint_interval = self.getLocalCheckpointInterval() - checkpoint_enabled = ( - local_checkpoint_interval is not None and local_checkpoint_interval > 0 - ) + checkpoint_interval = self.getCheckpointInterval() + checkpoint_enabled = checkpoint_interval is not None and checkpoint_interval > 0 + # Reliable checkpoint() requires a checkpoint directory; fail fast with a clear + # message rather than letting Spark raise mid-fit after work has been done. + if ( + checkpoint_enabled + and dataset.sparkSession.sparkContext.getCheckpointDir() is None + ): + raise ValueError( + "checkpointInterval > 0 requires a checkpoint directory. Set one via " + "spark.sparkContext.setCheckpointDir() before fitting." + ) + cache_enabled = self.getCacheIntermediateData() last_checkpoint_index = 0 + # Holds the single intermediate frame currently persisted (if any) so it can + # be unpersisted once superseded or once fitting completes. + cached_dataset: Optional[DataFrame] = None # Fit each stage, appending the transformer to the list of transformers # If the stage is a parent of an estimator, transform the dataset. transformers: List[BaseTransformer] = [] @@ -244,14 +320,25 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": # plan is physically bounded. eager=True forces materialisation now. if ( checkpoint_enabled - and index - last_checkpoint_index >= local_checkpoint_interval + and index - last_checkpoint_index >= checkpoint_interval ): - dataset = dataset.localCheckpoint(eager=True) + dataset = dataset.checkpoint(eager=True) last_checkpoint_index = index + # Persist the working frame so the fit action and any subsequent + # transform read a materialised result rather than re-scanning the + # upstream lineage from source. Only one frame is held at a time. + if cache_enabled: + new_cached = dataset.persist(StorageLevel.MEMORY_AND_DISK) + if cached_dataset is not None: + cached_dataset.unpersist() + cached_dataset = new_cached + dataset = new_cached model = stage.fit(dataset) transformers.append(model) if stage in estimator_parent_stages: dataset = model.transform(dataset) + if cached_dataset is not None: + cached_dataset.unpersist() return KamaeSparkPipelineModel(transformers) def copy(self, extra: Optional["ParamMap"] = None) -> "KamaeSparkPipeline": From 32502ccb98de7f72d481d6ecd9a5516411cd0082 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:36:49 +0100 Subject: [PATCH 03/36] Update single_feature_array_standard_scale.py --- .../single_feature_array_standard_scale.py | 78 ++++++++++--------- 1 file changed, 40 insertions(+), 38 deletions(-) diff --git a/src/kamae/spark/estimators/single_feature_array_standard_scale.py b/src/kamae/spark/estimators/single_feature_array_standard_scale.py index 4c209893..0207ed4f 100644 --- a/src/kamae/spark/estimators/single_feature_array_standard_scale.py +++ b/src/kamae/spark/estimators/single_feature_array_standard_scale.py @@ -19,13 +19,9 @@ from pyspark import keyword_only from pyspark.sql import DataFrame from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType +from pyspark.storagelevel import StorageLevel -from kamae.keras.core.backend import ALL_BACKENDS -from kamae.spark.params import ( - MaskValueParams, - SampleFractionParams, - SingleInputSingleOutputParams, -) +from kamae.spark.params import MaskValueParams, SingleInputSingleOutputParams from kamae.spark.transformers import StandardScaleTransformer from kamae.spark.utils import flatten_nested_arrays @@ -34,7 +30,6 @@ class SingleFeatureArrayStandardScaleEstimator( BaseEstimator, - SampleFractionParams, SingleInputSingleOutputParams, MaskValueParams, ): @@ -48,9 +43,6 @@ class SingleFeatureArrayStandardScaleEstimator( and standard deviation are calculated across all elements in all the arrays. """ - supported_backends = ALL_BACKENDS - jit_compatible = True - @keyword_only def __init__( self, @@ -60,7 +52,6 @@ def __init__( outputDtype: Optional[str] = None, layerName: Optional[str] = None, maskValue: Optional[float] = None, - sampleFraction: Optional[float] = None, ) -> None: """ Initializes a SingleFeatureArrayStandardScaleEstimator estimator. @@ -74,12 +65,10 @@ def __init__( transforming. :param layerName: Name of the layer. Used as the name of the tensorflow layer in the keras model. If not set, we use the uid of the Spark transformer. - :param sampleFraction: Fraction of data to sample for statistics - estimation (exclusive 0.0-1.0). Default None (no sampling). :returns: None - class instantiated. """ super().__init__() - self._setDefault(maskValue=None, sampleFraction=None) + self._setDefault(maskValue=None) kwargs = self._input_kwargs self.setParams(**kwargs) @@ -113,32 +102,45 @@ def _fit(self, dataset: DataFrame) -> "StandardScaleTransformer": Got {input_column_type} instead.""" ) - # Collect a single row to driver and get the length. - # We assume all subsequent rows have the same length. - array_size = np.array((dataset.select(self.getInputCol()).first()[0])).shape[-1] - - # Flatten the array to a single array. - # Will do nothing if the array is not nested. - flattened_array_col = flatten_nested_arrays( - column=F.col(self.getInputCol()), column_data_type=input_column_type - ) - - mean_and_stddev_dict: Dict[str, float] = ( - dataset.select(F.explode(flattened_array_col).alias(self.getInputCol())) - .withColumn( - "mask", - F.when( - F.col(self.getInputCol()) == F.lit(self.getMaskValue()), 1 - ).otherwise(0), + # Persist so the array-size probe and the moments aggregation reuse a + # materialised result instead of re-scanning the upstream lineage twice. + # Guarded so we do not double-persist data the caller already cached. + already_cached = dataset.storageLevel.useMemory or dataset.storageLevel.useDisk + if not already_cached: + dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) + + try: + # Collect a single row to driver and get the length. + # We assume all subsequent rows have the same length. + array_size = np.array( + (dataset.select(self.getInputCol()).first()[0]) + ).shape[-1] + + # Flatten the array to a single array. + # Will do nothing if the array is not nested. + flattened_array_col = flatten_nested_arrays( + column=F.col(self.getInputCol()), column_data_type=input_column_type ) - .filter(F.col("mask") == F.lit(0)) - .agg( - F.mean(self.getInputCol()).alias("mean"), - F.stddev_pop(self.getInputCol()).alias("stddev"), + + mean_and_stddev_dict: Dict[str, float] = ( + dataset.select(F.explode(flattened_array_col).alias(self.getInputCol())) + .withColumn( + "mask", + F.when( + F.col(self.getInputCol()) == F.lit(self.getMaskValue()), 1 + ).otherwise(0), + ) + .filter(F.col("mask") == F.lit(0)) + .agg( + F.mean(self.getInputCol()).alias("mean"), + F.stddev_pop(self.getInputCol()).alias("stddev"), + ) + .first() + .asDict() ) - .first() - .asDict() - ) + finally: + if not already_cached: + dataset.unpersist() mean: List[float] = [mean_and_stddev_dict["mean"] for _ in range(array_size)] stddev: List[float] = [ mean_and_stddev_dict["stddev"] for _ in range(array_size) From 4e367d3797f7effd981400122cd453d6351e854d Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:37:11 +0100 Subject: [PATCH 04/36] Update standard_scale.py --- src/kamae/spark/estimators/standard_scale.py | 106 +++++++++---------- 1 file changed, 52 insertions(+), 54 deletions(-) diff --git a/src/kamae/spark/estimators/standard_scale.py b/src/kamae/spark/estimators/standard_scale.py index a1c654ea..268a8ffb 100644 --- a/src/kamae/spark/estimators/standard_scale.py +++ b/src/kamae/spark/estimators/standard_scale.py @@ -21,15 +21,10 @@ import numpy as np import pyspark.sql.functions as F from pyspark import keyword_only -from pyspark.sql import DataFrame from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType +from pyspark.storagelevel import StorageLevel -from kamae.keras.core.backend import ALL_BACKENDS -from kamae.spark.params import ( - MaskValueParams, - SampleFractionParams, - SingleInputSingleOutputParams, -) +from kamae.spark.params import MaskValueParams, SingleInputSingleOutputParams from kamae.spark.transformers import StandardScaleTransformer from kamae.spark.utils import construct_nested_elements_for_scaling @@ -38,7 +33,6 @@ class StandardScaleEstimator( BaseEstimator, - SampleFractionParams, SingleInputSingleOutputParams, MaskValueParams, ): @@ -47,14 +41,8 @@ class StandardScaleEstimator( This estimator is used to calculate the mean and standard deviation of the input feature column. When fit is called it returns a StandardScaleTransformer which can be used to standardize/transform additional features. - - WARNING: If the input is an array, we assume that the array has a constant - shape across all rows. """ - supported_backends = ALL_BACKENDS - jit_compatible = True - @keyword_only def __init__( self, @@ -64,7 +52,6 @@ def __init__( outputDtype: Optional[str] = None, layerName: Optional[str] = None, maskValue: Optional[float] = None, - sampleFraction: Optional[float] = None, ) -> None: """ Initializes a StandardScaleEstimator estimator. @@ -78,12 +65,10 @@ def __init__( transforming. :param layerName: Name of the layer. Used as the name of the tensorflow layer in the keras model. If not set, we use the uid of the Spark transformer. - :param sampleFraction: Fraction of data to sample for statistics - estimation (exclusive 0.0-1.0). Default None (no sampling). :returns: None - class instantiated. """ super().__init__() - self._setDefault(maskValue=None, sampleFraction=None) + self._setDefault(maskValue=None) kwargs = self._input_kwargs self.setParams(**kwargs) @@ -97,7 +82,7 @@ def compatible_dtypes(self) -> Optional[List[DataType]]: """ return [FloatType(), DoubleType()] - def _fit(self, dataset: DataFrame) -> "StandardScaleTransformer": + def _fit(self, dataset) -> "StandardScaleTransformer": """ Fits the StandardScaleEstimator estimator to the given dataset. Calculates the mean and standard deviation of the input feature column and @@ -113,41 +98,54 @@ def _fit(self, dataset: DataFrame) -> "StandardScaleTransformer": else: input_col = F.col(self.getInputCol()) - # Collect a single row to driver and get the length. - # We assume all subsequent rows have the same length. - array_size = np.array((dataset.select(input_col).first()[0])).shape[-1] - - element_struct = construct_nested_elements_for_scaling( - column=input_col, - column_datatype=input_column_type, - array_dim=array_size, - ) - - mean_cols = [ - F.mean( - F.when( - F.col(f"element_struct.element_{i}") == F.lit(self.getMaskValue()), - F.lit(None), - ).otherwise(F.col(f"element_struct.element_{i}")) - ).alias(f"mean_{i}") - for i in range(1, array_size + 1) - ] - - stddev_cols = [ - F.stddev_pop( - F.when( - F.col(f"element_struct.element_{i}") == F.lit(self.getMaskValue()), - F.lit(None), - ).otherwise(F.col(f"element_struct.element_{i}")) - ).alias(f"stddev_{i}") - for i in range(1, array_size + 1) - ] - - metric_cols = mean_cols + stddev_cols - - mean_and_stddev_dict = ( - dataset.select(element_struct).agg(*metric_cols).first().asDict() - ) + # Persist so the array-size probe and the moments aggregation reuse a + # materialised result instead of re-scanning the upstream lineage twice. + # Guarded so we do not double-persist data the caller already cached. + already_cached = dataset.storageLevel.useMemory or dataset.storageLevel.useDisk + if not already_cached: + dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) + + try: + # Collect a single row to driver and get the length. + # We assume all subsequent rows have the same length. + array_size = np.array((dataset.select(input_col).first()[0])).shape[-1] + + element_struct = construct_nested_elements_for_scaling( + column=input_col, + column_datatype=input_column_type, + array_dim=array_size, + ) + + mean_cols = [ + F.mean( + F.when( + F.col(f"element_struct.element_{i}") + == F.lit(self.getMaskValue()), + F.lit(None), + ).otherwise(F.col(f"element_struct.element_{i}")) + ).alias(f"mean_{i}") + for i in range(1, array_size + 1) + ] + + stddev_cols = [ + F.stddev_pop( + F.when( + F.col(f"element_struct.element_{i}") + == F.lit(self.getMaskValue()), + F.lit(None), + ).otherwise(F.col(f"element_struct.element_{i}")) + ).alias(f"stddev_{i}") + for i in range(1, array_size + 1) + ] + + metric_cols = mean_cols + stddev_cols + + mean_and_stddev_dict = ( + dataset.select(element_struct).agg(*metric_cols).first().asDict() + ) + finally: + if not already_cached: + dataset.unpersist() mean = [mean_and_stddev_dict[f"mean_{i}"] for i in range(1, array_size + 1)] stddev = [mean_and_stddev_dict[f"stddev_{i}"] for i in range(1, array_size + 1)] From 2d2f75ae394609d630cf78895f04690e0589bb56 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:37:46 +0100 Subject: [PATCH 05/36] Update conditional_standard_scale.py --- .../estimators/conditional_standard_scale.py | 61 ++++++++++--------- 1 file changed, 32 insertions(+), 29 deletions(-) diff --git a/src/kamae/spark/estimators/conditional_standard_scale.py b/src/kamae/spark/estimators/conditional_standard_scale.py index b6edfb2f..5720a4c9 100644 --- a/src/kamae/spark/estimators/conditional_standard_scale.py +++ b/src/kamae/spark/estimators/conditional_standard_scale.py @@ -25,11 +25,10 @@ from pyspark.ml.param import Param, Params, TypeConverters from pyspark.sql import Column, DataFrame from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType +from pyspark.storagelevel import StorageLevel -from kamae.keras.core.backend import ALL_BACKENDS from kamae.spark.params import ( NanFillValueParams, - SampleFractionParams, SingleInputSingleOutputParams, StandardScaleSkipZerosParams, ) @@ -212,7 +211,6 @@ def getRelevanceCol(self) -> str: class ConditionalStandardScaleEstimator( BaseEstimator, - SampleFractionParams, SingleInputSingleOutputParams, ConditionalStandardScaleEstimatorParams, StandardScaleSkipZerosParams, @@ -233,14 +231,8 @@ class ConditionalStandardScaleEstimator( scalingFunction parameter. When fit is called it returns a ConditionalStandardScaleTransformer which can be used to standardize/transform the input data. - - WARNING: If the input is an array, we assume that the array has a constant - shape across all rows. """ - supported_backends = ALL_BACKENDS - jit_compatible = True - @keyword_only def __init__( self, @@ -257,7 +249,6 @@ def __init__( skipZeros: bool = False, epsilon: float = 0, nanFillValue: Optional[float] = None, - sampleFraction: Optional[float] = None, ) -> None: """ Initializes a ConditionalStandardScaleEstimator estimator. @@ -285,8 +276,6 @@ def __init__( when skipZeros is True. Defaults to 0. :param nanFillValue: Value to fill NaNs with after scaling. It is important to use it if epsilon filters out all the values. Defaults to None. - :param sampleFraction: Fraction of data to sample for statistics - estimation (exclusive 0.0-1.0). Default None (no sampling). :returns: None - class instantiated. """ super().__init__() @@ -299,7 +288,6 @@ def __init__( skipZeros=False, epsilon=0, nanFillValue=None, - sampleFraction=None, ) kwargs = self._input_kwargs self.setParams(**kwargs) @@ -378,22 +366,37 @@ def _fit(self, dataset: DataFrame) -> "ConditionalStandardScaleTransformer": mask_val = self.getMaskValues()[i] dataset = dataset.filter(mask_op(F.col(mask_col), mask_val)) - # Collect a single row to driver and get the length. - # We assume all subsequent rows have the same length. - row = dataset.select(input_col).first() - if row is None: - raise ValueError("No data left after application of mask conditions.") - array_size = np.array((row[0])).shape[-1] - - # Calculate the moments - if self.getScalingFunction().lower() == "standard": - return self._fit_standard( - dataset, input_col, input_column_dtype, array_size - ) - elif self.getScalingFunction().lower() == "binary": - return self._fit_binary(dataset, input_col, input_column_dtype, array_size) - else: - raise ValueError(f"Unknown scaling function: {self.getScalingFunction()}.") + # Persist so the array-size probe and the moments aggregation reuse a + # materialised result instead of re-scanning the (masked) upstream lineage + # twice. Guarded so we do not double-persist data the caller already cached. + already_cached = dataset.storageLevel.useMemory or dataset.storageLevel.useDisk + if not already_cached: + dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) + + try: + # Collect a single row to driver and get the length. + # We assume all subsequent rows have the same length. + row = dataset.select(input_col).first() + if row is None: + raise ValueError("No data left after application of mask conditions.") + array_size = np.array((row[0])).shape[-1] + + # Calculate the moments + if self.getScalingFunction().lower() == "standard": + return self._fit_standard( + dataset, input_col, input_column_dtype, array_size + ) + elif self.getScalingFunction().lower() == "binary": + return self._fit_binary( + dataset, input_col, input_column_dtype, array_size + ) + else: + raise ValueError( + f"Unknown scaling function: {self.getScalingFunction()}." + ) + finally: + if not already_cached: + dataset.unpersist() def _fit_binary( self, From f16c0ef819802e9b5740ce3c5bc99afb714158db Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:38:16 +0100 Subject: [PATCH 06/36] Update base.py --- src/kamae/spark/estimators/base.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/src/kamae/spark/estimators/base.py b/src/kamae/spark/estimators/base.py index dcd2a027..83734c3a 100644 --- a/src/kamae/spark/estimators/base.py +++ b/src/kamae/spark/estimators/base.py @@ -25,7 +25,7 @@ class BaseEstimator(Estimator, SparkOperation): - def __init__(self) -> None: + def __init__(self): """ Initializes the estimator. """ @@ -58,11 +58,6 @@ def fit( suffix=self.tmp_column_suffix, ) - if self.hasParam("sampleFraction"): - frac = self.getSampleFraction() - if frac is not None: - dataset = dataset.sample(fraction=frac) - # Replicate the logic from the existing abstract estimator fit method transformer = super().fit(dataset, params) @@ -86,9 +81,9 @@ def fit( param_dict = { param[0].name: param[1] for param in self.extractParamMap().items() } - raise e.__class__( + raise RuntimeError( f"Error in estimator: {self.uid} with params: {param_dict}" - ).with_traceback(e.__traceback__) + ) from e def construct_layer_info(self) -> Dict[str, Any]: """ From 8089c4fb91046a065081e647d2f25520cf722337 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:38:37 +0100 Subject: [PATCH 07/36] Update bucketize.py --- src/kamae/spark/transformers/bucketize.py | 52 +++++++++++------------ 1 file changed, 25 insertions(+), 27 deletions(-) diff --git a/src/kamae/spark/transformers/bucketize.py b/src/kamae/spark/transformers/bucketize.py index b639f0cc..9c706158 100644 --- a/src/kamae/spark/transformers/bucketize.py +++ b/src/kamae/spark/transformers/bucketize.py @@ -16,22 +16,21 @@ # pylint: disable=invalid-name # pylint: disable=too-many-ancestors # pylint: disable=no-member -from bisect import bisect_right -from typing import List, Optional, Union +from functools import reduce +from typing import List, Optional import pyspark.sql.functions as F import tensorflow as tf from pyspark import keyword_only from pyspark.ml.param import Param, Params, TypeConverters -from pyspark.sql import DataFrame +from pyspark.sql import Column, DataFrame from pyspark.sql.types import DataType, DoubleType, FloatType, IntegerType, LongType -from kamae.keras.core.backend import TENSORFLOW_ONLY -from kamae.keras.tensorflow.layers import BucketizeLayer from kamae.spark.params import SingleInputSingleOutputParams from kamae.spark.utils.transform_utils import ( - single_input_single_output_scalar_udf_transform, + single_input_single_output_scalar_transform, ) +from kamae.tensorflow.layers import BucketizeLayer from .base import BaseTransformer @@ -49,7 +48,7 @@ class BucketizeParams(Params): ) @staticmethod - def check_splits_sorted(splits: List[float]) -> None: + def check_splits_sorted(splits: List[float]): """ Checks that the splits parameter is sorted. @@ -90,10 +89,6 @@ class BucketizeTransformer( The 0 index is reserved for masking/padding. """ - jit_compatible = True - - supported_backends = TENSORFLOW_ONLY - @keyword_only def __init__( self, @@ -113,7 +108,7 @@ def __init__( transforming. :param outputDtype: Output data type to cast the output column to after transforming. - :param layerName: Name of the layer. Used as the name of the Keras layer + :param layerName: Name of the layer. Used as the name of the tensorflow layer in the keras model. If not set, we use the uid of the Spark transformer. :param splits: List of float values to use for bucketing. :returns: None - class instantiated. @@ -141,23 +136,26 @@ def _transform(self, dataset: DataFrame) -> DataFrame: :returns: Transformed pyspark dataframe. """ splits = self.getSplits() - # We need to create a UDF to perform binary search on the splits. - def bucketize(value: Optional[Union[float, int]]) -> Optional[int]: - # If null, keep null. There is no best bucket to place these into. - if value is None: - 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: + # Bucket index equals the number of splits <= value (matching + # bisect_right on a sorted splits list), plus 1 to reserve index 0 for + # mask/padding. Nulls are kept null - there is no best bucket for them. + bucket = reduce( + lambda acc, s: acc + + F.when(value >= F.lit(s), F.lit(1)).otherwise(F.lit(0)), + splits, + F.lit(1), + ) + return F.when(value.isNull(), F.lit(None)).otherwise(bucket) input_datatype = self.get_column_datatype( dataset=dataset, column_name=self.getInputCol() ) - output_col = single_input_single_output_scalar_udf_transform( + output_col = single_input_single_output_scalar_transform( input_col=F.col(self.getInputCol()), input_col_datatype=input_datatype, - func=lambda x: bucketize(x), - udf_return_element_datatype=IntegerType(), + func=bucketize, ) return dataset.withColumn( @@ -165,16 +163,16 @@ def bucketize(value: Optional[Union[float, int]]) -> Optional[int]: output_col, ) - def get_keras_layer(self) -> tf.keras.layers.Layer: + def get_tf_layer(self) -> tf.keras.layers.Layer: """ - Gets the Keras layer for the BucketizeLayer transformer. + Gets the tensorflow layer for the BucketizeLayer transformer. - :returns: Keras layer with name equal to the layerName parameter that + :returns: Tensorflow keras layer with name equal to the layerName parameter that performs a bucketing operation. """ return BucketizeLayer( name=self.getLayerName(), - input_dtype=self.getInputKerasDtype(), - output_dtype=self.getOutputKerasDtype(), + input_dtype=self.getInputTFDtype(), + output_dtype=self.getOutputTFDtype(), splits=self.getSplits(), ) From 94531191380bcf68a9ddb6690ed873ca98322b6b Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 1 Jul 2026 15:07:05 +0100 Subject: [PATCH 08/36] Update bucketize.py --- src/kamae/spark/transformers/bucketize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/kamae/spark/transformers/bucketize.py b/src/kamae/spark/transformers/bucketize.py index 9c706158..5ad3f618 100644 --- a/src/kamae/spark/transformers/bucketize.py +++ b/src/kamae/spark/transformers/bucketize.py @@ -30,7 +30,7 @@ from kamae.spark.utils.transform_utils import ( single_input_single_output_scalar_transform, ) -from kamae.tensorflow.layers import BucketizeLayer +from kamae.keras.tensorflow.layers import BucketizeLayer from .base import BaseTransformer From 06043b511f309586c138ec09910028ee77cf3247 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:27:33 +0100 Subject: [PATCH 09/36] Update conditional_standard_scale.py --- .../spark/estimators/conditional_standard_scale.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/kamae/spark/estimators/conditional_standard_scale.py b/src/kamae/spark/estimators/conditional_standard_scale.py index 5720a4c9..30f6ae65 100644 --- a/src/kamae/spark/estimators/conditional_standard_scale.py +++ b/src/kamae/spark/estimators/conditional_standard_scale.py @@ -27,8 +27,10 @@ from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType from pyspark.storagelevel import StorageLevel +from kamae.keras.core.backend import ALL_BACKENDS from kamae.spark.params import ( NanFillValueParams, + SampleFractionParams, SingleInputSingleOutputParams, StandardScaleSkipZerosParams, ) @@ -211,6 +213,7 @@ def getRelevanceCol(self) -> str: class ConditionalStandardScaleEstimator( BaseEstimator, + SampleFractionParams, SingleInputSingleOutputParams, ConditionalStandardScaleEstimatorParams, StandardScaleSkipZerosParams, @@ -231,8 +234,14 @@ class ConditionalStandardScaleEstimator( scalingFunction parameter. When fit is called it returns a ConditionalStandardScaleTransformer which can be used to standardize/transform the input data. + + WARNING: If the input is an array, we assume that the array has a constant + shape across all rows. """ + supported_backends = ALL_BACKENDS + jit_compatible = True + @keyword_only def __init__( self, @@ -249,6 +258,7 @@ def __init__( skipZeros: bool = False, epsilon: float = 0, nanFillValue: Optional[float] = None, + sampleFraction: Optional[float] = None, ) -> None: """ Initializes a ConditionalStandardScaleEstimator estimator. @@ -276,6 +286,8 @@ def __init__( when skipZeros is True. Defaults to 0. :param nanFillValue: Value to fill NaNs with after scaling. It is important to use it if epsilon filters out all the values. Defaults to None. + :param sampleFraction: Fraction of data to sample for statistics + estimation (exclusive 0.0-1.0). Default None (no sampling). :returns: None - class instantiated. """ super().__init__() @@ -288,6 +300,7 @@ def __init__( skipZeros=False, epsilon=0, nanFillValue=None, + sampleFraction=None, ) kwargs = self._input_kwargs self.setParams(**kwargs) From b1096ce12bd995b379f3ebfa67a8bf9650cb7fe5 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:27:52 +0100 Subject: [PATCH 10/36] Update single_feature_array_standard_scale.py --- .../single_feature_array_standard_scale.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/src/kamae/spark/estimators/single_feature_array_standard_scale.py b/src/kamae/spark/estimators/single_feature_array_standard_scale.py index 0207ed4f..a197ae4c 100644 --- a/src/kamae/spark/estimators/single_feature_array_standard_scale.py +++ b/src/kamae/spark/estimators/single_feature_array_standard_scale.py @@ -21,7 +21,12 @@ from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType from pyspark.storagelevel import StorageLevel -from kamae.spark.params import MaskValueParams, SingleInputSingleOutputParams +from kamae.keras.core.backend import ALL_BACKENDS +from kamae.spark.params import ( + MaskValueParams, + SampleFractionParams, + SingleInputSingleOutputParams, +) from kamae.spark.transformers import StandardScaleTransformer from kamae.spark.utils import flatten_nested_arrays @@ -30,6 +35,7 @@ class SingleFeatureArrayStandardScaleEstimator( BaseEstimator, + SampleFractionParams, SingleInputSingleOutputParams, MaskValueParams, ): @@ -43,6 +49,9 @@ class SingleFeatureArrayStandardScaleEstimator( and standard deviation are calculated across all elements in all the arrays. """ + supported_backends = ALL_BACKENDS + jit_compatible = True + @keyword_only def __init__( self, @@ -52,6 +61,7 @@ def __init__( outputDtype: Optional[str] = None, layerName: Optional[str] = None, maskValue: Optional[float] = None, + sampleFraction: Optional[float] = None, ) -> None: """ Initializes a SingleFeatureArrayStandardScaleEstimator estimator. @@ -65,10 +75,12 @@ def __init__( transforming. :param layerName: Name of the layer. Used as the name of the tensorflow layer in the keras model. If not set, we use the uid of the Spark transformer. + :param sampleFraction: Fraction of data to sample for statistics + estimation (exclusive 0.0-1.0). Default None (no sampling). :returns: None - class instantiated. """ super().__init__() - self._setDefault(maskValue=None) + self._setDefault(maskValue=None, sampleFraction=None) kwargs = self._input_kwargs self.setParams(**kwargs) From 70cd178a4aa00aa7f68c959d12f115a23839e1f9 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:28:09 +0100 Subject: [PATCH 11/36] Update standard_scale.py --- src/kamae/spark/estimators/standard_scale.py | 22 +++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/src/kamae/spark/estimators/standard_scale.py b/src/kamae/spark/estimators/standard_scale.py index 268a8ffb..379dce37 100644 --- a/src/kamae/spark/estimators/standard_scale.py +++ b/src/kamae/spark/estimators/standard_scale.py @@ -21,10 +21,16 @@ import numpy as np import pyspark.sql.functions as F from pyspark import keyword_only +from pyspark.sql import DataFrame from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType from pyspark.storagelevel import StorageLevel -from kamae.spark.params import MaskValueParams, SingleInputSingleOutputParams +from kamae.keras.core.backend import ALL_BACKENDS +from kamae.spark.params import ( + MaskValueParams, + SampleFractionParams, + SingleInputSingleOutputParams, +) from kamae.spark.transformers import StandardScaleTransformer from kamae.spark.utils import construct_nested_elements_for_scaling @@ -33,6 +39,7 @@ class StandardScaleEstimator( BaseEstimator, + SampleFractionParams, SingleInputSingleOutputParams, MaskValueParams, ): @@ -41,8 +48,14 @@ class StandardScaleEstimator( This estimator is used to calculate the mean and standard deviation of the input feature column. When fit is called it returns a StandardScaleTransformer which can be used to standardize/transform additional features. + + WARNING: If the input is an array, we assume that the array has a constant + shape across all rows. """ + supported_backends = ALL_BACKENDS + jit_compatible = True + @keyword_only def __init__( self, @@ -52,6 +65,7 @@ def __init__( outputDtype: Optional[str] = None, layerName: Optional[str] = None, maskValue: Optional[float] = None, + sampleFraction: Optional[float] = None, ) -> None: """ Initializes a StandardScaleEstimator estimator. @@ -65,10 +79,12 @@ def __init__( transforming. :param layerName: Name of the layer. Used as the name of the tensorflow layer in the keras model. If not set, we use the uid of the Spark transformer. + :param sampleFraction: Fraction of data to sample for statistics + estimation (exclusive 0.0-1.0). Default None (no sampling). :returns: None - class instantiated. """ super().__init__() - self._setDefault(maskValue=None) + self._setDefault(maskValue=None, sampleFraction=None) kwargs = self._input_kwargs self.setParams(**kwargs) @@ -82,7 +98,7 @@ def compatible_dtypes(self) -> Optional[List[DataType]]: """ return [FloatType(), DoubleType()] - def _fit(self, dataset) -> "StandardScaleTransformer": + def _fit(self, dataset: DataFrame) -> "StandardScaleTransformer": """ Fits the StandardScaleEstimator estimator to the given dataset. Calculates the mean and standard deviation of the input feature column and From c3db6b78b92ef02e67f6d6a30f95806fa5ae4713 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:28:45 +0100 Subject: [PATCH 12/36] Update bucketize.py --- src/kamae/spark/transformers/bucketize.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/src/kamae/spark/transformers/bucketize.py b/src/kamae/spark/transformers/bucketize.py index 5ad3f618..022a0042 100644 --- a/src/kamae/spark/transformers/bucketize.py +++ b/src/kamae/spark/transformers/bucketize.py @@ -19,18 +19,19 @@ from functools import reduce from typing import List, Optional +import keras import pyspark.sql.functions as F -import tensorflow as tf from pyspark import keyword_only from pyspark.ml.param import Param, Params, TypeConverters from pyspark.sql import Column, DataFrame from pyspark.sql.types import DataType, DoubleType, FloatType, IntegerType, LongType +from kamae.keras.core.backend import TENSORFLOW_ONLY +from kamae.keras.tensorflow.layers import BucketizeLayer from kamae.spark.params import SingleInputSingleOutputParams from kamae.spark.utils.transform_utils import ( single_input_single_output_scalar_transform, ) -from kamae.keras.tensorflow.layers import BucketizeLayer from .base import BaseTransformer @@ -89,6 +90,9 @@ class BucketizeTransformer( The 0 index is reserved for masking/padding. """ + supported_backends = TENSORFLOW_ONLY + jit_compatible = True + @keyword_only def __init__( self, @@ -163,16 +167,16 @@ def bucketize(value: Column) -> Column: output_col, ) - def get_tf_layer(self) -> tf.keras.layers.Layer: + def get_keras_layer(self) -> keras.layers.Layer: """ - Gets the tensorflow layer for the BucketizeLayer transformer. + Gets the Keras layer for the BucketizeLayer transformer. - :returns: Tensorflow keras layer with name equal to the layerName parameter that + :returns: Keras layer with name equal to the layerName parameter that performs a bucketing operation. """ return BucketizeLayer( name=self.getLayerName(), - input_dtype=self.getInputTFDtype(), - output_dtype=self.getOutputTFDtype(), + input_dtype=self.getInputKerasDtype(), + output_dtype=self.getOutputKerasDtype(), splits=self.getSplits(), ) From 847f181276feb3b2372c8bb67a43801a25757257 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:34:42 +0100 Subject: [PATCH 13/36] Update base.py --- src/kamae/spark/estimators/base.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/kamae/spark/estimators/base.py b/src/kamae/spark/estimators/base.py index 83734c3a..dcd2a027 100644 --- a/src/kamae/spark/estimators/base.py +++ b/src/kamae/spark/estimators/base.py @@ -25,7 +25,7 @@ class BaseEstimator(Estimator, SparkOperation): - def __init__(self): + def __init__(self) -> None: """ Initializes the estimator. """ @@ -58,6 +58,11 @@ def fit( suffix=self.tmp_column_suffix, ) + if self.hasParam("sampleFraction"): + frac = self.getSampleFraction() + if frac is not None: + dataset = dataset.sample(fraction=frac) + # Replicate the logic from the existing abstract estimator fit method transformer = super().fit(dataset, params) @@ -81,9 +86,9 @@ def fit( param_dict = { param[0].name: param[1] for param in self.extractParamMap().items() } - raise RuntimeError( + raise e.__class__( f"Error in estimator: {self.uid} with params: {param_dict}" - ) from e + ).with_traceback(e.__traceback__) def construct_layer_info(self) -> Dict[str, Any]: """ From 6d5964a893f04be4a53b953707d00a5c81b714f5 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:35:06 +0100 Subject: [PATCH 14/36] Update pipeline.py --- src/kamae/spark/pipeline/pipeline.py | 107 +++++++++++---------------- 1 file changed, 44 insertions(+), 63 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index e1d972ea..5e8d3629 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -40,25 +40,18 @@ class KamaeSparkPipeline(Pipeline): together BaseTransformers. It maintains the same functionality as pyspark.ml.Pipeline e.g. serialisation. - The `checkpointInterval` param optionally bounds the depth of the Spark + The `localCheckpointInterval` param optionally bounds the depth of the Spark logical plan built up while fitting a multi-estimator pipeline. When set to a - positive integer it triggers a reliable `DataFrame.checkpoint(eager=True)` - every `checkpointInterval` stages (evaluated at estimator-fit action + positive integer it triggers an ephemeral `DataFrame.localCheckpoint(eager=True)` + every `localCheckpointInterval` stages (evaluated at estimator-fit action boundaries), physically truncating the accumulated lineage. This is a depth-bounding / reliability feature: it guards against deep-plan failures such as "plan too large", 64KB codegen, and CodeCache-full errors, and avoids - re-executing the full upstream lineage on every estimator fit. - - Reliable checkpointing writes the intermediate DataFrame to the checkpoint - directory configured via `spark.sparkContext.setCheckpointDir()`, which - must point at fault-tolerant storage (DFS/cloud storage). Unlike local - checkpointing it survives executor loss (e.g. autoscaling, spot reclaim, OOM), - at the cost of writing to remote storage rather than executor-local disk. A - checkpoint directory MUST be set before fitting with a positive interval. Its - throughput impact is data-dependent and NOT guaranteed positive (the full, wide - intermediate DataFrame is persisted with no column pruning), so benchmark before - relying on it for speed. The default of 0 disables checkpointing entirely, - leaving fit behaviour byte-for-byte unchanged. + re-executing the full upstream lineage on every estimator fit. Its throughput + impact is data-dependent and NOT guaranteed positive (localCheckpoint persists + the full, wide intermediate DataFrame to executor local disk with no column + pruning), so benchmark before relying on it for speed. The default of 0 disables + checkpointing entirely, leaving fit behaviour byte-for-byte unchanged. The `cacheIntermediateData` param optionally persists the working DataFrame (MEMORY_AND_DISK) at each estimator-fit boundary so that the estimator's fit @@ -66,18 +59,17 @@ class KamaeSparkPipeline(Pipeline): re-executing (and re-reading from source) the full upstream lineage on every estimator. Only one intermediate frame is held at a time: each new persist unpersists the one it supersedes, and the final frame is released before - returning. Unlike `checkpointInterval` it does not truncate the logical plan or - require a checkpoint directory; it is purely a re-scan-avoidance optimisation. - It preserves data exactly, so fitted results are identical to the default. The - default of False leaves fit behaviour unchanged. + returning. Unlike `localCheckpointInterval` it does not truncate the logical + plan; it is purely a re-scan-avoidance optimisation. It preserves data exactly, + so fitted results are identical to the default. The default of False leaves fit + behaviour unchanged. """ - checkpointInterval = Param( + localCheckpointInterval = Param( Params._dummy(), - "checkpointInterval", - "Number of stages between reliable checkpoint(eager=True) calls during " - "fit, used to bound logical-plan depth. Requires a checkpoint directory set " - "via spark.sparkContext.setCheckpointDir. 0 (the default) disables " + "localCheckpointInterval", + "Number of stages between ephemeral localCheckpoint(eager=True) calls during " + "fit, used to bound logical-plan depth. 0 (the default) disables " "checkpointing and leaves fit behaviour exactly unchanged.", typeConverter=TypeConverters.toInt, ) @@ -97,15 +89,15 @@ def __init__( self, *, stages: Optional[List["KamaePipelineStage"]] = None, - checkpointInterval: int = 0, + localCheckpointInterval: int = 0, cacheIntermediateData: bool = False, ) -> None: """ Initialises the KamaeSparkPipeline object. :param stages: List of LayerTransformers to chain together. - :param checkpointInterval: Number of stages between reliable - checkpoint(eager=True) calls during fit. 0 (default) disables it. + :param localCheckpointInterval: Number of stages between ephemeral + localCheckpoint(eager=True) calls during fit. 0 (default) disables it. :param cacheIntermediateData: If True, persist the working DataFrame at each estimator-fit boundary to avoid re-scanning the upstream lineage. False (default) disables it. @@ -113,7 +105,7 @@ def __init__( """ kwargs = self._input_kwargs super().__init__() - self._setDefault(checkpointInterval=0, cacheIntermediateData=False) + self._setDefault(localCheckpointInterval=0, cacheIntermediateData=False) self.setParams(**kwargs) def setStages(self, value: List["KamaePipelineStage"]) -> "KamaeSparkPipeline": @@ -133,23 +125,23 @@ def getStages(self) -> List["KamaePipelineStage"]: """ return self.getOrDefault("stages") - def setCheckpointInterval(self, value: int) -> "KamaeSparkPipeline": + def setLocalCheckpointInterval(self, value: int) -> "KamaeSparkPipeline": """ - Sets the `checkpointInterval` parameter. + Sets the `localCheckpointInterval` parameter. - :param value: Number of stages between reliable checkpoint calls during + :param value: Number of stages between ephemeral localCheckpoint calls during fit. 0 (or None) disables checkpointing. - :returns: KamaeSparkPipeline object with checkpointInterval set. + :returns: KamaeSparkPipeline object with localCheckpointInterval set. """ - return self._set(checkpointInterval=value) + return self._set(localCheckpointInterval=value) - def getCheckpointInterval(self) -> int: + def getLocalCheckpointInterval(self) -> int: """ - Gets the value of the `checkpointInterval` parameter. + Gets the value of the `localCheckpointInterval` parameter. - :returns: The checkpointInterval value. + :returns: The localCheckpointInterval value. """ - return self.getOrDefault(self.checkpointInterval) + return self.getOrDefault(self.localCheckpointInterval) def setCacheIntermediateData(self, value: bool) -> "KamaeSparkPipeline": """ @@ -174,15 +166,15 @@ def setParams( self, *, stages: Optional["KamaePipelineStage"] = None, - checkpointInterval: int = 0, + localCheckpointInterval: int = 0, cacheIntermediateData: bool = False, ) -> "KamaeSparkPipeline": """ Sets the keyword arguments of the pipeline. :param stages: List of pipeline stages. - :param checkpointInterval: Number of stages between reliable - checkpoint(eager=True) calls during fit. 0 (default) disables it. + :param localCheckpointInterval: Number of stages between ephemeral + localCheckpoint(eager=True) calls during fit. 0 (default) disables it. :param cacheIntermediateData: If True, persist the working DataFrame at each estimator-fit boundary. False (default) disables it. :returns: KamaeSparkPipeline object with params set. @@ -249,14 +241,13 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": Calls the super fit method of the pyspark.ml.Pipeline class and then constructs a KamaeSparkPipelineModel uses the stages from the fit pipeline. - If `checkpointInterval` is a positive integer, the working DataFrame is - reliably checkpointed via `checkpoint(eager=True)` roughly every - `checkpointInterval` stages (at estimator-fit action boundaries) to bound - logical-plan depth. checkpoint(eager=True) preserves the data exactly and + If `localCheckpointInterval` is a positive integer, the working DataFrame is + ephemerally checkpointed via `localCheckpoint(eager=True)` roughly every + `localCheckpointInterval` stages (at estimator-fit action boundaries) to bound + logical-plan depth. localCheckpoint(eager=True) preserves the data exactly and only truncates lineage, so fitted results are numerically identical to the - default (interval=0) behaviour. A checkpoint directory must be configured via - `spark.sparkContext.setCheckpointDir` before fitting with a positive interval. - The default of 0 (or None) disables checkpointing entirely. + default (interval=0) behaviour. The default of 0 (or None) disables + checkpointing entirely. If `cacheIntermediateData` is True, the working DataFrame is persisted (MEMORY_AND_DISK) at each estimator-fit boundary so the fit action and any @@ -266,8 +257,6 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": :param dataset: PySpark DataFrame to fit the pipeline to. :returns: KamaeSparkPipelineModel object. - :raises ValueError: If checkpointing is enabled but no checkpoint directory - has been set on the SparkContext. """ expanded_pipeline_stages = self.expand_pipeline_stages() @@ -290,18 +279,10 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": expanded_pipeline_stages ) # Optional, opt-in plan-depth bounding. 0 (or None) keeps behaviour unchanged. - checkpoint_interval = self.getCheckpointInterval() - checkpoint_enabled = checkpoint_interval is not None and checkpoint_interval > 0 - # Reliable checkpoint() requires a checkpoint directory; fail fast with a clear - # message rather than letting Spark raise mid-fit after work has been done. - if ( - checkpoint_enabled - and dataset.sparkSession.sparkContext.getCheckpointDir() is None - ): - raise ValueError( - "checkpointInterval > 0 requires a checkpoint directory. Set one via " - "spark.sparkContext.setCheckpointDir() before fitting." - ) + local_checkpoint_interval = self.getLocalCheckpointInterval() + checkpoint_enabled = ( + local_checkpoint_interval is not None and local_checkpoint_interval > 0 + ) cache_enabled = self.getCacheIntermediateData() last_checkpoint_index = 0 # Holds the single intermediate frame currently persisted (if any) so it can @@ -320,9 +301,9 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": # plan is physically bounded. eager=True forces materialisation now. if ( checkpoint_enabled - and index - last_checkpoint_index >= checkpoint_interval + and index - last_checkpoint_index >= local_checkpoint_interval ): - dataset = dataset.checkpoint(eager=True) + dataset = dataset.localCheckpoint(eager=True) last_checkpoint_index = index # Persist the working frame so the fit action and any subsequent # transform read a materialised result rather than re-scanning the From ff3e9f7f446feb2027ab5cf24239df70bea3cb1b Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 1 Jul 2026 17:35:05 +0100 Subject: [PATCH 15/36] Update pipeline.py --- src/kamae/spark/pipeline/pipeline.py | 107 ++++++++++++++++----------- 1 file changed, 63 insertions(+), 44 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index 5e8d3629..e1d972ea 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -40,18 +40,25 @@ class KamaeSparkPipeline(Pipeline): together BaseTransformers. It maintains the same functionality as pyspark.ml.Pipeline e.g. serialisation. - The `localCheckpointInterval` param optionally bounds the depth of the Spark + The `checkpointInterval` param optionally bounds the depth of the Spark logical plan built up while fitting a multi-estimator pipeline. When set to a - positive integer it triggers an ephemeral `DataFrame.localCheckpoint(eager=True)` - every `localCheckpointInterval` stages (evaluated at estimator-fit action + positive integer it triggers a reliable `DataFrame.checkpoint(eager=True)` + every `checkpointInterval` stages (evaluated at estimator-fit action boundaries), physically truncating the accumulated lineage. This is a depth-bounding / reliability feature: it guards against deep-plan failures such as "plan too large", 64KB codegen, and CodeCache-full errors, and avoids - re-executing the full upstream lineage on every estimator fit. Its throughput - impact is data-dependent and NOT guaranteed positive (localCheckpoint persists - the full, wide intermediate DataFrame to executor local disk with no column - pruning), so benchmark before relying on it for speed. The default of 0 disables - checkpointing entirely, leaving fit behaviour byte-for-byte unchanged. + re-executing the full upstream lineage on every estimator fit. + + Reliable checkpointing writes the intermediate DataFrame to the checkpoint + directory configured via `spark.sparkContext.setCheckpointDir()`, which + must point at fault-tolerant storage (DFS/cloud storage). Unlike local + checkpointing it survives executor loss (e.g. autoscaling, spot reclaim, OOM), + at the cost of writing to remote storage rather than executor-local disk. A + checkpoint directory MUST be set before fitting with a positive interval. Its + throughput impact is data-dependent and NOT guaranteed positive (the full, wide + intermediate DataFrame is persisted with no column pruning), so benchmark before + relying on it for speed. The default of 0 disables checkpointing entirely, + leaving fit behaviour byte-for-byte unchanged. The `cacheIntermediateData` param optionally persists the working DataFrame (MEMORY_AND_DISK) at each estimator-fit boundary so that the estimator's fit @@ -59,17 +66,18 @@ class KamaeSparkPipeline(Pipeline): re-executing (and re-reading from source) the full upstream lineage on every estimator. Only one intermediate frame is held at a time: each new persist unpersists the one it supersedes, and the final frame is released before - returning. Unlike `localCheckpointInterval` it does not truncate the logical - plan; it is purely a re-scan-avoidance optimisation. It preserves data exactly, - so fitted results are identical to the default. The default of False leaves fit - behaviour unchanged. + returning. Unlike `checkpointInterval` it does not truncate the logical plan or + require a checkpoint directory; it is purely a re-scan-avoidance optimisation. + It preserves data exactly, so fitted results are identical to the default. The + default of False leaves fit behaviour unchanged. """ - localCheckpointInterval = Param( + checkpointInterval = Param( Params._dummy(), - "localCheckpointInterval", - "Number of stages between ephemeral localCheckpoint(eager=True) calls during " - "fit, used to bound logical-plan depth. 0 (the default) disables " + "checkpointInterval", + "Number of stages between reliable checkpoint(eager=True) calls during " + "fit, used to bound logical-plan depth. Requires a checkpoint directory set " + "via spark.sparkContext.setCheckpointDir. 0 (the default) disables " "checkpointing and leaves fit behaviour exactly unchanged.", typeConverter=TypeConverters.toInt, ) @@ -89,15 +97,15 @@ def __init__( self, *, stages: Optional[List["KamaePipelineStage"]] = None, - localCheckpointInterval: int = 0, + checkpointInterval: int = 0, cacheIntermediateData: bool = False, ) -> None: """ Initialises the KamaeSparkPipeline object. :param stages: List of LayerTransformers to chain together. - :param localCheckpointInterval: Number of stages between ephemeral - localCheckpoint(eager=True) calls during fit. 0 (default) disables it. + :param checkpointInterval: Number of stages between reliable + checkpoint(eager=True) calls during fit. 0 (default) disables it. :param cacheIntermediateData: If True, persist the working DataFrame at each estimator-fit boundary to avoid re-scanning the upstream lineage. False (default) disables it. @@ -105,7 +113,7 @@ def __init__( """ kwargs = self._input_kwargs super().__init__() - self._setDefault(localCheckpointInterval=0, cacheIntermediateData=False) + self._setDefault(checkpointInterval=0, cacheIntermediateData=False) self.setParams(**kwargs) def setStages(self, value: List["KamaePipelineStage"]) -> "KamaeSparkPipeline": @@ -125,23 +133,23 @@ def getStages(self) -> List["KamaePipelineStage"]: """ return self.getOrDefault("stages") - def setLocalCheckpointInterval(self, value: int) -> "KamaeSparkPipeline": + def setCheckpointInterval(self, value: int) -> "KamaeSparkPipeline": """ - Sets the `localCheckpointInterval` parameter. + Sets the `checkpointInterval` parameter. - :param value: Number of stages between ephemeral localCheckpoint calls during + :param value: Number of stages between reliable checkpoint calls during fit. 0 (or None) disables checkpointing. - :returns: KamaeSparkPipeline object with localCheckpointInterval set. + :returns: KamaeSparkPipeline object with checkpointInterval set. """ - return self._set(localCheckpointInterval=value) + return self._set(checkpointInterval=value) - def getLocalCheckpointInterval(self) -> int: + def getCheckpointInterval(self) -> int: """ - Gets the value of the `localCheckpointInterval` parameter. + Gets the value of the `checkpointInterval` parameter. - :returns: The localCheckpointInterval value. + :returns: The checkpointInterval value. """ - return self.getOrDefault(self.localCheckpointInterval) + return self.getOrDefault(self.checkpointInterval) def setCacheIntermediateData(self, value: bool) -> "KamaeSparkPipeline": """ @@ -166,15 +174,15 @@ def setParams( self, *, stages: Optional["KamaePipelineStage"] = None, - localCheckpointInterval: int = 0, + checkpointInterval: int = 0, cacheIntermediateData: bool = False, ) -> "KamaeSparkPipeline": """ Sets the keyword arguments of the pipeline. :param stages: List of pipeline stages. - :param localCheckpointInterval: Number of stages between ephemeral - localCheckpoint(eager=True) calls during fit. 0 (default) disables it. + :param checkpointInterval: Number of stages between reliable + checkpoint(eager=True) calls during fit. 0 (default) disables it. :param cacheIntermediateData: If True, persist the working DataFrame at each estimator-fit boundary. False (default) disables it. :returns: KamaeSparkPipeline object with params set. @@ -241,13 +249,14 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": Calls the super fit method of the pyspark.ml.Pipeline class and then constructs a KamaeSparkPipelineModel uses the stages from the fit pipeline. - If `localCheckpointInterval` is a positive integer, the working DataFrame is - ephemerally checkpointed via `localCheckpoint(eager=True)` roughly every - `localCheckpointInterval` stages (at estimator-fit action boundaries) to bound - logical-plan depth. localCheckpoint(eager=True) preserves the data exactly and + If `checkpointInterval` is a positive integer, the working DataFrame is + reliably checkpointed via `checkpoint(eager=True)` roughly every + `checkpointInterval` stages (at estimator-fit action boundaries) to bound + logical-plan depth. checkpoint(eager=True) preserves the data exactly and only truncates lineage, so fitted results are numerically identical to the - default (interval=0) behaviour. The default of 0 (or None) disables - checkpointing entirely. + default (interval=0) behaviour. A checkpoint directory must be configured via + `spark.sparkContext.setCheckpointDir` before fitting with a positive interval. + The default of 0 (or None) disables checkpointing entirely. If `cacheIntermediateData` is True, the working DataFrame is persisted (MEMORY_AND_DISK) at each estimator-fit boundary so the fit action and any @@ -257,6 +266,8 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": :param dataset: PySpark DataFrame to fit the pipeline to. :returns: KamaeSparkPipelineModel object. + :raises ValueError: If checkpointing is enabled but no checkpoint directory + has been set on the SparkContext. """ expanded_pipeline_stages = self.expand_pipeline_stages() @@ -279,10 +290,18 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": expanded_pipeline_stages ) # Optional, opt-in plan-depth bounding. 0 (or None) keeps behaviour unchanged. - local_checkpoint_interval = self.getLocalCheckpointInterval() - checkpoint_enabled = ( - local_checkpoint_interval is not None and local_checkpoint_interval > 0 - ) + checkpoint_interval = self.getCheckpointInterval() + checkpoint_enabled = checkpoint_interval is not None and checkpoint_interval > 0 + # Reliable checkpoint() requires a checkpoint directory; fail fast with a clear + # message rather than letting Spark raise mid-fit after work has been done. + if ( + checkpoint_enabled + and dataset.sparkSession.sparkContext.getCheckpointDir() is None + ): + raise ValueError( + "checkpointInterval > 0 requires a checkpoint directory. Set one via " + "spark.sparkContext.setCheckpointDir() before fitting." + ) cache_enabled = self.getCacheIntermediateData() last_checkpoint_index = 0 # Holds the single intermediate frame currently persisted (if any) so it can @@ -301,9 +320,9 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": # plan is physically bounded. eager=True forces materialisation now. if ( checkpoint_enabled - and index - last_checkpoint_index >= local_checkpoint_interval + and index - last_checkpoint_index >= checkpoint_interval ): - dataset = dataset.localCheckpoint(eager=True) + dataset = dataset.checkpoint(eager=True) last_checkpoint_index = index # Persist the working frame so the fit action and any subsequent # transform read a materialised result rather than re-scanning the From f8cff4d21d2e1a21fad61ded45f0f528fc69bb21 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Thu, 2 Jul 2026 14:17:13 +0100 Subject: [PATCH 16/36] Update standard_scale.py --- src/kamae/spark/estimators/standard_scale.py | 86 +++++++++----------- 1 file changed, 37 insertions(+), 49 deletions(-) diff --git a/src/kamae/spark/estimators/standard_scale.py b/src/kamae/spark/estimators/standard_scale.py index 379dce37..ebf70f2a 100644 --- a/src/kamae/spark/estimators/standard_scale.py +++ b/src/kamae/spark/estimators/standard_scale.py @@ -23,7 +23,6 @@ from pyspark import keyword_only from pyspark.sql import DataFrame from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType -from pyspark.storagelevel import StorageLevel from kamae.keras.core.backend import ALL_BACKENDS from kamae.spark.params import ( @@ -114,54 +113,43 @@ def _fit(self, dataset: DataFrame) -> "StandardScaleTransformer": else: input_col = F.col(self.getInputCol()) - # Persist so the array-size probe and the moments aggregation reuse a - # materialised result instead of re-scanning the upstream lineage twice. - # Guarded so we do not double-persist data the caller already cached. - already_cached = dataset.storageLevel.useMemory or dataset.storageLevel.useDisk - if not already_cached: - dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) - - try: - # Collect a single row to driver and get the length. - # We assume all subsequent rows have the same length. - array_size = np.array((dataset.select(input_col).first()[0])).shape[-1] - - element_struct = construct_nested_elements_for_scaling( - column=input_col, - column_datatype=input_column_type, - array_dim=array_size, - ) - - mean_cols = [ - F.mean( - F.when( - F.col(f"element_struct.element_{i}") - == F.lit(self.getMaskValue()), - F.lit(None), - ).otherwise(F.col(f"element_struct.element_{i}")) - ).alias(f"mean_{i}") - for i in range(1, array_size + 1) - ] - - stddev_cols = [ - F.stddev_pop( - F.when( - F.col(f"element_struct.element_{i}") - == F.lit(self.getMaskValue()), - F.lit(None), - ).otherwise(F.col(f"element_struct.element_{i}")) - ).alias(f"stddev_{i}") - for i in range(1, array_size + 1) - ] - - metric_cols = mean_cols + stddev_cols - - mean_and_stddev_dict = ( - dataset.select(element_struct).agg(*metric_cols).first().asDict() - ) - finally: - if not already_cached: - dataset.unpersist() + # Collect a single row to driver and get the length. + # We assume all subsequent rows have the same length. + array_size = np.array((dataset.select(input_col).first()[0])).shape[-1] + + element_struct = construct_nested_elements_for_scaling( + column=input_col, + column_datatype=input_column_type, + array_dim=array_size, + ) + + mean_cols = [ + F.mean( + F.when( + F.col(f"element_struct.element_{i}") + == F.lit(self.getMaskValue()), + F.lit(None), + ).otherwise(F.col(f"element_struct.element_{i}")) + ).alias(f"mean_{i}") + for i in range(1, array_size + 1) + ] + + stddev_cols = [ + F.stddev_pop( + F.when( + F.col(f"element_struct.element_{i}") + == F.lit(self.getMaskValue()), + F.lit(None), + ).otherwise(F.col(f"element_struct.element_{i}")) + ).alias(f"stddev_{i}") + for i in range(1, array_size + 1) + ] + + metric_cols = mean_cols + stddev_cols + + mean_and_stddev_dict = ( + dataset.select(element_struct).agg(*metric_cols).first().asDict() + ) mean = [mean_and_stddev_dict[f"mean_{i}"] for i in range(1, array_size + 1)] stddev = [mean_and_stddev_dict[f"stddev_{i}"] for i in range(1, array_size + 1)] From 5bb8cd944f13a3576677ac5d64932a2f21a1da1c Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Thu, 2 Jul 2026 14:17:40 +0100 Subject: [PATCH 17/36] Update single_feature_array_standard_scale.py --- .../single_feature_array_standard_scale.py | 64 ++++++++----------- 1 file changed, 26 insertions(+), 38 deletions(-) diff --git a/src/kamae/spark/estimators/single_feature_array_standard_scale.py b/src/kamae/spark/estimators/single_feature_array_standard_scale.py index a197ae4c..74ed5eed 100644 --- a/src/kamae/spark/estimators/single_feature_array_standard_scale.py +++ b/src/kamae/spark/estimators/single_feature_array_standard_scale.py @@ -19,7 +19,6 @@ from pyspark import keyword_only from pyspark.sql import DataFrame from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType -from pyspark.storagelevel import StorageLevel from kamae.keras.core.backend import ALL_BACKENDS from kamae.spark.params import ( @@ -114,45 +113,34 @@ def _fit(self, dataset: DataFrame) -> "StandardScaleTransformer": Got {input_column_type} instead.""" ) - # Persist so the array-size probe and the moments aggregation reuse a - # materialised result instead of re-scanning the upstream lineage twice. - # Guarded so we do not double-persist data the caller already cached. - already_cached = dataset.storageLevel.useMemory or dataset.storageLevel.useDisk - if not already_cached: - dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) - - try: - # Collect a single row to driver and get the length. - # We assume all subsequent rows have the same length. - array_size = np.array( - (dataset.select(self.getInputCol()).first()[0]) - ).shape[-1] - - # Flatten the array to a single array. - # Will do nothing if the array is not nested. - flattened_array_col = flatten_nested_arrays( - column=F.col(self.getInputCol()), column_data_type=input_column_type - ) + # Collect a single row to driver and get the length. + # We assume all subsequent rows have the same length. + array_size = np.array( + (dataset.select(self.getInputCol()).first()[0]) + ).shape[-1] + + # Flatten the array to a single array. + # Will do nothing if the array is not nested. + flattened_array_col = flatten_nested_arrays( + column=F.col(self.getInputCol()), column_data_type=input_column_type + ) - mean_and_stddev_dict: Dict[str, float] = ( - dataset.select(F.explode(flattened_array_col).alias(self.getInputCol())) - .withColumn( - "mask", - F.when( - F.col(self.getInputCol()) == F.lit(self.getMaskValue()), 1 - ).otherwise(0), - ) - .filter(F.col("mask") == F.lit(0)) - .agg( - F.mean(self.getInputCol()).alias("mean"), - F.stddev_pop(self.getInputCol()).alias("stddev"), - ) - .first() - .asDict() + mean_and_stddev_dict: Dict[str, float] = ( + dataset.select(F.explode(flattened_array_col).alias(self.getInputCol())) + .withColumn( + "mask", + F.when( + F.col(self.getInputCol()) == F.lit(self.getMaskValue()), 1 + ).otherwise(0), ) - finally: - if not already_cached: - dataset.unpersist() + .filter(F.col("mask") == F.lit(0)) + .agg( + F.mean(self.getInputCol()).alias("mean"), + F.stddev_pop(self.getInputCol()).alias("stddev"), + ) + .first() + .asDict() + ) mean: List[float] = [mean_and_stddev_dict["mean"] for _ in range(array_size)] stddev: List[float] = [ mean_and_stddev_dict["stddev"] for _ in range(array_size) From 4c470d47dfd856510ef74dfdd70f18b69a55a4f8 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Thu, 2 Jul 2026 14:19:39 +0100 Subject: [PATCH 18/36] Update conditional_standard_scale.py --- .../estimators/conditional_standard_scale.py | 50 ++++++++----------- 1 file changed, 21 insertions(+), 29 deletions(-) diff --git a/src/kamae/spark/estimators/conditional_standard_scale.py b/src/kamae/spark/estimators/conditional_standard_scale.py index 30f6ae65..0af4abc5 100644 --- a/src/kamae/spark/estimators/conditional_standard_scale.py +++ b/src/kamae/spark/estimators/conditional_standard_scale.py @@ -25,7 +25,6 @@ from pyspark.ml.param import Param, Params, TypeConverters from pyspark.sql import Column, DataFrame from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType -from pyspark.storagelevel import StorageLevel from kamae.keras.core.backend import ALL_BACKENDS from kamae.spark.params import ( @@ -382,34 +381,27 @@ def _fit(self, dataset: DataFrame) -> "ConditionalStandardScaleTransformer": # Persist so the array-size probe and the moments aggregation reuse a # materialised result instead of re-scanning the (masked) upstream lineage # twice. Guarded so we do not double-persist data the caller already cached. - already_cached = dataset.storageLevel.useMemory or dataset.storageLevel.useDisk - if not already_cached: - dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) - - try: - # Collect a single row to driver and get the length. - # We assume all subsequent rows have the same length. - row = dataset.select(input_col).first() - if row is None: - raise ValueError("No data left after application of mask conditions.") - array_size = np.array((row[0])).shape[-1] - - # Calculate the moments - if self.getScalingFunction().lower() == "standard": - return self._fit_standard( - dataset, input_col, input_column_dtype, array_size - ) - elif self.getScalingFunction().lower() == "binary": - return self._fit_binary( - dataset, input_col, input_column_dtype, array_size - ) - else: - raise ValueError( - f"Unknown scaling function: {self.getScalingFunction()}." - ) - finally: - if not already_cached: - dataset.unpersist() + # Collect a single row to driver and get the length. + # We assume all subsequent rows have the same length. + row = dataset.select(input_col).first() + if row is None: + raise ValueError("No data left after application of mask conditions.") + array_size = np.array((row[0])).shape[-1] + + # Calculate the moments + if self.getScalingFunction().lower() == "standard": + return self._fit_standard( + dataset, input_col, input_column_dtype, array_size + ) + elif self.getScalingFunction().lower() == "binary": + return self._fit_binary( + dataset, input_col, input_column_dtype, array_size + ) + else: + raise ValueError( + f"Unknown scaling function: {self.getScalingFunction()}." + ) + def _fit_binary( self, From a8a8750c88114e277d6c2bacc6eb19f2bf5d3818 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Thu, 2 Jul 2026 14:22:42 +0100 Subject: [PATCH 19/36] Update transform_utils.py --- src/kamae/spark/utils/transform_utils.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/kamae/spark/utils/transform_utils.py b/src/kamae/spark/utils/transform_utils.py index c67d8773..c2be3050 100644 --- a/src/kamae/spark/utils/transform_utils.py +++ b/src/kamae/spark/utils/transform_utils.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - +import pandas as pd from typing import Callable, List import pyspark.sql.functions as F @@ -150,10 +150,23 @@ def _single_input_single_output_udf_transform( func=func, nest_level=nested_level, ) + # Scalar (non-array) columns transfer as a flat Arrow batch, so a pandas_udf + # that maps the same per-element func avoids the per-row pickling of a plain + # Python UDF (~1.4x faster). Nested-array columns are kept on the row-wise UDF: + # Arrow (de)serialisation of nested lists there costs more than it saves. + if not isinstance(input_col_datatype, ArrayType): + + def _vectorized_func(series: pd.Series) -> pd.Series: + return series.map(nested_lambda_func) + + udf_func = F.pandas_udf(_vectorized_func, udf_return_type) + return udf_func(input_col) + udf_func = F.udf(nested_lambda_func, udf_return_type) return udf_func(input_col) + def single_input_single_output_scalar_udf_transform( input_col: Column, input_col_datatype: DataType, From 71cf01a32799f08205a277bc82eaf48aaef565ff Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Thu, 2 Jul 2026 14:23:24 +0100 Subject: [PATCH 20/36] Update pyproject.toml --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index ed5d9050..0c175c4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ requires-python = ">=3.10,<3.13" dependencies = [ "pyspark>=3.4.0,<4.0.0", "pandas>=1.3.4,<3.0.0", + "pyarrow>=4.0.0", "networkx>=2.6.3,<3.0.0", "pyfarmhash>=0.3.2,<0.4.0", "keras>=3.0.0,<4.0.0", From 6c4e18f8007d34b4c95377a24cfbe826dc4b028b Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Mon, 6 Jul 2026 15:44:54 +0100 Subject: [PATCH 21/36] Update pipeline.py --- src/kamae/spark/pipeline/pipeline.py | 55 +++++++++++++++++++++++++++- 1 file changed, 54 insertions(+), 1 deletion(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index e1d972ea..95c14eba 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import TYPE_CHECKING, List, Optional, Type +from typing import TYPE_CHECKING, List, Optional, Set, Type import networkx as nx from pyspark import keyword_only @@ -242,6 +242,51 @@ def collect_estimator_parents( ] return estimator_parent_stages + @staticmethod + def collect_required_input_columns( + stages: List["KamaePipelineStage"], + ) -> Set[str]: + """ + Collects every column read as an input by any stage in the pipeline. + + A raw input-DataFrame column absent from this set is consumed by no stage + and can be dropped before fitting, so it is not carried through every + transform (and every materialisation) below. + + :param stages: List of pipeline stages. + :returns: Set of column names read by at least one stage. + """ + required_input_columns: Set[str] = set() + for stage in stages: + inputs, _ = stage.get_layer_inputs_outputs() + required_input_columns.update(inputs) + return required_input_columns + + def prune_unused_input_columns( + self, + dataset: DataFrame, + stages: List["KamaePipelineStage"], + ) -> DataFrame: + """ + Projects the input DataFrame down to only the columns the pipeline reads. + + Columns produced by stages are created downstream via `withColumn`, so only + the pipeline's source columns need to be present up front. Pruning here + keeps the frame narrow before any expansion, reducing the cost of every + subsequent transform and materialisation. If no unused columns are found + (or the pipeline reads none of the DataFrame's columns) the DataFrame is + returned unchanged. + + :param dataset: Input DataFrame to prune. + :param stages: Expanded pipeline stages. + :returns: DataFrame projected to the columns the pipeline consumes. + """ + 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): + return dataset.select(*columns_to_keep) + return dataset + def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": """ Fits the pipeline to the dataset. Returns a KamaeSparkPipelineModel object. @@ -249,6 +294,10 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": Calls the super fit method of the pyspark.ml.Pipeline class and then constructs a KamaeSparkPipelineModel uses the stages from the fit pipeline. + Before fitting, the input DataFrame is projected down to only the columns + the pipeline reads (see `prune_unused_input_columns`), so columns no stage + consumes are not carried through every transform and materialisation. + If `checkpointInterval` is a positive integer, the working DataFrame is reliably checkpointed via `checkpoint(eager=True)` roughly every `checkpointInterval` stages (at estimator-fit action boundaries) to bound @@ -279,6 +328,10 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": "Cannot recognize a pipeline stage of type %s." % type(stage) ) + # Drop input columns no stage consumes before any expansion, so dead + # columns are not carried through every transform and materialisation below. + dataset = self.prune_unused_input_columns(dataset, expanded_pipeline_stages) + # Native Spark checks for the last estimator and executes all transformers # before it, regardless whether there is a dependency between them. See here: # https://github.com/apache/spark/blob/master/python/pyspark/ml/pipeline.py#L120 From 771ab9cde181962fe218b2930d719a22f1bff71e Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Mon, 6 Jul 2026 17:58:44 +0100 Subject: [PATCH 22/36] Update pipeline.py --- src/kamae/spark/pipeline/pipeline.py | 90 +++++++++++++++++++++++++--- 1 file changed, 81 insertions(+), 9 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index 95c14eba..4d4aeaae 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -70,6 +70,14 @@ class KamaeSparkPipeline(Pipeline): require a checkpoint directory; it is purely a re-scan-avoidance optimisation. It preserves data exactly, so fitted results are identical to the default. The default of False leaves fit behaviour unchanged. + + The `pruneInputColumns` param optionally projects the input DataFrame down to + only the columns the pipeline reads before fitting, dropping columns no stage + consumes so they are not carried through every transform and materialisation. + The set of columns to keep is computed generously (see + `collect_required_input_columns`) to avoid dropping columns referenced via + params other than `inputCol(s)`. The default of False leaves fit behaviour + unchanged. """ checkpointInterval = Param( @@ -92,6 +100,15 @@ class KamaeSparkPipeline(Pipeline): typeConverter=TypeConverters.toBoolean, ) + pruneInputColumns = Param( + Params._dummy(), + "pruneInputColumns", + "If True, project the input DataFrame down to only the columns the " + "pipeline reads before fitting, dropping columns no stage consumes. False " + "(the default) leaves fit behaviour exactly unchanged.", + typeConverter=TypeConverters.toBoolean, + ) + @keyword_only def __init__( self, @@ -99,6 +116,7 @@ def __init__( stages: Optional[List["KamaePipelineStage"]] = None, checkpointInterval: int = 0, cacheIntermediateData: bool = False, + pruneInputColumns: bool = False, ) -> None: """ Initialises the KamaeSparkPipeline object. @@ -109,11 +127,17 @@ def __init__( :param cacheIntermediateData: If True, persist the working DataFrame at each estimator-fit boundary to avoid re-scanning the upstream lineage. False (default) disables it. + :param pruneInputColumns: If True, drop input columns no stage consumes + before fitting. False (default) disables it. :returns: None - class instantiated. """ kwargs = self._input_kwargs super().__init__() - self._setDefault(checkpointInterval=0, cacheIntermediateData=False) + self._setDefault( + checkpointInterval=0, + cacheIntermediateData=False, + pruneInputColumns=False, + ) self.setParams(**kwargs) def setStages(self, value: List["KamaePipelineStage"]) -> "KamaeSparkPipeline": @@ -169,6 +193,23 @@ def getCacheIntermediateData(self) -> bool: """ return self.getOrDefault(self.cacheIntermediateData) + def setPruneInputColumns(self, value: bool) -> "KamaeSparkPipeline": + """ + Sets the `pruneInputColumns` parameter. + + :param value: Whether to drop input columns no stage consumes before fitting. + :returns: KamaeSparkPipeline object with pruneInputColumns set. + """ + return self._set(pruneInputColumns=value) + + def getPruneInputColumns(self) -> bool: + """ + Gets the value of the `pruneInputColumns` parameter. + + :returns: The pruneInputColumns value. + """ + return self.getOrDefault(self.pruneInputColumns) + @keyword_only def setParams( self, @@ -176,6 +217,7 @@ def setParams( stages: Optional["KamaePipelineStage"] = None, checkpointInterval: int = 0, cacheIntermediateData: bool = False, + pruneInputColumns: bool = False, ) -> "KamaeSparkPipeline": """ Sets the keyword arguments of the pipeline. @@ -185,6 +227,8 @@ def setParams( checkpoint(eager=True) calls during fit. 0 (default) disables it. :param cacheIntermediateData: If True, persist the working DataFrame at each estimator-fit boundary. False (default) disables it. + :param pruneInputColumns: If True, drop input columns no stage consumes + before fitting. False (default) disables it. :returns: KamaeSparkPipeline object with params set. """ kwargs = self._input_kwargs @@ -247,19 +291,43 @@ def collect_required_input_columns( stages: List["KamaePipelineStage"], ) -> Set[str]: """ - Collects every column read as an input by any stage in the pipeline. + Collects every column potentially read by any stage in the pipeline. A raw input-DataFrame column absent from this set is consumed by no stage and can be dropped before fitting, so it is not carried through every transform (and every materialisation) below. + This is deliberately generous: as well as the canonical inputs + (`inputCol`/`inputCols` from `get_layer_inputs_outputs`), it unions the + value(s) of every param whose name ends in `Col`/`Cols`. Some stages read + extra columns during fit through such params (e.g. `maskCols` and + `relevanceCol` on ConditionalStandardScaleEstimator, `queryIdCol` on + listwise transformers) that `inputCol(s)` does not capture. The name + convention holds across all estimators and transformers, so this + self-maintaining heuristic covers future aux column params without a + per-stage allowlist. Over-inclusion is harmless - a name that is not a real + input column simply never matches `dataset.columns` - whereas omitting a + referenced column would wrongly drop data the pipeline needs at fit time. + :param stages: List of pipeline stages. - :returns: Set of column names read by at least one stage. + :returns: Set of column names potentially read by at least one stage. """ required_input_columns: Set[str] = set() for stage in stages: inputs, _ = stage.get_layer_inputs_outputs() required_input_columns.update(inputs) + 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) + ) return required_input_columns def prune_unused_input_columns( @@ -294,9 +362,11 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": Calls the super fit method of the pyspark.ml.Pipeline class and then constructs a KamaeSparkPipelineModel uses the stages from the fit pipeline. - Before fitting, the input DataFrame is projected down to only the columns - the pipeline reads (see `prune_unused_input_columns`), so columns no stage - consumes are not carried through every transform and materialisation. + If `pruneInputColumns` is True, the input DataFrame is projected down to + only the columns the pipeline reads (see `prune_unused_input_columns`) + before fitting, so columns no stage consumes are not carried through every + transform and materialisation. The default of False leaves fit behaviour + unchanged. If `checkpointInterval` is a positive integer, the working DataFrame is reliably checkpointed via `checkpoint(eager=True)` roughly every @@ -328,9 +398,11 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": "Cannot recognize a pipeline stage of type %s." % type(stage) ) - # Drop input columns no stage consumes before any expansion, so dead - # columns are not carried through every transform and materialisation below. - dataset = self.prune_unused_input_columns(dataset, expanded_pipeline_stages) + # Optional, opt-in. Drop input columns no stage consumes before any + # expansion, so dead columns are not carried through every transform and + # materialisation below. Default False keeps behaviour unchanged. + if self.getPruneInputColumns(): + dataset = self.prune_unused_input_columns(dataset, expanded_pipeline_stages) # Native Spark checks for the last estimator and executes all transformers # before it, regardless whether there is a dependency between them. See here: From d57d2a42cb775b3b11660f7e8bcf8efe283aa6cd Mon Sep 17 00:00:00 2001 From: cworthington Date: Wed, 5 Aug 2026 14:51:03 +0100 Subject: [PATCH 23/36] perf: persist during scaler fit and add pipeline fit-optimisation tests 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 --- src/kamae/spark/estimators/base.py | 4 +- .../estimators/conditional_standard_scale.py | 50 ++-- .../single_feature_array_standard_scale.py | 64 ++-- src/kamae/spark/estimators/standard_scale.py | 86 +++--- tests/kamae/spark/conftest.py | 7 +- tests/kamae/spark/pipeline/test_pipeline.py | 277 +++++++++++++++++- 6 files changed, 399 insertions(+), 89 deletions(-) diff --git a/src/kamae/spark/estimators/base.py b/src/kamae/spark/estimators/base.py index dcd2a027..90f651b1 100644 --- a/src/kamae/spark/estimators/base.py +++ b/src/kamae/spark/estimators/base.py @@ -86,9 +86,9 @@ def fit( param_dict = { param[0].name: param[1] for param in self.extractParamMap().items() } - raise e.__class__( + raise RuntimeError( f"Error in estimator: {self.uid} with params: {param_dict}" - ).with_traceback(e.__traceback__) + ) from e def construct_layer_info(self) -> Dict[str, Any]: """ diff --git a/src/kamae/spark/estimators/conditional_standard_scale.py b/src/kamae/spark/estimators/conditional_standard_scale.py index 0af4abc5..30f6ae65 100644 --- a/src/kamae/spark/estimators/conditional_standard_scale.py +++ b/src/kamae/spark/estimators/conditional_standard_scale.py @@ -25,6 +25,7 @@ from pyspark.ml.param import Param, Params, TypeConverters from pyspark.sql import Column, DataFrame from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType +from pyspark.storagelevel import StorageLevel from kamae.keras.core.backend import ALL_BACKENDS from kamae.spark.params import ( @@ -381,27 +382,34 @@ def _fit(self, dataset: DataFrame) -> "ConditionalStandardScaleTransformer": # Persist so the array-size probe and the moments aggregation reuse a # materialised result instead of re-scanning the (masked) upstream lineage # twice. Guarded so we do not double-persist data the caller already cached. - # Collect a single row to driver and get the length. - # We assume all subsequent rows have the same length. - row = dataset.select(input_col).first() - if row is None: - raise ValueError("No data left after application of mask conditions.") - array_size = np.array((row[0])).shape[-1] - - # Calculate the moments - if self.getScalingFunction().lower() == "standard": - return self._fit_standard( - dataset, input_col, input_column_dtype, array_size - ) - elif self.getScalingFunction().lower() == "binary": - return self._fit_binary( - dataset, input_col, input_column_dtype, array_size - ) - else: - raise ValueError( - f"Unknown scaling function: {self.getScalingFunction()}." - ) - + already_cached = dataset.storageLevel.useMemory or dataset.storageLevel.useDisk + if not already_cached: + dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) + + try: + # Collect a single row to driver and get the length. + # We assume all subsequent rows have the same length. + row = dataset.select(input_col).first() + if row is None: + raise ValueError("No data left after application of mask conditions.") + array_size = np.array((row[0])).shape[-1] + + # Calculate the moments + if self.getScalingFunction().lower() == "standard": + return self._fit_standard( + dataset, input_col, input_column_dtype, array_size + ) + elif self.getScalingFunction().lower() == "binary": + return self._fit_binary( + dataset, input_col, input_column_dtype, array_size + ) + else: + raise ValueError( + f"Unknown scaling function: {self.getScalingFunction()}." + ) + finally: + if not already_cached: + dataset.unpersist() def _fit_binary( self, diff --git a/src/kamae/spark/estimators/single_feature_array_standard_scale.py b/src/kamae/spark/estimators/single_feature_array_standard_scale.py index 74ed5eed..a197ae4c 100644 --- a/src/kamae/spark/estimators/single_feature_array_standard_scale.py +++ b/src/kamae/spark/estimators/single_feature_array_standard_scale.py @@ -19,6 +19,7 @@ from pyspark import keyword_only from pyspark.sql import DataFrame from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType +from pyspark.storagelevel import StorageLevel from kamae.keras.core.backend import ALL_BACKENDS from kamae.spark.params import ( @@ -113,34 +114,45 @@ def _fit(self, dataset: DataFrame) -> "StandardScaleTransformer": Got {input_column_type} instead.""" ) - # Collect a single row to driver and get the length. - # We assume all subsequent rows have the same length. - array_size = np.array( - (dataset.select(self.getInputCol()).first()[0]) - ).shape[-1] - - # Flatten the array to a single array. - # Will do nothing if the array is not nested. - flattened_array_col = flatten_nested_arrays( - column=F.col(self.getInputCol()), column_data_type=input_column_type - ) - - mean_and_stddev_dict: Dict[str, float] = ( - dataset.select(F.explode(flattened_array_col).alias(self.getInputCol())) - .withColumn( - "mask", - F.when( - F.col(self.getInputCol()) == F.lit(self.getMaskValue()), 1 - ).otherwise(0), + # Persist so the array-size probe and the moments aggregation reuse a + # materialised result instead of re-scanning the upstream lineage twice. + # Guarded so we do not double-persist data the caller already cached. + already_cached = dataset.storageLevel.useMemory or dataset.storageLevel.useDisk + if not already_cached: + dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) + + try: + # Collect a single row to driver and get the length. + # We assume all subsequent rows have the same length. + array_size = np.array( + (dataset.select(self.getInputCol()).first()[0]) + ).shape[-1] + + # Flatten the array to a single array. + # Will do nothing if the array is not nested. + flattened_array_col = flatten_nested_arrays( + column=F.col(self.getInputCol()), column_data_type=input_column_type ) - .filter(F.col("mask") == F.lit(0)) - .agg( - F.mean(self.getInputCol()).alias("mean"), - F.stddev_pop(self.getInputCol()).alias("stddev"), + + mean_and_stddev_dict: Dict[str, float] = ( + dataset.select(F.explode(flattened_array_col).alias(self.getInputCol())) + .withColumn( + "mask", + F.when( + F.col(self.getInputCol()) == F.lit(self.getMaskValue()), 1 + ).otherwise(0), + ) + .filter(F.col("mask") == F.lit(0)) + .agg( + F.mean(self.getInputCol()).alias("mean"), + F.stddev_pop(self.getInputCol()).alias("stddev"), + ) + .first() + .asDict() ) - .first() - .asDict() - ) + finally: + if not already_cached: + dataset.unpersist() mean: List[float] = [mean_and_stddev_dict["mean"] for _ in range(array_size)] stddev: List[float] = [ mean_and_stddev_dict["stddev"] for _ in range(array_size) diff --git a/src/kamae/spark/estimators/standard_scale.py b/src/kamae/spark/estimators/standard_scale.py index ebf70f2a..379dce37 100644 --- a/src/kamae/spark/estimators/standard_scale.py +++ b/src/kamae/spark/estimators/standard_scale.py @@ -23,6 +23,7 @@ from pyspark import keyword_only from pyspark.sql import DataFrame from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType +from pyspark.storagelevel import StorageLevel from kamae.keras.core.backend import ALL_BACKENDS from kamae.spark.params import ( @@ -113,43 +114,54 @@ def _fit(self, dataset: DataFrame) -> "StandardScaleTransformer": else: input_col = F.col(self.getInputCol()) - # Collect a single row to driver and get the length. - # We assume all subsequent rows have the same length. - array_size = np.array((dataset.select(input_col).first()[0])).shape[-1] - - element_struct = construct_nested_elements_for_scaling( - column=input_col, - column_datatype=input_column_type, - array_dim=array_size, - ) - - mean_cols = [ - F.mean( - F.when( - F.col(f"element_struct.element_{i}") - == F.lit(self.getMaskValue()), - F.lit(None), - ).otherwise(F.col(f"element_struct.element_{i}")) - ).alias(f"mean_{i}") - for i in range(1, array_size + 1) - ] - - stddev_cols = [ - F.stddev_pop( - F.when( - F.col(f"element_struct.element_{i}") - == F.lit(self.getMaskValue()), - F.lit(None), - ).otherwise(F.col(f"element_struct.element_{i}")) - ).alias(f"stddev_{i}") - for i in range(1, array_size + 1) - ] - - metric_cols = mean_cols + stddev_cols - - mean_and_stddev_dict = ( - dataset.select(element_struct).agg(*metric_cols).first().asDict() - ) + # Persist so the array-size probe and the moments aggregation reuse a + # materialised result instead of re-scanning the upstream lineage twice. + # Guarded so we do not double-persist data the caller already cached. + already_cached = dataset.storageLevel.useMemory or dataset.storageLevel.useDisk + if not already_cached: + dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) + + try: + # Collect a single row to driver and get the length. + # We assume all subsequent rows have the same length. + array_size = np.array((dataset.select(input_col).first()[0])).shape[-1] + + element_struct = construct_nested_elements_for_scaling( + column=input_col, + column_datatype=input_column_type, + array_dim=array_size, + ) + + mean_cols = [ + F.mean( + F.when( + F.col(f"element_struct.element_{i}") + == F.lit(self.getMaskValue()), + F.lit(None), + ).otherwise(F.col(f"element_struct.element_{i}")) + ).alias(f"mean_{i}") + for i in range(1, array_size + 1) + ] + + stddev_cols = [ + F.stddev_pop( + F.when( + F.col(f"element_struct.element_{i}") + == F.lit(self.getMaskValue()), + F.lit(None), + ).otherwise(F.col(f"element_struct.element_{i}")) + ).alias(f"stddev_{i}") + for i in range(1, array_size + 1) + ] + + metric_cols = mean_cols + stddev_cols + + mean_and_stddev_dict = ( + dataset.select(element_struct).agg(*metric_cols).first().asDict() + ) + finally: + if not already_cached: + dataset.unpersist() mean = [mean_and_stddev_dict[f"mean_{i}"] for i in range(1, array_size + 1)] stddev = [mean_and_stddev_dict[f"stddev_{i}"] for i in range(1, array_size + 1)] diff --git a/tests/kamae/spark/conftest.py b/tests/kamae/spark/conftest.py index 0d356a68..ce7b2b84 100644 --- a/tests/kamae/spark/conftest.py +++ b/tests/kamae/spark/conftest.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import tempfile from typing import List, Optional import pytest @@ -43,8 +44,10 @@ def spark_session(): .config("spark.driver.memory", "2g") .getOrCreate() ) - yield spark - spark.stop() + with tempfile.TemporaryDirectory() as checkpoint_dir: + spark.sparkContext.setCheckpointDir(checkpoint_dir) + yield spark + spark.stop() @pytest.fixture(scope="module") diff --git a/tests/kamae/spark/pipeline/test_pipeline.py b/tests/kamae/spark/pipeline/test_pipeline.py index 03c7ddbd..3e0314a4 100644 --- a/tests/kamae/spark/pipeline/test_pipeline.py +++ b/tests/kamae/spark/pipeline/test_pipeline.py @@ -14,12 +14,18 @@ import os from shutil import rmtree +from unittest.mock import patch import pytest import tensorflow as tf +from pyspark.sql import DataFrame from pyspark.sql.types import DoubleType -from kamae.spark.estimators import StandardScaleEstimator, StringIndexEstimator +from kamae.spark.estimators import ( + ConditionalStandardScaleEstimator, + StandardScaleEstimator, + StringIndexEstimator, +) from kamae.spark.pipeline import KamaeSparkPipeline, KamaeSparkPipelineModel from kamae.spark.transformers import ( ArrayConcatenateTransformer, @@ -27,6 +33,7 @@ BucketizeTransformer, HashIndexTransformer, IdentityTransformer, + ListMeanTransformer, LogTransformer, SubtractTransformer, ) @@ -552,6 +559,274 @@ def test_spark_pipeline( diff = transformed_df.exceptAll(request.getfixturevalue(expected_dataframe)) assert diff.isEmpty(), f"PipelineKeras output is not the same as expected." + @pytest.mark.parametrize( + "stages", + [ + "valid_stages_1", + "valid_stages_2", + ], + ) + def test_spark_pipeline_checkpoint_is_transparent( + self, stages, example_dataframe, request + ): + """ + checkpoint(eager=True) only truncates lineage, so fitting with a positive + checkpointInterval must yield results identical to the default of 0. + """ + stages = request.getfixturevalue(stages) + + baseline_model = KamaeSparkPipeline(stages=stages, checkpointInterval=0).fit( + example_dataframe + ) + checkpointed_model = KamaeSparkPipeline( + stages=stages, checkpointInterval=2 + ).fit(example_dataframe) + + baseline_df = baseline_model.transform(example_dataframe) + checkpointed_df = checkpointed_model.transform(example_dataframe) + + assert baseline_df.schema == checkpointed_df.schema + assert baseline_df.exceptAll(checkpointed_df).isEmpty() + assert checkpointed_df.exceptAll(baseline_df).isEmpty() + + def test_spark_pipeline_checkpoint_invocation( + self, valid_stages_1, example_dataframe + ): + """ + checkpoint must be invoked during fit only when checkpointInterval > 0. + """ + original_checkpoint = DataFrame.checkpoint + + with patch.object( + DataFrame, + "checkpoint", + autospec=True, + side_effect=original_checkpoint, + ) as mock_checkpoint: + KamaeSparkPipeline(stages=valid_stages_1, checkpointInterval=0).fit( + example_dataframe + ) + assert mock_checkpoint.call_count == 0 + + mock_checkpoint.reset_mock() + KamaeSparkPipeline(stages=valid_stages_1, checkpointInterval=2).fit( + example_dataframe + ) + assert mock_checkpoint.call_count > 0 + + def test_spark_pipeline_checkpoint_bounds_plan_depth(self, spark_session): + """ + The point of checkpointInterval is to bound logical-plan depth. We build a + deep, linearly-dependent pipeline (every stage is an ancestor of the next, so + the working DataFrame is advanced at every fit and lineage keeps growing) and + capture the logical-plan size of the DataFrame handed to each estimator fit. + With a positive interval the plan must stay markedly smaller than the default. + """ + df = spark_session.createDataFrame( + [(1.0,), (4.0,), (7.0,), (2.0,), (9.0,)], + ["col0"], + ) + + num_blocks = 4 + transforms_per_block = 4 + + def build_stages(): + stages = [] + prev = "col0" + for b in range(num_blocks): + for t in range(transforms_per_block): + out = f"t_{b}_{t}" + stages.append( + SubtractTransformer( + inputCol=prev, outputCol=out, mathFloatConstant=1.0 + ) + ) + prev = out + out = f"s_{b}" + stages.append(StandardScaleEstimator(inputCol=prev, outputCol=out)) + prev = out + return stages + + def max_fit_plan_length(interval): + plan_lengths = [] + original_fit = StandardScaleEstimator.fit + + def spy_fit(estimator, dataset, *args, **kwargs): + plan = dataset._jdf.queryExecution().logical().toString() + plan_lengths.append(len(plan)) + return original_fit(estimator, dataset, *args, **kwargs) + + with patch.object(StandardScaleEstimator, "fit", spy_fit): + KamaeSparkPipeline( + stages=build_stages(), checkpointInterval=interval + ).fit(df) + return max(plan_lengths) + + baseline_max = max_fit_plan_length(0) + checkpointed_max = max_fit_plan_length(transforms_per_block + 1) + + # Checkpointing must keep the deepest fit-time plan well below the un-bounded + # baseline. A strict 2x margin is robust to Spark-version plan-string changes. + assert checkpointed_max * 2 < baseline_max, ( + f"plan not bounded: baseline_max={baseline_max}, " + f"checkpointed_max={checkpointed_max}" + ) + + def test_spark_pipeline_prunes_unused_input_columns( + self, valid_stages_1, example_dataframe + ): + """ + prune_unused_input_columns must keep the columns the pipeline reads + (col1/col2/col3 via ArrayConcatenate, col4 via StringIndex) and drop the + unused col5 and col1_col2_col3, while preserving every row. + """ + pipeline = KamaeSparkPipeline(stages=valid_stages_1) + + # The required set is generous (a superset), so assert containment rather + # than equality - it also carries output/param strings that harmlessly do + # not match any input-DataFrame column. + required = pipeline.collect_required_input_columns(valid_stages_1) + assert {"col1", "col2", "col3", "col4"}.issubset(required) + + pruned = pipeline.prune_unused_input_columns(example_dataframe, valid_stages_1) + + assert pruned.columns == ["col1", "col2", "col3", "col4"] + assert pruned.exceptAll( + example_dataframe.select("col1", "col2", "col3", "col4") + ).isEmpty() + + def test_collect_required_input_columns_includes_aux_columns(self): + """ + Aux columns read at fit time via params other than inputCol(s) - here + maskCols and relevanceCol on ConditionalStandardScaleEstimator, and + queryIdCol on a listwise transformer - must be reported by the collector so + pruning does not drop them. No Spark session needed. + """ + stages = [ + ConditionalStandardScaleEstimator( + inputCol="x", + outputCol="x_scaled", + maskCols=["m"], + relevanceCol="r", + ), + ListMeanTransformer( + inputCol="p", + outputCol="p_list_mean", + queryIdCol="q", + ), + ] + + required = KamaeSparkPipeline.collect_required_input_columns(stages) + + assert {"x", "m", "r", "p", "q"} <= required + + def test_collect_required_input_columns_plain_estimator(self): + """ + A stage with no aux column params must still report its inputCol and must + not gain spurious columns - confirms the aux sweep does not regress the + simple case. + """ + stages = [StandardScaleEstimator(inputCol="x", outputCol="x_scaled")] + + required = KamaeSparkPipeline.collect_required_input_columns(stages) + + assert "x" in required + + def test_spark_pipeline_prune_input_columns_is_opt_in( + self, valid_stages_1, example_dataframe + ): + """ + Pruning must only happen when pruneInputColumns is True. With the default + (False) the input DataFrame is not projected during fit. + """ + with patch.object( + KamaeSparkPipeline, + "prune_unused_input_columns", + autospec=True, + side_effect=KamaeSparkPipeline.prune_unused_input_columns, + ) as mock_prune: + KamaeSparkPipeline(stages=valid_stages_1).fit(example_dataframe) + assert mock_prune.call_count == 0 + + mock_prune.reset_mock() + KamaeSparkPipeline(stages=valid_stages_1, pruneInputColumns=True).fit( + example_dataframe + ) + assert mock_prune.call_count == 1 + + def test_spark_pipeline_prune_keeps_aux_fit_columns(self, spark_session): + """ + Regression: pruning must not drop columns an estimator reads at fit time + through params other than inputCol (here maskCols). The fit must not raise, + the genuinely-unused column must be pruned, and the fitted moments must be + identical to a prune-disabled baseline (numerically transparent). + """ + df = spark_session.createDataFrame( + [(1.0, 1, 3.0, 99.0), (2.0, 0, 1.0, 99.0), (3.0, 1, 2.0, 99.0)], + ["x", "m", "r", "junk"], + ) + + def build_pipeline(prune): + return KamaeSparkPipeline( + stages=[ + ConditionalStandardScaleEstimator( + inputCol="x", + outputCol="x_scaled", + maskCols=["m"], + maskOperators=["eq"], + maskValues=[1.0], + relevanceCol="r", + ), + ], + pruneInputColumns=prune, + ) + + pruned_pipeline = build_pipeline(prune=True) + + # Aux fit columns kept, genuinely-unused column dropped. + required = pruned_pipeline.collect_required_input_columns( + pruned_pipeline.getStages() + ) + assert {"x", "m", "r"} <= required + assert "junk" not in required + + # Must NOT raise UNRESOLVED_COLUMN / "Mask column m not found". + pruned_model = pruned_pipeline.fit(df) + baseline_model = build_pipeline(prune=False).fit(df) + + pruned_scaler = pruned_model.stages[-1] + baseline_scaler = baseline_model.stages[-1] + + assert pruned_scaler.getMean() == baseline_scaler.getMean() + assert pruned_scaler.getStddev() == baseline_scaler.getStddev() + + def test_spark_pipeline_prune_is_transparent_to_fit( + self, valid_stages_1, example_dataframe + ): + """ + Pruning drops only columns no stage reads, so a fitted model - and its + transform output - must be identical whether or not the input carries an + extra unused column when pruneInputColumns is enabled. + """ + with_extra = example_dataframe.withColumn( + "unused", example_dataframe["col1"] * 100.0 + ) + + baseline_out = ( + KamaeSparkPipeline(stages=valid_stages_1, pruneInputColumns=True) + .fit(example_dataframe) + .transform(example_dataframe) + ) + with_extra_out = ( + KamaeSparkPipeline(stages=valid_stages_1, pruneInputColumns=True) + .fit(with_extra) + .transform(example_dataframe) + ) + + assert baseline_out.schema == with_extra_out.schema + assert baseline_out.exceptAll(with_extra_out).isEmpty() + assert with_extra_out.exceptAll(baseline_out).isEmpty() + @pytest.mark.parametrize( "stages, input_col, original_dtype", [ From 94dea3994d79cd813a169a68d29fc9794ef69c0d Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:56:12 +0100 Subject: [PATCH 24/36] Patch black --- src/kamae/spark/utils/transform_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/kamae/spark/utils/transform_utils.py b/src/kamae/spark/utils/transform_utils.py index c2be3050..87425aeb 100644 --- a/src/kamae/spark/utils/transform_utils.py +++ b/src/kamae/spark/utils/transform_utils.py @@ -11,9 +11,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import pandas as pd from typing import Callable, List +import pandas as pd import pyspark.sql.functions as F from pyspark.sql import Column from pyspark.sql.types import ArrayType, DataType @@ -166,7 +166,6 @@ def _vectorized_func(series: pd.Series) -> pd.Series: return udf_func(input_col) - def single_input_single_output_scalar_udf_transform( input_col: Column, input_col_datatype: DataType, From 4f5cea285f0555a63e103a03903bc951ff3e3653 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:57:02 +0100 Subject: [PATCH 25/36] Patch error --- src/kamae/spark/estimators/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/kamae/spark/estimators/base.py b/src/kamae/spark/estimators/base.py index 90f651b1..dcd2a027 100644 --- a/src/kamae/spark/estimators/base.py +++ b/src/kamae/spark/estimators/base.py @@ -86,9 +86,9 @@ def fit( param_dict = { param[0].name: param[1] for param in self.extractParamMap().items() } - raise RuntimeError( + raise e.__class__( f"Error in estimator: {self.uid} with params: {param_dict}" - ) from e + ).with_traceback(e.__traceback__) def construct_layer_info(self) -> Dict[str, Any]: """ From 3878d2f308c74c1dd7a24e28e7652850cfbff564 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:57:54 +0100 Subject: [PATCH 26/36] Patch test case --- tests/kamae/spark/pipeline/test_pipeline.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/kamae/spark/pipeline/test_pipeline.py b/tests/kamae/spark/pipeline/test_pipeline.py index 3e0314a4..36607ed4 100644 --- a/tests/kamae/spark/pipeline/test_pipeline.py +++ b/tests/kamae/spark/pipeline/test_pipeline.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os +import tempfile from shutil import rmtree from unittest.mock import patch @@ -46,8 +46,7 @@ class TestPipeline: @pytest.fixture def test_dir(self): - path = "./tmp_test" - os.makedirs(path, exist_ok=True) + path = tempfile.mkdtemp() yield path rmtree(path) From 19180900e1c0664ae4465769ab9428ffbe79d84c Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:44:07 +0100 Subject: [PATCH 27/36] Remove complexity, remove comments --- src/kamae/spark/pipeline/pipeline.py | 188 ++++++++++----------------- 1 file changed, 69 insertions(+), 119 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index 4d4aeaae..a72a1797 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -40,53 +40,18 @@ class KamaeSparkPipeline(Pipeline): together BaseTransformers. It maintains the same functionality as pyspark.ml.Pipeline e.g. serialisation. - The `checkpointInterval` param optionally bounds the depth of the Spark - logical plan built up while fitting a multi-estimator pipeline. When set to a - positive integer it triggers a reliable `DataFrame.checkpoint(eager=True)` - every `checkpointInterval` stages (evaluated at estimator-fit action - boundaries), physically truncating the accumulated lineage. This is a - depth-bounding / reliability feature: it guards against deep-plan failures such - as "plan too large", 64KB codegen, and CodeCache-full errors, and avoids - re-executing the full upstream lineage on every estimator fit. - - Reliable checkpointing writes the intermediate DataFrame to the checkpoint - directory configured via `spark.sparkContext.setCheckpointDir()`, which - must point at fault-tolerant storage (DFS/cloud storage). Unlike local - checkpointing it survives executor loss (e.g. autoscaling, spot reclaim, OOM), - at the cost of writing to remote storage rather than executor-local disk. A - checkpoint directory MUST be set before fitting with a positive interval. Its - throughput impact is data-dependent and NOT guaranteed positive (the full, wide - intermediate DataFrame is persisted with no column pruning), so benchmark before - relying on it for speed. The default of 0 disables checkpointing entirely, - leaving fit behaviour byte-for-byte unchanged. - - The `cacheIntermediateData` param optionally persists the working DataFrame - (MEMORY_AND_DISK) at each estimator-fit boundary so that the estimator's fit - action - and any subsequent transforms - reuse a materialised result instead of - re-executing (and re-reading from source) the full upstream lineage on every - estimator. Only one intermediate frame is held at a time: each new persist - unpersists the one it supersedes, and the final frame is released before - returning. Unlike `checkpointInterval` it does not truncate the logical plan or - require a checkpoint directory; it is purely a re-scan-avoidance optimisation. - It preserves data exactly, so fitted results are identical to the default. The - default of False leaves fit behaviour unchanged. - - The `pruneInputColumns` param optionally projects the input DataFrame down to - only the columns the pipeline reads before fitting, dropping columns no stage - consumes so they are not carried through every transform and materialisation. - The set of columns to keep is computed generously (see - `collect_required_input_columns`) to avoid dropping columns referenced via - params other than `inputCol(s)`. The default of False leaves fit behaviour - unchanged. + Three opt-in fit optimisations are available, all defaulting off (fit behaviour + unchanged): `checkpointInterval` reliably checkpoints every N stages to bound + logical-plan depth (requires a checkpoint dir); `cacheIntermediateData` persists + the working DataFrame at each estimator-fit boundary to avoid re-scanning the + upstream lineage; `pruneInputColumns` drops input columns no stage consumes. """ checkpointInterval = Param( Params._dummy(), "checkpointInterval", - "Number of stages between reliable checkpoint(eager=True) calls during " - "fit, used to bound logical-plan depth. Requires a checkpoint directory set " - "via spark.sparkContext.setCheckpointDir. 0 (the default) disables " - "checkpointing and leaves fit behaviour exactly unchanged.", + "Stages between reliable checkpoint(eager=True) calls during fit, to bound " + "logical-plan depth. Requires a checkpoint dir. 0 (default) disables it.", typeConverter=TypeConverters.toInt, ) @@ -94,18 +59,16 @@ class KamaeSparkPipeline(Pipeline): Params._dummy(), "cacheIntermediateData", "If True, persist the working DataFrame (MEMORY_AND_DISK) at each " - "estimator-fit boundary so estimator fits reuse a materialised result " - "rather than re-scanning the upstream lineage from source. False (the " - "default) leaves fit behaviour exactly unchanged.", + "estimator-fit boundary to avoid re-scanning the upstream lineage. False " + "(default) disables it.", typeConverter=TypeConverters.toBoolean, ) pruneInputColumns = Param( Params._dummy(), "pruneInputColumns", - "If True, project the input DataFrame down to only the columns the " - "pipeline reads before fitting, dropping columns no stage consumes. False " - "(the default) leaves fit behaviour exactly unchanged.", + "If True, drop input columns no stage consumes before fitting. False " + "(default) disables it.", typeConverter=TypeConverters.toBoolean, ) @@ -293,21 +256,10 @@ def collect_required_input_columns( """ Collects every column potentially read by any stage in the pipeline. - A raw input-DataFrame column absent from this set is consumed by no stage - and can be dropped before fitting, so it is not carried through every - transform (and every materialisation) below. - - This is deliberately generous: as well as the canonical inputs - (`inputCol`/`inputCols` from `get_layer_inputs_outputs`), it unions the - value(s) of every param whose name ends in `Col`/`Cols`. Some stages read - extra columns during fit through such params (e.g. `maskCols` and - `relevanceCol` on ConditionalStandardScaleEstimator, `queryIdCol` on - listwise transformers) that `inputCol(s)` does not capture. The name - convention holds across all estimators and transformers, so this - self-maintaining heuristic covers future aux column params without a - per-stage allowlist. Over-inclusion is harmless - a name that is not a real - input column simply never matches `dataset.columns` - whereas omitting a - referenced column would wrongly drop data the pipeline needs at fit time. + Generous by design: unions canonical inputs with the value(s) of every param + whose name ends in `Col`/`Cols`, so aux columns read during fit are not + missed. Over-inclusion is harmless (names not matching `dataset.columns` are + ignored); omission would wrongly drop data the pipeline needs. :param stages: List of pipeline stages. :returns: Set of column names potentially read by at least one stage. @@ -338,12 +290,7 @@ def prune_unused_input_columns( """ Projects the input DataFrame down to only the columns the pipeline reads. - Columns produced by stages are created downstream via `withColumn`, so only - the pipeline's source columns need to be present up front. Pruning here - keeps the frame narrow before any expansion, reducing the cost of every - subsequent transform and materialisation. If no unused columns are found - (or the pipeline reads none of the DataFrame's columns) the DataFrame is - returned unchanged. + Returned unchanged if there are no unused columns to drop. :param dataset: Input DataFrame to prune. :param stages: Expanded pipeline stages. @@ -355,6 +302,45 @@ def prune_unused_input_columns( return dataset.select(*columns_to_keep) return dataset + @staticmethod + def _validate_stage_types(stages: List["KamaePipelineStage"]) -> None: + """ + Ensures every expanded stage is a recognised estimator or transformer. + + :param stages: Expanded pipeline stages. + :raises TypeError: If any stage is not a BaseEstimator or BaseTransformer. + """ + for stage in stages: + if not isinstance(stage, (BaseEstimator, BaseTransformer)): + raise TypeError( + "Cannot recognize a pipeline stage of type %s." % type(stage) + ) + + @staticmethod + def _resolve_checkpoint_enabled( + dataset: DataFrame, checkpoint_interval: Optional[int] + ) -> bool: + """ + Determines whether checkpointing is enabled and validates its prerequisites. + + Fails fast if enabled without a checkpoint dir, rather than raising mid-fit. + + :param dataset: DataFrame whose SparkContext is checked for a checkpoint dir. + :param checkpoint_interval: Configured checkpoint interval (0/None disables). + :returns: True if checkpointing is enabled, False otherwise. + :raises ValueError: If enabled but no checkpoint directory has been set. + """ + checkpoint_enabled = checkpoint_interval is not None and checkpoint_interval > 0 + if ( + checkpoint_enabled + and dataset.sparkSession.sparkContext.getCheckpointDir() is None + ): + raise ValueError( + "checkpointInterval > 0 requires a checkpoint directory. Set one via " + "spark.sparkContext.setCheckpointDir() before fitting." + ) + return checkpoint_enabled + def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": """ Fits the pipeline to the dataset. Returns a KamaeSparkPipelineModel object. @@ -362,26 +348,9 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": Calls the super fit method of the pyspark.ml.Pipeline class and then constructs a KamaeSparkPipelineModel uses the stages from the fit pipeline. - If `pruneInputColumns` is True, the input DataFrame is projected down to - only the columns the pipeline reads (see `prune_unused_input_columns`) - before fitting, so columns no stage consumes are not carried through every - transform and materialisation. The default of False leaves fit behaviour - unchanged. - - If `checkpointInterval` is a positive integer, the working DataFrame is - reliably checkpointed via `checkpoint(eager=True)` roughly every - `checkpointInterval` stages (at estimator-fit action boundaries) to bound - logical-plan depth. checkpoint(eager=True) preserves the data exactly and - only truncates lineage, so fitted results are numerically identical to the - default (interval=0) behaviour. A checkpoint directory must be configured via - `spark.sparkContext.setCheckpointDir` before fitting with a positive interval. - The default of 0 (or None) disables checkpointing entirely. - - If `cacheIntermediateData` is True, the working DataFrame is persisted - (MEMORY_AND_DISK) at each estimator-fit boundary so the fit action and any - subsequent transform reuse a materialised result rather than re-scanning the - upstream lineage. Persistence preserves data exactly, so fitted results are - identical to the default (False) behaviour. The default of False disables it. + Optionally applies the opt-in fit optimisations (`pruneInputColumns`, + `checkpointInterval`, `cacheIntermediateData`); see the class docstring. All + preserve data exactly, so fitted results match the defaults-off behaviour. :param dataset: PySpark DataFrame to fit the pipeline to. :returns: KamaeSparkPipelineModel object. @@ -389,18 +358,9 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": has been set on the SparkContext. """ expanded_pipeline_stages = self.expand_pipeline_stages() + self._validate_stage_types(expanded_pipeline_stages) - for stage in expanded_pipeline_stages: - if not ( - isinstance(stage, BaseEstimator) or isinstance(stage, BaseTransformer) - ): - raise TypeError( - "Cannot recognize a pipeline stage of type %s." % type(stage) - ) - - # Optional, opt-in. Drop input columns no stage consumes before any - # expansion, so dead columns are not carried through every transform and - # materialisation below. Default False keeps behaviour unchanged. + # Opt-in: drop input columns no stage consumes. Default False = no change. if self.getPruneInputColumns(): dataset = self.prune_unused_input_columns(dataset, expanded_pipeline_stages) @@ -414,23 +374,14 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": estimator_parent_stages = self.collect_estimator_parents( expanded_pipeline_stages ) - # Optional, opt-in plan-depth bounding. 0 (or None) keeps behaviour unchanged. + # Opt-in plan-depth bounding. 0 (or None) = no change. checkpoint_interval = self.getCheckpointInterval() - checkpoint_enabled = checkpoint_interval is not None and checkpoint_interval > 0 - # Reliable checkpoint() requires a checkpoint directory; fail fast with a clear - # message rather than letting Spark raise mid-fit after work has been done. - if ( - checkpoint_enabled - and dataset.sparkSession.sparkContext.getCheckpointDir() is None - ): - raise ValueError( - "checkpointInterval > 0 requires a checkpoint directory. Set one via " - "spark.sparkContext.setCheckpointDir() before fitting." - ) + checkpoint_enabled = self._resolve_checkpoint_enabled( + dataset, checkpoint_interval + ) cache_enabled = self.getCacheIntermediateData() last_checkpoint_index = 0 - # Holds the single intermediate frame currently persisted (if any) so it can - # be unpersisted once superseded or once fitting completes. + # The single persisted frame (if any), unpersisted once superseded or done. cached_dataset: Optional[DataFrame] = None # Fit each stage, appending the transformer to the list of transformers # If the stage is a parent of an estimator, transform the dataset. @@ -441,17 +392,16 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": if stage in estimator_parent_stages: dataset = stage.transform(dataset) else: - # Truncate the accumulated lineage just before the fit action so the - # plan is physically bounded. eager=True forces materialisation now. + # Truncate accumulated lineage before the fit action to bound plan + # depth. eager=True materialises now. if ( checkpoint_enabled and index - last_checkpoint_index >= checkpoint_interval ): dataset = dataset.checkpoint(eager=True) last_checkpoint_index = index - # Persist the working frame so the fit action and any subsequent - # transform read a materialised result rather than re-scanning the - # upstream lineage from source. Only one frame is held at a time. + # Persist so the fit action and downstream transforms reuse a + # materialised frame instead of re-scanning. One frame held at a time. if cache_enabled: new_cached = dataset.persist(StorageLevel.MEMORY_AND_DISK) if cached_dataset is not None: From 96fc18444d4a28a733fbbaca97df6d9af4543a56 Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:53:46 +0100 Subject: [PATCH 28/36] Patch for lint --- src/kamae/spark/pipeline/pipeline.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index a72a1797..c1901ca9 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -454,7 +454,7 @@ class KamaeSparkPipelineReader(PipelineReader): Util class for reading a pipeline from a persistent storage path. """ - def __init__(self, cls: Type[KamaeSparkPipeline]): + def __init__(self, cls: Type[KamaeSparkPipeline]) -> None: super().__init__(cls=cls) def load(self, path: str) -> KamaeSparkPipeline: @@ -474,5 +474,5 @@ class KamaeSparkPipelineWriter(PipelineWriter): Util class for writing a pipeline to a persistent storage path. """ - def __init__(self, instance: KamaeSparkPipeline): + def __init__(self, instance: KamaeSparkPipeline) -> None: super().__init__(instance=instance) From 6af4ba4a4a0e4d17602aa2b3364010d55fa6d73a Mon Sep 17 00:00:00 2001 From: Conor worthington <45695214+ConorWorthington@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:54:15 +0100 Subject: [PATCH 29/36] Patch for linter --- src/kamae/spark/transformers/bucketize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/kamae/spark/transformers/bucketize.py b/src/kamae/spark/transformers/bucketize.py index 022a0042..db743f53 100644 --- a/src/kamae/spark/transformers/bucketize.py +++ b/src/kamae/spark/transformers/bucketize.py @@ -49,7 +49,7 @@ class BucketizeParams(Params): ) @staticmethod - def check_splits_sorted(splits: List[float]): + def check_splits_sorted(splits: List[float]) -> None: """ Checks that the splits parameter is sorted. From 78fef71a4856d1187258e207136df9576c68413a Mon Sep 17 00:00:00 2001 From: cworthington Date: Wed, 12 Aug 2026 13:43:07 +0100 Subject: [PATCH 30/36] refactor: address PR review feedback on pipeline fit optimisations - 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 --- src/kamae/spark/pipeline/pipeline.py | 46 +++++++++++++-------- tests/kamae/spark/pipeline/test_pipeline.py | 16 +++++-- uv.lock | 31 ++++++++++++++ 3 files changed, 72 insertions(+), 21 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index c1901ca9..22a22bb0 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -51,7 +51,7 @@ class KamaeSparkPipeline(Pipeline): Params._dummy(), "checkpointInterval", "Stages between reliable checkpoint(eager=True) calls during fit, to bound " - "logical-plan depth. Requires a checkpoint dir. 0 (default) disables it.", + "logical-plan depth. Requires a checkpoint dir. None (default) disables it.", typeConverter=TypeConverters.toInt, ) @@ -77,7 +77,7 @@ def __init__( self, *, stages: Optional[List["KamaePipelineStage"]] = None, - checkpointInterval: int = 0, + checkpointInterval: Optional[int] = None, cacheIntermediateData: bool = False, pruneInputColumns: bool = False, ) -> None: @@ -86,7 +86,7 @@ def __init__( :param stages: List of LayerTransformers to chain together. :param checkpointInterval: Number of stages between reliable - checkpoint(eager=True) calls during fit. 0 (default) disables it. + checkpoint(eager=True) calls during fit. None (default) disables it. :param cacheIntermediateData: If True, persist the working DataFrame at each estimator-fit boundary to avoid re-scanning the upstream lineage. False (default) disables it. @@ -95,9 +95,9 @@ def __init__( :returns: None - class instantiated. """ kwargs = self._input_kwargs - super().__init__() + super().__init__(stages=stages) self._setDefault( - checkpointInterval=0, + checkpointInterval=None, cacheIntermediateData=False, pruneInputColumns=False, ) @@ -120,17 +120,23 @@ def getStages(self) -> List["KamaePipelineStage"]: """ return self.getOrDefault("stages") - def setCheckpointInterval(self, value: int) -> "KamaeSparkPipeline": + def setCheckpointInterval(self, value: Optional[int]) -> "KamaeSparkPipeline": """ Sets the `checkpointInterval` parameter. - :param value: Number of stages between reliable checkpoint calls during - fit. 0 (or None) disables checkpointing. + :param value: Positive number of stages between reliable checkpoint calls + during fit. None disables checkpointing. :returns: KamaeSparkPipeline object with checkpointInterval set. + :raises ValueError: If value is not None and not a positive integer. """ + if value is not None and value < 1: + raise ValueError( + "checkpointInterval must be a positive integer or None, got " + f"{value}." + ) return self._set(checkpointInterval=value) - def getCheckpointInterval(self) -> int: + def getCheckpointInterval(self) -> Optional[int]: """ Gets the value of the `checkpointInterval` parameter. @@ -178,24 +184,29 @@ def setParams( self, *, stages: Optional["KamaePipelineStage"] = None, - checkpointInterval: int = 0, + checkpointInterval: Optional[int] = None, cacheIntermediateData: bool = False, pruneInputColumns: bool = False, ) -> "KamaeSparkPipeline": """ Sets the keyword arguments of the pipeline. + Routes each supplied param through its setter so setter-level validation + (e.g. checkpointInterval) runs. + :param stages: List of pipeline stages. :param checkpointInterval: Number of stages between reliable - checkpoint(eager=True) calls during fit. 0 (default) disables it. + checkpoint(eager=True) calls during fit. None (default) disables it. :param cacheIntermediateData: If True, persist the working DataFrame at each estimator-fit boundary. False (default) disables it. :param pruneInputColumns: If True, drop input columns no stage consumes before fitting. False (default) disables it. :returns: KamaeSparkPipeline object with params set. """ - kwargs = self._input_kwargs - return self._set(**kwargs) + for param_name, param_value in self._input_kwargs.items(): + setter = getattr(self, f"set{param_name[0].upper()}{param_name[1:]}") + setter(param_value) + return self def expand_pipeline_stages(self) -> List["KamaePipelineStage"]: """ @@ -257,9 +268,10 @@ def collect_required_input_columns( Collects every column potentially read by any stage in the pipeline. Generous by design: unions canonical inputs with the value(s) of every param - whose name ends in `Col`/`Cols`, so aux columns read during fit are not - missed. Over-inclusion is harmless (names not matching `dataset.columns` are - ignored); omission would wrongly drop data the pipeline needs. + whose name ends in `Col`/`Cols`, so aux columns read during fit (e.g. + maskCols, relevanceCol, queryIdCol) are not missed. Over-inclusion is + harmless (names not matching `dataset.columns` are ignored); omission would + wrongly drop data the pipeline needs at fit time. :param stages: List of pipeline stages. :returns: Set of column names potentially read by at least one stage. @@ -298,7 +310,7 @@ def prune_unused_input_columns( """ 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): + if columns_to_keep: return dataset.select(*columns_to_keep) return dataset diff --git a/tests/kamae/spark/pipeline/test_pipeline.py b/tests/kamae/spark/pipeline/test_pipeline.py index 36607ed4..997ba854 100644 --- a/tests/kamae/spark/pipeline/test_pipeline.py +++ b/tests/kamae/spark/pipeline/test_pipeline.py @@ -570,11 +570,11 @@ def test_spark_pipeline_checkpoint_is_transparent( ): """ checkpoint(eager=True) only truncates lineage, so fitting with a positive - checkpointInterval must yield results identical to the default of 0. + checkpointInterval must yield results identical to the default (None). """ stages = request.getfixturevalue(stages) - baseline_model = KamaeSparkPipeline(stages=stages, checkpointInterval=0).fit( + baseline_model = KamaeSparkPipeline(stages=stages, checkpointInterval=None).fit( example_dataframe ) checkpointed_model = KamaeSparkPipeline( @@ -602,7 +602,7 @@ def test_spark_pipeline_checkpoint_invocation( autospec=True, side_effect=original_checkpoint, ) as mock_checkpoint: - KamaeSparkPipeline(stages=valid_stages_1, checkpointInterval=0).fit( + KamaeSparkPipeline(stages=valid_stages_1, checkpointInterval=None).fit( example_dataframe ) assert mock_checkpoint.call_count == 0 @@ -613,6 +613,14 @@ def test_spark_pipeline_checkpoint_invocation( ) assert mock_checkpoint.call_count > 0 + @pytest.mark.parametrize("bad_value", [0, -1, -5]) + def test_spark_pipeline_checkpoint_interval_rejects_non_positive(self, bad_value): + """ + checkpointInterval must be a positive integer or None; 0 and negatives raise. + """ + with pytest.raises(ValueError): + KamaeSparkPipeline(checkpointInterval=bad_value) + def test_spark_pipeline_checkpoint_bounds_plan_depth(self, spark_session): """ The point of checkpointInterval is to bound logical-plan depth. We build a @@ -661,7 +669,7 @@ def spy_fit(estimator, dataset, *args, **kwargs): ).fit(df) return max(plan_lengths) - baseline_max = max_fit_plan_length(0) + baseline_max = max_fit_plan_length(None) checkpointed_max = max_fit_plan_length(transforms_per_block + 1) # Checkpointing must keep the deepest fit-time plan well below the un-bounded diff --git a/uv.lock b/uv.lock index 68ac512a..ab5c5c8d 100644 --- a/uv.lock +++ b/uv.lock @@ -873,6 +873,7 @@ dependencies = [ { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pandas", version = "1.5.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "pyarrow" }, { name = "pyfarmhash" }, { name = "pyspark" }, { name = "tensorflow" }, @@ -922,6 +923,7 @@ requires-dist = [ { name = "networkx", specifier = ">=2.6.3,<3.0.0" }, { name = "numpy", specifier = ">=1.22.0,<2.0.0" }, { name = "pandas", specifier = ">=1.3.4,<3.0.0" }, + { name = "pyarrow", specifier = ">=4.0.0" }, { name = "pyfarmhash", specifier = ">=0.3.2,<0.4.0" }, { name = "pyspark", specifier = ">=3.4.0,<4.0.0" }, { name = "tensorflow", specifier = ">=2.16.0,<3.0.0" }, @@ -1959,6 +1961,35 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/30/a58b32568f1623aaad7db22aa9eafc4c6c194b429ff35bdc55ca2726da47/py4j-0.10.9.7-py2.py3-none-any.whl", hash = "sha256:85defdfd2b2376eb3abf5ca6474b51ab7e0de341c75a02f46dc9b5976f5a5c1b", size = 200481 }, ] +[[package]] +name = "pyarrow" +version = "25.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3d/e3/27f57f80141379d60defe6703eb50a707325706f07fedfd1312c7a751995/pyarrow-25.0.1.tar.gz", hash = "sha256:9150a83248bfed9813ea3c3af74c3856c1984d444aa28e58bf7733b9750ddf6a", size = 1201653 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0a/3e/5cd70becb51e1d044c54ba5e627424a6e87df5b98008cbd22cc6abd409ca/pyarrow-25.0.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:0b1edbb2f385a6a65e9711b62ba86ac54a7816a3f8d17bb3e8a5929d65fb2485", size = 35954271 }, + { url = "https://files.pythonhosted.org/packages/64/be/17599e086df264ea7dc221d1101e3131e181e00da428a2f9bd0358f0d06b/pyarrow-25.0.1-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:a4dd8bf99a8fac133efc0ed6a92f5fddbe2adba0d0f6dd720e39ba9855cea85c", size = 37647543 }, + { url = "https://files.pythonhosted.org/packages/42/34/e138b451fd3970a6eda4599f68ae3b2b32b661bc958de3239d54a0bf6575/pyarrow-25.0.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:bddd0c4f7630c2a3ddf6347c1bdaa79d97bcf6bd445f9e60c816b7d77c85a5ae", size = 46837120 }, + { url = "https://files.pythonhosted.org/packages/57/5c/f8fc0eb2de03464a557d5a4d0c15e972d73362414696618833b771f7eddd/pyarrow-25.0.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:a4d6d5e9a3d1879a97c08ded0c797579b7965eafd0f0c26c30b45ccc06db939b", size = 50066460 }, + { url = "https://files.pythonhosted.org/packages/3f/d1/0dd64fd06de0333b808a02f60981635f067b71aad3a30698a9a104fae778/pyarrow-25.0.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:514ddb60285631af068875550c90eddc181db3e8e63a032b1559be189e82f056", size = 49937892 }, + { url = "https://files.pythonhosted.org/packages/cb/3c/f89d1bd76d5f3284c2a44d7d7ebbd8204535e5ae2b41f4077069b4ff2ec6/pyarrow-25.0.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:cab40b1edfef0262e0e5251aa2c58d75630f24d06dd7794480243acc001a1d7d", size = 53107240 }, + { url = "https://files.pythonhosted.org/packages/67/67/b554a8e09f3f3decccf405eb8fbe86696321cbcb5b62d18b4a5057a4c113/pyarrow-25.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:60e89d8f13861a1f7f8d950fa54aebb8023b30734d0ac51ffa80beabe2df4bba", size = 27848683 }, + { url = "https://files.pythonhosted.org/packages/ee/8b/0d23b47702fcfe8b3618d5292035099675c5a1c48258932350c08020f7b5/pyarrow-25.0.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:51093dd9e10325fbdb3c10a2ae7c4806e5c822d94e74ae4938b26524a3323fee", size = 35946180 }, + { url = "https://files.pythonhosted.org/packages/d8/17/707d17a5476c55a9541fde0db8213ac30979a792864d72415f176ba50c45/pyarrow-25.0.1-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:eb6203482ff3746a5632303a7279ae0b5a304c46985b49ed1378cb350ea6728d", size = 37644787 }, + { url = "https://files.pythonhosted.org/packages/c1/b2/cdc98ecf1a6408280bc3a6a07054cdd99a3f4670acc0545d383ce113e87d/pyarrow-25.0.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:880523be3d29efcf83d3998835d206118ccf35e3871dbd2fb60408cf6b007a80", size = 46834633 }, + { url = "https://files.pythonhosted.org/packages/c8/6e/d3fafc41f378b2c65be43b827798c0fae42049a641c8526633ed3eb573e2/pyarrow-25.0.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:25f8720bf6387d5dc2ebd2622112de630760419e4b66134405dd24110d15f37e", size = 50065507 }, + { url = "https://files.pythonhosted.org/packages/d5/12/8d0698954b8c3001844a898e0a6900bebe83d7ee40c11195174c5122f324/pyarrow-25.0.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:4facd65742a024a4a366328a1d2292062d72d6e023c1b7dda8d4c37544933a25", size = 49955690 }, + { url = "https://files.pythonhosted.org/packages/d3/0b/1ecb936ac6409e90a34d58eea1c7cec09a9ae6d2141b9e49ad01a2b1ea47/pyarrow-25.0.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:aa0559502e1cd6254d6814614085dd9c5a3dd0419362978a936a3f68a9e5c3df", size = 53128198 }, + { url = "https://files.pythonhosted.org/packages/8e/1c/5236033550633c9b7377b2a53660b2bbb06cb06dc09c4356332d67643ca1/pyarrow-25.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:62cd0d785b8aa6675ee355f9fc02252a340f4441257c42674937826fd7594325", size = 27857263 }, + { url = "https://files.pythonhosted.org/packages/a6/e2/9ab15b88cbfac28e16419ce5439ec29234c5172cb8259301b4ba639bdec0/pyarrow-25.0.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:df961f2e7ae9cf496459259d798652c70625f6c080650d6952f8c04053c58ee9", size = 35861559 }, + { url = "https://files.pythonhosted.org/packages/58/79/a0036dbe1eabe1f73127427342f1d99982584c4a2cde2651d6c93499c6f6/pyarrow-25.0.1-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:cc4aa407fde9fc660be3939e49ea31f50f3e9fec17c0ec63159f7711edd3efc9", size = 37628383 }, + { url = "https://files.pythonhosted.org/packages/13/49/d93a57d375f4bf0cf82913dd6bb54acafde83dd993be2282c81ac5616cad/pyarrow-25.0.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:4340f0ba6c1d2e13f21658de1d7c662ca2545018568d0030a1e9afca159d87e3", size = 46820190 }, + { url = "https://files.pythonhosted.org/packages/60/c9/711ca85d79f1ec98f29a5eae2b051e25b4ecec5de3e3c0e2d5c5dcb15664/pyarrow-25.0.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:5389cdf79447ed1515c9e31620e6e1e2302249564d603f2ad727d4f6d313e4c3", size = 50102437 }, + { url = "https://files.pythonhosted.org/packages/80/53/8fb8359ff17cfb6263a1cf3ebf7caec9fe197de118719e84fcb1d0618026/pyarrow-25.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d51592cb7561e87877c506113e7adbf1342ab579e6c21f0ef44b8ba41cb74c80", size = 49942424 }, + { url = "https://files.pythonhosted.org/packages/e8/83/4e5ae02a9341571b18a6fca380ac7a58ce6ddae7ab3c060208c0a1e79f02/pyarrow-25.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6109c94d8b9f3b17a041daca16cacb2f651ad8f1ef70a4232c2c0f37a23da2a8", size = 53144206 }, + { url = "https://files.pythonhosted.org/packages/65/ee/197cbf47e49f83e6ebeb946a5259a48a638dea27ac774db42fe78022179d/pyarrow-25.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:8858d7bfc22e3f51529aeaa4077225029724623e4595dc9eff8c793935c34140", size = 27953934 }, +] + [[package]] name = "pycodestyle" version = "2.12.1" From 60fa70bacdcc13f3e2c2b3d26b822f1e95d7e640 Mon Sep 17 00:00:00 2001 From: cworthington Date: Mon, 17 Aug 2026 14:35:52 +0100 Subject: [PATCH 31/36] refactor: annotate bucketize get_keras_layer return as tf.keras.layers.Layer BucketizeLayer is TensorFlow-only, so the tf-specific return type is accurate. Co-Authored-By: Claude Opus 4.7 --- src/kamae/spark/transformers/bucketize.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/kamae/spark/transformers/bucketize.py b/src/kamae/spark/transformers/bucketize.py index db743f53..e9055cf7 100644 --- a/src/kamae/spark/transformers/bucketize.py +++ b/src/kamae/spark/transformers/bucketize.py @@ -19,8 +19,8 @@ from functools import reduce from typing import List, Optional -import keras import pyspark.sql.functions as F +import tensorflow as tf from pyspark import keyword_only from pyspark.ml.param import Param, Params, TypeConverters from pyspark.sql import Column, DataFrame @@ -167,7 +167,7 @@ def bucketize(value: Column) -> Column: output_col, ) - def get_keras_layer(self) -> keras.layers.Layer: + def get_keras_layer(self) -> tf.keras.layers.Layer: """ Gets the Keras layer for the BucketizeLayer transformer. From 36859af0c66be4ce45e3f52a5ed77bac15a93777 Mon Sep 17 00:00:00 2001 From: cworthington Date: Tue, 18 Aug 2026 18:09:21 +0100 Subject: [PATCH 32/36] feat: add opt-in cacheEstimatorInput fit optimisation to KamaeSparkPipeline 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 --- src/kamae/spark/pipeline/pipeline.py | 88 ++++++++++++++++++++- tests/kamae/spark/pipeline/test_pipeline.py | 80 +++++++++++++++++++ 2 files changed, 164 insertions(+), 4 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index 22a22bb0..57b5aee4 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import warnings from typing import TYPE_CHECKING, List, Optional, Set, Type import networkx as nx @@ -40,11 +41,14 @@ class KamaeSparkPipeline(Pipeline): together BaseTransformers. It maintains the same functionality as pyspark.ml.Pipeline e.g. serialisation. - Three opt-in fit optimisations are available, all defaulting off (fit behaviour + Four opt-in fit optimisations are available, all defaulting off (fit behaviour unchanged): `checkpointInterval` reliably checkpoints every N stages to bound logical-plan depth (requires a checkpoint dir); `cacheIntermediateData` persists the working DataFrame at each estimator-fit boundary to avoid re-scanning the - upstream lineage; `pruneInputColumns` drops input columns no stage consumes. + upstream lineage; `pruneInputColumns` drops input columns no stage consumes; + `cacheEstimatorInput` projects to the columns still read downstream at the first + estimator boundary and persists that narrow frame once, so independent sibling + estimators reuse it instead of re-scanning the wide input. """ checkpointInterval = Param( @@ -72,6 +76,16 @@ class KamaeSparkPipeline(Pipeline): typeConverter=TypeConverters.toBoolean, ) + cacheEstimatorInput = Param( + Params._dummy(), + "cacheEstimatorInput", + "If True, at the first estimator-fit boundary project the working DataFrame " + "to the columns still read downstream and persist (MEMORY_AND_DISK) that " + "narrow frame once, reused by all subsequent estimators. Competes with " + "cacheIntermediateData; enable at most one. False (default) disables it.", + typeConverter=TypeConverters.toBoolean, + ) + @keyword_only def __init__( self, @@ -80,6 +94,7 @@ def __init__( checkpointInterval: Optional[int] = None, cacheIntermediateData: bool = False, pruneInputColumns: bool = False, + cacheEstimatorInput: bool = False, ) -> None: """ Initialises the KamaeSparkPipeline object. @@ -92,6 +107,9 @@ def __init__( False (default) disables it. :param pruneInputColumns: If True, drop input columns no stage consumes before fitting. False (default) disables it. + :param cacheEstimatorInput: If True, project to the columns still read + downstream at the first estimator boundary and persist that narrow frame + once for reuse by subsequent estimators. False (default) disables it. :returns: None - class instantiated. """ kwargs = self._input_kwargs @@ -100,6 +118,7 @@ def __init__( checkpointInterval=None, cacheIntermediateData=False, pruneInputColumns=False, + cacheEstimatorInput=False, ) self.setParams(**kwargs) @@ -179,6 +198,24 @@ def getPruneInputColumns(self) -> bool: """ return self.getOrDefault(self.pruneInputColumns) + def setCacheEstimatorInput(self, value: bool) -> "KamaeSparkPipeline": + """ + Sets the `cacheEstimatorInput` parameter. + + :param value: Whether to project and persist a narrow estimator-input + frame once at the first estimator boundary during fit. + :returns: KamaeSparkPipeline object with cacheEstimatorInput set. + """ + return self._set(cacheEstimatorInput=value) + + def getCacheEstimatorInput(self) -> bool: + """ + Gets the value of the `cacheEstimatorInput` parameter. + + :returns: The cacheEstimatorInput value. + """ + return self.getOrDefault(self.cacheEstimatorInput) + @keyword_only def setParams( self, @@ -187,6 +224,7 @@ def setParams( checkpointInterval: Optional[int] = None, cacheIntermediateData: bool = False, pruneInputColumns: bool = False, + cacheEstimatorInput: bool = False, ) -> "KamaeSparkPipeline": """ Sets the keyword arguments of the pipeline. @@ -201,6 +239,9 @@ def setParams( each estimator-fit boundary. False (default) disables it. :param pruneInputColumns: If True, drop input columns no stage consumes before fitting. False (default) disables it. + :param cacheEstimatorInput: If True, project to the columns still read + downstream at the first estimator boundary and persist that narrow frame + once for reuse by subsequent estimators. False (default) disables it. :returns: KamaeSparkPipeline object with params set. """ for param_name, param_value in self._input_kwargs.items(): @@ -361,8 +402,14 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": then constructs a KamaeSparkPipelineModel uses the stages from the fit pipeline. Optionally applies the opt-in fit optimisations (`pruneInputColumns`, - `checkpointInterval`, `cacheIntermediateData`); see the class docstring. All - preserve data exactly, so fitted results match the defaults-off behaviour. + `checkpointInterval`, `cacheIntermediateData`, `cacheEstimatorInput`); see the + class docstring. All preserve data exactly, so fitted results match the + defaults-off behaviour. + + If both `cacheIntermediateData` and `cacheEstimatorInput` are enabled, + `cacheEstimatorInput` takes precedence (a warning is emitted) and + `cacheIntermediateData` is ignored, since the narrow frame is a strictly + smaller cache and the intermediate cache would evict it. :param dataset: PySpark DataFrame to fit the pipeline to. :returns: KamaeSparkPipelineModel object. @@ -392,9 +439,24 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": dataset, checkpoint_interval ) cache_enabled = self.getCacheIntermediateData() + cache_estimator_input = self.getCacheEstimatorInput() + # Competing caching strategies - both would persist the wide frame, and the + # intermediate cache's per-boundary unpersist would evict the narrow frame. + # cacheEstimatorInput wins: it persists a strictly narrower frame once. + if cache_enabled and cache_estimator_input: + warnings.warn( + "cacheIntermediateData and cacheEstimatorInput are competing " + "caching strategies; cacheEstimatorInput takes precedence and " + "cacheIntermediateData is ignored.", + stacklevel=2, + ) + cache_enabled = False last_checkpoint_index = 0 # The single persisted frame (if any), unpersisted once superseded or done. cached_dataset: Optional[DataFrame] = None + # The narrow estimator-input frame (if any), persisted once and reused. + estimator_input_cache: Optional[DataFrame] = None + estimator_input_cached = False # Fit each stage, appending the transformer to the list of transformers # If the stage is a parent of an estimator, transform the dataset. transformers: List[BaseTransformer] = [] @@ -404,6 +466,22 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": if stage in estimator_parent_stages: dataset = stage.transform(dataset) else: + # Opt-in: at the first estimator boundary, project to the columns + # still read downstream and persist that narrow frame once. + # Independent sibling estimators then fit against the cached narrow + # frame instead of re-scanning the wide input. The persist sits + # below each estimator's in-fit sample, so sampling is unchanged. + if cache_estimator_input and not estimator_input_cached: + estimator_input_cached = True + live_columns = self.collect_required_input_columns( + expanded_pipeline_stages[index:] + ) + keep_columns = [c for c in dataset.columns if c in live_columns] + if keep_columns and len(keep_columns) < len(dataset.columns): + estimator_input_cache = dataset.select(*keep_columns).persist( + StorageLevel.MEMORY_AND_DISK + ) + dataset = estimator_input_cache # Truncate accumulated lineage before the fit action to bound plan # depth. eager=True materialises now. if ( @@ -426,6 +504,8 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": dataset = model.transform(dataset) if cached_dataset is not None: cached_dataset.unpersist() + if estimator_input_cache is not None: + estimator_input_cache.unpersist() return KamaeSparkPipelineModel(transformers) def copy(self, extra: Optional["ParamMap"] = None) -> "KamaeSparkPipeline": diff --git a/tests/kamae/spark/pipeline/test_pipeline.py b/tests/kamae/spark/pipeline/test_pipeline.py index 997ba854..dd5b040a 100644 --- a/tests/kamae/spark/pipeline/test_pipeline.py +++ b/tests/kamae/spark/pipeline/test_pipeline.py @@ -834,6 +834,86 @@ def test_spark_pipeline_prune_is_transparent_to_fit( assert baseline_out.exceptAll(with_extra_out).isEmpty() assert with_extra_out.exceptAll(baseline_out).isEmpty() + def test_spark_pipeline_cache_estimator_input_is_transparent_to_fit( + self, spark_session + ): + """ + cacheEstimatorInput projects to still-needed columns and persists that + narrow frame once; independent sibling estimators must fit to byte-identical + params whether it is on or off, with the genuinely-unused column dropped. + """ + df = spark_session.createDataFrame( + [ + (1.0, 2.0, 3.0, 99.0), + (2.0, 4.0, 6.0, 99.0), + (3.0, 6.0, 9.0, 99.0), + (4.0, 8.0, 12.0, 99.0), + ], + ["x1", "x2", "x3", "junk"], + ) + + def build_pipeline(cache): + return KamaeSparkPipeline( + stages=[ + StandardScaleEstimator(inputCol="x1", outputCol="x1_scaled"), + StandardScaleEstimator(inputCol="x2", outputCol="x2_scaled"), + StandardScaleEstimator(inputCol="x3", outputCol="x3_scaled"), + ], + cacheEstimatorInput=cache, + ) + + cached_model = build_pipeline(cache=True).fit(df) + baseline_model = build_pipeline(cache=False).fit(df) + + for cached_scaler, baseline_scaler in zip( + cached_model.stages, baseline_model.stages + ): + assert cached_scaler.getMean() == baseline_scaler.getMean() + assert cached_scaler.getStddev() == baseline_scaler.getStddev() + + def test_spark_pipeline_cache_estimator_input_is_opt_in(self, spark_session): + """ + The narrow-cache projection (which computes the live keep-set via + collect_required_input_columns) must only run when cacheEstimatorInput is + True. Pruning is left off so the collector is not called for that reason. + """ + df = spark_session.createDataFrame( + [(1.0, 2.0, 99.0), (2.0, 4.0, 99.0), (3.0, 6.0, 99.0)], + ["x1", "x2", "junk"], + ) + stages = [ + StandardScaleEstimator(inputCol="x1", outputCol="x1_scaled"), + StandardScaleEstimator(inputCol="x2", outputCol="x2_scaled"), + ] + + with patch.object( + KamaeSparkPipeline, + "collect_required_input_columns", + wraps=KamaeSparkPipeline.collect_required_input_columns, + ) as mock_collect: + KamaeSparkPipeline(stages=stages).fit(df) + assert mock_collect.call_count == 0 + + mock_collect.reset_mock() + KamaeSparkPipeline(stages=stages, cacheEstimatorInput=True).fit(df) + assert mock_collect.call_count == 1 + + def test_spark_pipeline_cache_estimator_input_mutually_exclusive( + self, valid_stages_1, example_dataframe + ): + """ + cacheEstimatorInput and cacheIntermediateData are competing strategies; + enabling both warns, prefers cacheEstimatorInput, and still fits. + """ + pipeline = KamaeSparkPipeline( + stages=valid_stages_1, + cacheIntermediateData=True, + cacheEstimatorInput=True, + ) + with pytest.warns(UserWarning, match="takes precedence"): + pipeline_model = pipeline.fit(example_dataframe) + assert pipeline_model is not None + @pytest.mark.parametrize( "stages, input_col, original_dtype", [ From 2e0a0cd06d7c1b281ebb74f5673262e357ce886e Mon Sep 17 00:00:00 2001 From: cworthington Date: Thu, 20 Aug 2026 17:53:21 +0100 Subject: [PATCH 33/36] feat: add opt-in fitSampleFraction fit optimisation to KamaeSparkPipeline 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 --- src/kamae/spark/pipeline/pipeline.py | 254 +++++++++++++++++--- tests/kamae/spark/pipeline/test_pipeline.py | 128 ++++++++++ 2 files changed, 353 insertions(+), 29 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index 57b5aee4..e23a34fb 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -13,7 +13,7 @@ # limitations under the License. import warnings -from typing import TYPE_CHECKING, List, Optional, Set, Type +from typing import TYPE_CHECKING, Any, List, Optional, Set, Tuple, Type import networkx as nx from pyspark import keyword_only @@ -41,14 +41,16 @@ class KamaeSparkPipeline(Pipeline): together BaseTransformers. It maintains the same functionality as pyspark.ml.Pipeline e.g. serialisation. - Four opt-in fit optimisations are available, all defaulting off (fit behaviour + Five opt-in fit optimisations are available, all defaulting off (fit behaviour unchanged): `checkpointInterval` reliably checkpoints every N stages to bound logical-plan depth (requires a checkpoint dir); `cacheIntermediateData` persists the working DataFrame at each estimator-fit boundary to avoid re-scanning the upstream lineage; `pruneInputColumns` drops input columns no stage consumes; `cacheEstimatorInput` projects to the columns still read downstream at the first estimator boundary and persists that narrow frame once, so independent sibling - estimators reuse it instead of re-scanning the wide input. + estimators reuse it instead of re-scanning the wide input; `fitSampleFraction` + draws a single persisted sample of the input up-front and fits every estimator + from it (see its param docstring for the correctness caveat). """ checkpointInterval = Param( @@ -86,6 +88,29 @@ class KamaeSparkPipeline(Pipeline): typeConverter=TypeConverters.toBoolean, ) + fitSampleFraction = Param( + Params._dummy(), + "fitSampleFraction", + "If set to a float in (0, 1], draw a single sample of the input up-front, " + "persist (MEMORY_AND_DISK) and materialise it once, and fit every estimator " + "from that shared sample (each estimator's own sampleFraction is ignored for " + "the fit). Avoids re-scanning or persisting the wide source once per " + "estimator. ONLY correct when every estimator computes sample-robust " + "statistics (mean/std/quantiles, e.g. ConditionalStandardScale); vocabulary " + "builders (StringIndexer/OneHot), min/max scalers and distinct counts need " + "exact/global statistics and will be inaccurate on a sample. Disables " + "cacheIntermediateData and cacheEstimatorInput. None (default) disables it.", + typeConverter=TypeConverters.toFloat, + ) + + fitSampleSeed = Param( + Params._dummy(), + "fitSampleSeed", + "Optional integer seed passed to the fitSampleFraction sample for " + "reproducibility. None (default) leaves the sample unseeded.", + typeConverter=TypeConverters.toInt, + ) + @keyword_only def __init__( self, @@ -95,6 +120,8 @@ def __init__( cacheIntermediateData: bool = False, pruneInputColumns: bool = False, cacheEstimatorInput: bool = False, + fitSampleFraction: Optional[float] = None, + fitSampleSeed: Optional[int] = None, ) -> None: """ Initialises the KamaeSparkPipeline object. @@ -110,6 +137,11 @@ def __init__( :param cacheEstimatorInput: If True, project to the columns still read downstream at the first estimator boundary and persist that narrow frame once for reuse by subsequent estimators. False (default) disables it. + :param fitSampleFraction: If set to a float in (0, 1], fit every estimator + from a single up-front persisted sample of the input. Only correct when all + estimators compute sample-robust statistics. None (default) disables it. + :param fitSampleSeed: Optional integer seed for the fitSampleFraction + sample. None (default) leaves it unseeded. :returns: None - class instantiated. """ kwargs = self._input_kwargs @@ -119,6 +151,8 @@ def __init__( cacheIntermediateData=False, pruneInputColumns=False, cacheEstimatorInput=False, + fitSampleFraction=None, + fitSampleSeed=None, ) self.setParams(**kwargs) @@ -216,6 +250,46 @@ def getCacheEstimatorInput(self) -> bool: """ return self.getOrDefault(self.cacheEstimatorInput) + def setFitSampleFraction(self, value: Optional[float]) -> "KamaeSparkPipeline": + """ + Sets the `fitSampleFraction` parameter. + + :param value: Fraction in (0, 1] of the input to sample once up-front and + fit every estimator from. None disables it. + :returns: KamaeSparkPipeline object with fitSampleFraction set. + :raises ValueError: If value is not None and not in the range (0, 1]. + """ + if value is not None and not 0.0 < value <= 1.0: + raise ValueError( + f"fitSampleFraction must be in the range (0, 1] or None, got {value}." + ) + return self._set(fitSampleFraction=value) + + def getFitSampleFraction(self) -> Optional[float]: + """ + Gets the value of the `fitSampleFraction` parameter. + + :returns: The fitSampleFraction value. + """ + return self.getOrDefault(self.fitSampleFraction) + + def setFitSampleSeed(self, value: Optional[int]) -> "KamaeSparkPipeline": + """ + Sets the `fitSampleSeed` parameter. + + :param value: Integer seed for the fitSampleFraction sample, or None. + :returns: KamaeSparkPipeline object with fitSampleSeed set. + """ + return self._set(fitSampleSeed=value) + + def getFitSampleSeed(self) -> Optional[int]: + """ + Gets the value of the `fitSampleSeed` parameter. + + :returns: The fitSampleSeed value. + """ + return self.getOrDefault(self.fitSampleSeed) + @keyword_only def setParams( self, @@ -225,6 +299,8 @@ def setParams( cacheIntermediateData: bool = False, pruneInputColumns: bool = False, cacheEstimatorInput: bool = False, + fitSampleFraction: Optional[float] = None, + fitSampleSeed: Optional[int] = None, ) -> "KamaeSparkPipeline": """ Sets the keyword arguments of the pipeline. @@ -242,6 +318,11 @@ def setParams( :param cacheEstimatorInput: If True, project to the columns still read downstream at the first estimator boundary and persist that narrow frame once for reuse by subsequent estimators. False (default) disables it. + :param fitSampleFraction: If set to a float in (0, 1], fit every estimator + from a single up-front persisted sample of the input. None (default) + disables it. + :param fitSampleSeed: Optional integer seed for the fitSampleFraction + sample. None (default) leaves it unseeded. :returns: KamaeSparkPipeline object with params set. """ for param_name, param_value in self._input_kwargs.items(): @@ -402,8 +483,9 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": then constructs a KamaeSparkPipelineModel uses the stages from the fit pipeline. Optionally applies the opt-in fit optimisations (`pruneInputColumns`, - `checkpointInterval`, `cacheIntermediateData`, `cacheEstimatorInput`); see the - class docstring. All preserve data exactly, so fitted results match the + `checkpointInterval`, `cacheIntermediateData`, `cacheEstimatorInput`, + `fitSampleFraction`); see the class docstring. With the exception of + `fitSampleFraction`, all preserve data exactly, so fitted results match the defaults-off behaviour. If both `cacheIntermediateData` and `cacheEstimatorInput` are enabled, @@ -411,6 +493,13 @@ class docstring. All preserve data exactly, so fitted results match the `cacheIntermediateData` is ignored, since the narrow frame is a strictly smaller cache and the intermediate cache would evict it. + If `fitSampleFraction` is set, the input is sampled once, persisted and + materialised, and every estimator is fit from that shared sample with its + own `sampleFraction` temporarily disabled; `cacheIntermediateData` and + `cacheEstimatorInput` are disabled (with a warning) as they persist frames + this option avoids. Only correct for sample-robust estimators (a runtime + warning is emitted); see the `fitSampleFraction` param docstring. + :param dataset: PySpark DataFrame to fit the pipeline to. :returns: KamaeSparkPipelineModel object. :raises ValueError: If checkpointing is enabled but no checkpoint directory @@ -451,26 +540,135 @@ class docstring. All preserve data exactly, so fitted results match the stacklevel=2, ) cache_enabled = False + + # Opt-in: fit every estimator from one shared sample drawn up-front, instead + # of re-scanning or persisting the wide source once per estimator. + fit_sample_fraction = self.getFitSampleFraction() + sampled_dataset: Optional[DataFrame] = None + # (stage, param, was_explicitly_set, original_value) for restoration. + overridden_sample_fractions: List[Tuple[Any, Any, bool, Any]] = [] + if fit_sample_fraction is not None: + warnings.warn( + "fitSampleFraction fits every estimator on one shared sample of the " + "input. This is only correct when all estimators compute " + "sample-robust statistics (mean/std/quantiles, e.g. " + "ConditionalStandardScale). Vocabulary builders " + "(StringIndexer/OneHot), min/max scalers and distinct counts require " + "exact/global statistics and will be inaccurate on a sample. It also " + "replaces the independent per-estimator samples with one shared " + "sample.", + stacklevel=2, + ) + # fitSampleFraction persists a tiny sample instead of the wide frame, so + # the caching strategies it supersedes are turned off. + if cache_enabled or cache_estimator_input: + warnings.warn( + "fitSampleFraction is incompatible with cacheIntermediateData " + "and cacheEstimatorInput (which persist frames this option " + "avoids); the caching options are disabled.", + stacklevel=2, + ) + cache_enabled = False + cache_estimator_input = False + sampled_dataset = dataset.sample( + fraction=fit_sample_fraction, seed=self.getFitSampleSeed() + ).persist(StorageLevel.MEMORY_AND_DISK) + # Force one materialisation so the sample is computed exactly once and + # every estimator fit reads the cached rows instead of re-scanning. + sampled_dataset.count() + dataset = sampled_dataset + # The shared sample is already drawn, so each estimator must not sample + # again (fraction-of-a-fraction would leave far too few rows). Disable + # each estimator's sampleFraction for this fit, restoring it in finally. + for stage in expanded_pipeline_stages: + if isinstance(stage, BaseEstimator) and stage.hasParam( + "sampleFraction" + ): + param = stage.getParam("sampleFraction") + was_set = stage.isSet(param) + original = stage.getOrDefault(param) if was_set else None + overridden_sample_fractions.append( + (stage, param, was_set, original) + ) + stage.set(param, None) + + # Fit each stage, appending the transformer to the list of transformers. + # If the stage is a parent of an estimator, transform the dataset. + transformers: List[BaseTransformer] = [] + try: + fitted_pipeline_model = self._run_fit_loop( + expanded_pipeline_stages=expanded_pipeline_stages, + dataset=dataset, + estimator_parent_stages=estimator_parent_stages, + transformers=transformers, + checkpoint_enabled=checkpoint_enabled, + checkpoint_interval=checkpoint_interval, + cache_enabled=cache_enabled, + cache_estimator_input=cache_estimator_input, + ) + finally: + # Restore each estimator's original sampleFraction and release the + # shared sample, whether or not the fit succeeded. + for stage, param, was_set, original in overridden_sample_fractions: + if was_set: + stage.set(param, original) + else: + stage.clear(param) + if sampled_dataset is not None: + sampled_dataset.unpersist() + return fitted_pipeline_model + + def _run_fit_loop( + self, + *, + expanded_pipeline_stages: List["KamaePipelineStage"], + dataset: DataFrame, + estimator_parent_stages: List["KamaePipelineStage"], + transformers: List[BaseTransformer], + checkpoint_enabled: bool, + checkpoint_interval: Optional[int], + cache_enabled: bool, + cache_estimator_input: bool, + ) -> "KamaeSparkPipelineModel": + """ + Runs the stage-by-stage fit loop, applying the checkpoint/cache optimisations. + + Extracted from `_fit` so the loop can run inside a try/finally that restores + estimator sampling and releases the shared sample when fitSampleFraction is + used. Behaviour is identical to the previous inline loop. + + `cacheIntermediateData` and `cacheEstimatorInput` are mutually exclusive (see + `_fit`), so a single `persisted_frame` handle tracks whichever frame is + persisted. It is kept separate from `dataset` because `dataset` is reassigned + by `model.transform(...)` between boundaries; the handle is what lets us + unpersist the actual persisted frame at the end. + + :param expanded_pipeline_stages: Flattened pipeline stages to fit. + :param dataset: DataFrame (possibly sampled) to fit the stages against. + :param estimator_parent_stages: Stages whose output an estimator consumes. + :param transformers: Accumulator list the fitted stages are appended to. + :param checkpoint_enabled: Whether reliable checkpointing is enabled. + :param checkpoint_interval: Stages between checkpoints (when enabled). + :param cache_enabled: Whether to persist at each estimator boundary. + :param cache_estimator_input: Whether to persist a narrow frame once. + :returns: KamaeSparkPipelineModel object. + """ last_checkpoint_index = 0 # The single persisted frame (if any), unpersisted once superseded or done. - cached_dataset: Optional[DataFrame] = None - # The narrow estimator-input frame (if any), persisted once and reused. - estimator_input_cache: Optional[DataFrame] = None + # Only one of the mutually-exclusive cache strategies ever populates it. + persisted_frame: Optional[DataFrame] = None estimator_input_cached = False - # Fit each stage, appending the transformer to the list of transformers - # If the stage is a parent of an estimator, transform the dataset. - transformers: List[BaseTransformer] = [] for index, stage in enumerate(expanded_pipeline_stages): if isinstance(stage, BaseTransformer): transformers.append(stage) if stage in estimator_parent_stages: dataset = stage.transform(dataset) else: - # Opt-in: at the first estimator boundary, project to the columns - # still read downstream and persist that narrow frame once. - # Independent sibling estimators then fit against the cached narrow - # frame instead of re-scanning the wide input. The persist sits - # below each estimator's in-fit sample, so sampling is unchanged. + # cacheEstimatorInput: at the first estimator boundary, project to + # the columns still read downstream and persist that narrow frame + # once. Independent sibling estimators then fit against it instead of + # re-scanning the wide input. The persist sits below each estimator's + # in-fit sample, so sampling is unchanged. if cache_estimator_input and not estimator_input_cached: estimator_input_cached = True live_columns = self.collect_required_input_columns( @@ -478,10 +676,10 @@ class docstring. All preserve data exactly, so fitted results match the ) keep_columns = [c for c in dataset.columns if c in live_columns] if keep_columns and len(keep_columns) < len(dataset.columns): - estimator_input_cache = dataset.select(*keep_columns).persist( + dataset = dataset.select(*keep_columns).persist( StorageLevel.MEMORY_AND_DISK ) - dataset = estimator_input_cache + persisted_frame = dataset # Truncate accumulated lineage before the fit action to bound plan # depth. eager=True materialises now. if ( @@ -490,22 +688,20 @@ class docstring. All preserve data exactly, so fitted results match the ): dataset = dataset.checkpoint(eager=True) last_checkpoint_index = index - # Persist so the fit action and downstream transforms reuse a - # materialised frame instead of re-scanning. One frame held at a time. + # cacheIntermediateData: persist so the fit action and downstream + # transforms reuse a materialised frame instead of re-scanning, + # releasing the previous frame first. One frame held at a time. if cache_enabled: - new_cached = dataset.persist(StorageLevel.MEMORY_AND_DISK) - if cached_dataset is not None: - cached_dataset.unpersist() - cached_dataset = new_cached - dataset = new_cached + if persisted_frame is not None: + persisted_frame.unpersist() + dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) + persisted_frame = dataset model = stage.fit(dataset) transformers.append(model) if stage in estimator_parent_stages: dataset = model.transform(dataset) - if cached_dataset is not None: - cached_dataset.unpersist() - if estimator_input_cache is not None: - estimator_input_cache.unpersist() + if persisted_frame is not None: + persisted_frame.unpersist() return KamaeSparkPipelineModel(transformers) def copy(self, extra: Optional["ParamMap"] = None) -> "KamaeSparkPipeline": diff --git a/tests/kamae/spark/pipeline/test_pipeline.py b/tests/kamae/spark/pipeline/test_pipeline.py index dd5b040a..8e983cfb 100644 --- a/tests/kamae/spark/pipeline/test_pipeline.py +++ b/tests/kamae/spark/pipeline/test_pipeline.py @@ -914,6 +914,134 @@ def test_spark_pipeline_cache_estimator_input_mutually_exclusive( pipeline_model = pipeline.fit(example_dataframe) assert pipeline_model is not None + def test_spark_pipeline_fit_sample_fraction_none_is_unchanged( + self, valid_stages_1, example_dataframe + ): + """ + With fitSampleFraction=None (the default), fit is unchanged: a smoke fit + succeeds and the resulting model transforms the full dataset. + """ + model = KamaeSparkPipeline(stages=valid_stages_1, fitSampleFraction=None).fit( + example_dataframe + ) + assert isinstance(model, KamaeSparkPipelineModel) + # The model still applies to the full dataset at transform time. + assert model.transform(example_dataframe).count() == example_dataframe.count() + + def test_spark_pipeline_fit_sample_fraction_matches_full_within_tolerance( + self, spark_session + ): + """ + fitSampleFraction fits sample-robust estimators (ConditionalStandardScale) + from a single shared sample; on enough seeded data the fitted mean/stddev + must stay within a loose statistical tolerance of a full-data fit. + """ + from pyspark.sql import functions as F + + df = ( + spark_session.range(0, 40000) + .withColumn("x1", F.randn(seed=42)) + .withColumn("x2", 5.0 + 2.0 * F.randn(seed=7)) + .select("x1", "x2") + ).persist() + df.count() + + def build(fraction): + return KamaeSparkPipeline( + stages=[ + ConditionalStandardScaleEstimator( + inputCol="x1", outputCol="x1_scaled" + ), + ConditionalStandardScaleEstimator( + inputCol="x2", outputCol="x2_scaled" + ), + ], + fitSampleFraction=fraction, + fitSampleSeed=13, + ) + + full_model = build(None).fit(df) + with pytest.warns(UserWarning): + sampled_model = build(0.2).fit(df) + + for full_stage, sampled_stage in zip(full_model.stages, sampled_model.stages): + assert abs(full_stage.getMean()[0] - sampled_stage.getMean()[0]) < 0.15 + assert abs(full_stage.getStddev()[0] - sampled_stage.getStddev()[0]) < 0.15 + df.unpersist() + + def test_spark_pipeline_fit_sample_fraction_scans_source_once(self, spark_session): + """ + fitSampleFraction materialises one shared sample, so the source is scanned + exactly once regardless of how many estimators fit from it - unlike the + default, where each independent estimator rescans the source. + """ + from pyspark.sql import functions as F + from pyspark.sql.types import DoubleType as SparkDoubleType + + n_rows = 500 + source = ( + spark_session.range(0, n_rows) + .select(F.col("id").cast("double").alias("raw")) + .persist() + ) + source.count() + + def counting_column(accumulator): + def _count(value): + accumulator.add(1) + return value + + udf = F.udf(_count, SparkDoubleType()).asNondeterministic() + return source.withColumn("x", udf(F.col("raw"))).drop("raw") + + def estimators(): + return [ + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_a"), + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_b"), + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_c"), + ] + + sampled_accum = spark_session.sparkContext.accumulator(0) + with pytest.warns(UserWarning): + KamaeSparkPipeline( + stages=estimators(), fitSampleFraction=1.0, fitSampleSeed=1 + ).fit(counting_column(sampled_accum)) + # One shared-sample materialisation => exactly one pass over the source. + assert sampled_accum.value == n_rows + + baseline_accum = spark_session.sparkContext.accumulator(0) + KamaeSparkPipeline(stages=estimators()).fit(counting_column(baseline_accum)) + # Default: each of the three estimators rescans the source. + assert baseline_accum.value > n_rows + + source.unpersist() + + def test_spark_pipeline_fit_sample_fraction_disables_caching(self, spark_session): + """ + fitSampleFraction persists a tiny sample instead of the wide frame, so it is + incompatible with cacheEstimatorInput: enabling both must warn and disable + the cache (its narrow-projection path never runs). + """ + df = spark_session.createDataFrame([(1.0,), (2.0,), (3.0,), (4.0,)], ["x"]) + pipeline = KamaeSparkPipeline( + stages=[ + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_scaled") + ], + cacheEstimatorInput=True, + fitSampleFraction=1.0, + fitSampleSeed=1, + ) + with patch.object( + KamaeSparkPipeline, + "collect_required_input_columns", + wraps=KamaeSparkPipeline.collect_required_input_columns, + ) as mock_collect: + with pytest.warns(UserWarning, match="incompatible"): + model = pipeline.fit(df) + # cacheEstimatorInput disabled => its keep-set collector is never called. + assert mock_collect.call_count == 0 + assert model is not None + @pytest.mark.parametrize( "stages, input_col, original_dtype", [ From ece629dd295591426ca1f27fd11dfdbd6f2553e4 Mon Sep 17 00:00:00 2001 From: cworthington Date: Fri, 21 Aug 2026 15:32:40 +0100 Subject: [PATCH 34/36] feat: make fitSampleFraction per-estimator opt-in via useFitSample 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 --- .../estimators/conditional_standard_scale.py | 4 + src/kamae/spark/estimators/impute.py | 4 + src/kamae/spark/estimators/min_max_scale.py | 5 +- .../single_feature_array_standard_scale.py | 5 +- src/kamae/spark/estimators/standard_scale.py | 5 +- src/kamae/spark/params/base.py | 27 ++ src/kamae/spark/pipeline/pipeline.py | 240 ++++++++++++------ tests/kamae/spark/pipeline/test_pipeline.py | 145 ++++++++++- 8 files changed, 338 insertions(+), 97 deletions(-) diff --git a/src/kamae/spark/estimators/conditional_standard_scale.py b/src/kamae/spark/estimators/conditional_standard_scale.py index 30f6ae65..aa0e6098 100644 --- a/src/kamae/spark/estimators/conditional_standard_scale.py +++ b/src/kamae/spark/estimators/conditional_standard_scale.py @@ -259,6 +259,7 @@ def __init__( epsilon: float = 0, nanFillValue: Optional[float] = None, sampleFraction: Optional[float] = None, + useFitSample: bool = False, ) -> None: """ Initializes a ConditionalStandardScaleEstimator estimator. @@ -288,6 +289,8 @@ def __init__( to use it if epsilon filters out all the values. Defaults to None. :param sampleFraction: Fraction of data to sample for statistics estimation (exclusive 0.0-1.0). Default None (no sampling). + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. :returns: None - class instantiated. """ super().__init__() @@ -301,6 +304,7 @@ def __init__( epsilon=0, nanFillValue=None, sampleFraction=None, + useFitSample=False, ) kwargs = self._input_kwargs self.setParams(**kwargs) diff --git a/src/kamae/spark/estimators/impute.py b/src/kamae/spark/estimators/impute.py index 56a4e83e..8498a57e 100644 --- a/src/kamae/spark/estimators/impute.py +++ b/src/kamae/spark/estimators/impute.py @@ -66,6 +66,7 @@ def __init__( maskValue: Optional[Union[float, int, str]] = None, imputeMethod: Optional[str] = None, sampleFraction: Optional[float] = None, + useFitSample: bool = False, ) -> None: """ Initializes a ImputeEstimator estimator. @@ -86,12 +87,15 @@ def __init__( Valid values are "mean" or "median". :param sampleFraction: Fraction of data to sample for statistics estimation (exclusive 0.0-1.0). Default None (no sampling). + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. :returns: None - class instantiated. """ super().__init__() self._setDefault( imputeMethod="mean", sampleFraction=None, + useFitSample=False, ) self.valid_impute_methods = ["mean", "median"] kwargs = self._input_kwargs diff --git a/src/kamae/spark/estimators/min_max_scale.py b/src/kamae/spark/estimators/min_max_scale.py index 36a6430f..31c878f8 100644 --- a/src/kamae/spark/estimators/min_max_scale.py +++ b/src/kamae/spark/estimators/min_max_scale.py @@ -65,6 +65,7 @@ def __init__( layerName: Optional[str] = None, maskValue: Optional[float] = None, sampleFraction: Optional[float] = None, + useFitSample: bool = False, ) -> None: """ Initializes a MinMaxScaleEstimator estimator. @@ -82,10 +83,12 @@ def __init__( during the computation of the min and max values. :param sampleFraction: Fraction of data to sample for statistics estimation (exclusive 0.0-1.0). Default None (no sampling). + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. :returns: None - class instantiated. """ super().__init__() - self._setDefault(maskValue=None, sampleFraction=None) + self._setDefault(maskValue=None, sampleFraction=None, useFitSample=False) kwargs = self._input_kwargs self.setParams(**kwargs) diff --git a/src/kamae/spark/estimators/single_feature_array_standard_scale.py b/src/kamae/spark/estimators/single_feature_array_standard_scale.py index a197ae4c..121a0fbe 100644 --- a/src/kamae/spark/estimators/single_feature_array_standard_scale.py +++ b/src/kamae/spark/estimators/single_feature_array_standard_scale.py @@ -62,6 +62,7 @@ def __init__( layerName: Optional[str] = None, maskValue: Optional[float] = None, sampleFraction: Optional[float] = None, + useFitSample: bool = False, ) -> None: """ Initializes a SingleFeatureArrayStandardScaleEstimator estimator. @@ -77,10 +78,12 @@ def __init__( in the keras model. If not set, we use the uid of the Spark transformer. :param sampleFraction: Fraction of data to sample for statistics estimation (exclusive 0.0-1.0). Default None (no sampling). + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. :returns: None - class instantiated. """ super().__init__() - self._setDefault(maskValue=None, sampleFraction=None) + self._setDefault(maskValue=None, sampleFraction=None, useFitSample=False) kwargs = self._input_kwargs self.setParams(**kwargs) diff --git a/src/kamae/spark/estimators/standard_scale.py b/src/kamae/spark/estimators/standard_scale.py index 379dce37..0d0b44aa 100644 --- a/src/kamae/spark/estimators/standard_scale.py +++ b/src/kamae/spark/estimators/standard_scale.py @@ -66,6 +66,7 @@ def __init__( layerName: Optional[str] = None, maskValue: Optional[float] = None, sampleFraction: Optional[float] = None, + useFitSample: bool = False, ) -> None: """ Initializes a StandardScaleEstimator estimator. @@ -81,10 +82,12 @@ def __init__( in the keras model. If not set, we use the uid of the Spark transformer. :param sampleFraction: Fraction of data to sample for statistics estimation (exclusive 0.0-1.0). Default None (no sampling). + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. :returns: None - class instantiated. """ super().__init__() - self._setDefault(maskValue=None, sampleFraction=None) + self._setDefault(maskValue=None, sampleFraction=None, useFitSample=False) kwargs = self._input_kwargs self.setParams(**kwargs) diff --git a/src/kamae/spark/params/base.py b/src/kamae/spark/params/base.py index 071305ad..a96967e2 100644 --- a/src/kamae/spark/params/base.py +++ b/src/kamae/spark/params/base.py @@ -144,6 +144,15 @@ class SampleFractionParams(Params): "Default None (no sampling).", ) + useFitSample = Param( + Params._dummy(), + "useFitSample", + "If True, and the enclosing KamaeSparkPipeline has fitSampleFraction set, fit " + "this estimator on the pipeline's shared sample instead of the full input. " + "Provides opt-in per estimator without a second fraction value. Default False.", + typeConverter=TypeConverters.toBoolean, + ) + def setSampleFraction(self, value: float) -> "SampleFractionParams": """ Sets the parameter sampleFraction to the given float value. @@ -167,6 +176,24 @@ def getSampleFraction(self) -> Optional[float]: """ return self.getOrDefault(self.sampleFraction) + def setUseFitSample(self, value: bool) -> "SampleFractionParams": + """ + Sets the parameter useFitSample to the given boolean value. + + :param value: Whether to fit this estimator on the pipeline's shared sample + when fitSampleFraction is set. + :returns: Instance of class mixed in. + """ + return self._set(useFitSample=value) + + def getUseFitSample(self) -> bool: + """ + Gets the value of the useFitSample parameter. + + :returns: Whether this estimator opts in to the pipeline's shared sample. + """ + return self.getOrDefault(self.useFitSample) + class SingleInputParams(HasInputCol): """ diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index e23a34fb..c4a97eba 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -49,8 +49,9 @@ class KamaeSparkPipeline(Pipeline): `cacheEstimatorInput` projects to the columns still read downstream at the first estimator boundary and persists that narrow frame once, so independent sibling estimators reuse it instead of re-scanning the wide input; `fitSampleFraction` - draws a single persisted sample of the input up-front and fits every estimator - from it (see its param docstring for the correctness caveat). + draws a single persisted sample of the input up-front and fits the estimators + that opt in (those with `useFitSample=True`) from that shared sample, while + estimators without it keep fitting on the full input. """ checkpointInterval = Param( @@ -92,14 +93,15 @@ class KamaeSparkPipeline(Pipeline): Params._dummy(), "fitSampleFraction", "If set to a float in (0, 1], draw a single sample of the input up-front, " - "persist (MEMORY_AND_DISK) and materialise it once, and fit every estimator " - "from that shared sample (each estimator's own sampleFraction is ignored for " - "the fit). Avoids re-scanning or persisting the wide source once per " - "estimator. ONLY correct when every estimator computes sample-robust " - "statistics (mean/std/quantiles, e.g. ConditionalStandardScale); vocabulary " - "builders (StringIndexer/OneHot), min/max scalers and distinct counts need " - "exact/global statistics and will be inaccurate on a sample. Disables " - "cacheIntermediateData and cacheEstimatorInput. None (default) disables it.", + "persist (MEMORY_AND_DISK) and materialise it once, and fit the estimators " + "that opt in from that shared sample. An estimator opts in by setting its " + "boolean useFitSample param to True. Estimators without useFitSample fit on " + "the full input, so leave it False for estimators needing exact/global " + "statistics (vocabulary builders like StringIndexer/OneHot, min/max scalers, " + "distinct counts) and set it True on sample-robust estimators (mean/std/" + "quantiles, e.g. ConditionalStandardScale). Avoids re-scanning or persisting " + "the wide source once per opted-in estimator. Disables cacheIntermediateData " + "and cacheEstimatorInput. None (default) disables it.", typeConverter=TypeConverters.toFloat, ) @@ -137,9 +139,10 @@ def __init__( :param cacheEstimatorInput: If True, project to the columns still read downstream at the first estimator boundary and persist that narrow frame once for reuse by subsequent estimators. False (default) disables it. - :param fitSampleFraction: If set to a float in (0, 1], fit every estimator - from a single up-front persisted sample of the input. Only correct when all - estimators compute sample-robust statistics. None (default) disables it. + :param fitSampleFraction: If set to a float in (0, 1], fit the estimators + that opt in (those with useFitSample=True) from a single up-front persisted + sample of the input; estimators without useFitSample fit on the full input. + None (default) disables it. :param fitSampleSeed: Optional integer seed for the fitSampleFraction sample. None (default) leaves it unseeded. :returns: None - class instantiated. @@ -494,11 +497,12 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": smaller cache and the intermediate cache would evict it. If `fitSampleFraction` is set, the input is sampled once, persisted and - materialised, and every estimator is fit from that shared sample with its - own `sampleFraction` temporarily disabled; `cacheIntermediateData` and - `cacheEstimatorInput` are disabled (with a warning) as they persist frames - this option avoids. Only correct for sample-robust estimators (a runtime - warning is emitted); see the `fitSampleFraction` param docstring. + materialised, and the estimators that opt in (those with `useFitSample=True`) + are fit from that shared sample with any `sampleFraction` they also set + temporarily disabled; every other estimator still fits on the full input. + `cacheIntermediateData` and `cacheEstimatorInput` are disabled (with a + warning) as they persist frames this option avoids; a warning also names the + opted-in estimators. See the `fitSampleFraction` param docstring. :param dataset: PySpark DataFrame to fit the pipeline to. :returns: KamaeSparkPipelineModel object. @@ -541,56 +545,98 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": ) cache_enabled = False - # Opt-in: fit every estimator from one shared sample drawn up-front, instead - # of re-scanning or persisting the wide source once per estimator. + # Opt-in: draw one shared sample up-front and fit only the estimators that + # opt in against it, instead of each re-scanning or persisting the wide + # source. An estimator opts in via its boolean `useFitSample` param. + # Estimators without it set still fit on the full input, so exact/global- + # statistic estimators (vocabulary builders, min/max) stay correct. fit_sample_fraction = self.getFitSampleFraction() sampled_dataset: Optional[DataFrame] = None - # (stage, param, was_explicitly_set, original_value) for restoration. - overridden_sample_fractions: List[Tuple[Any, Any, bool, Any]] = [] + # (stage, param, original_value) for restoration of overridden sampleFractions. + overridden_sample_fractions: List[Tuple[Any, Any, Any]] = [] + # Identities of the estimators that fit on the shared sample. + use_sample_stage_ids: Set[int] = set() if fit_sample_fraction is not None: - warnings.warn( - "fitSampleFraction fits every estimator on one shared sample of the " - "input. This is only correct when all estimators compute " - "sample-robust statistics (mean/std/quantiles, e.g. " - "ConditionalStandardScale). Vocabulary builders " - "(StringIndexer/OneHot), min/max scalers and distinct counts require " - "exact/global statistics and will be inaccurate on a sample. It also " - "replaces the independent per-estimator samples with one shared " - "sample.", - stacklevel=2, - ) - # fitSampleFraction persists a tiny sample instead of the wide frame, so - # the caching strategies it supersedes are turned off. - if cache_enabled or cache_estimator_input: + sampling_estimators = [ + stage + for stage in expanded_pipeline_stages + if isinstance(stage, BaseEstimator) + and stage.hasParam("useFitSample") + and stage.getUseFitSample() + ] + if not sampling_estimators: warnings.warn( - "fitSampleFraction is incompatible with cacheIntermediateData " - "and cacheEstimatorInput (which persist frames this option " - "avoids); the caching options are disabled.", + "fitSampleFraction is set but no estimator has useFitSample=True, " + "so nothing opts in to the shared sample and fitSampleFraction " + "has no effect. Set useFitSample=True on the estimators that " + "should fit on the sample.", stacklevel=2, ) - cache_enabled = False - cache_estimator_input = False - sampled_dataset = dataset.sample( - fraction=fit_sample_fraction, seed=self.getFitSampleSeed() - ).persist(StorageLevel.MEMORY_AND_DISK) - # Force one materialisation so the sample is computed exactly once and - # every estimator fit reads the cached rows instead of re-scanning. - sampled_dataset.count() - dataset = sampled_dataset - # The shared sample is already drawn, so each estimator must not sample - # again (fraction-of-a-fraction would leave far too few rows). Disable - # each estimator's sampleFraction for this fit, restoring it in finally. - for stage in expanded_pipeline_stages: - if isinstance(stage, BaseEstimator) and stage.hasParam( - "sampleFraction" - ): + else: + warnings.warn( + "fitSampleFraction fits " + f"{sorted({type(s).__name__ for s in sampling_estimators})} on " + "one shared pipeline sample (useFitSample=True). Estimators " + "without useFitSample set fit on the full input.", + stacklevel=2, + ) + # fitSampleFraction persists a tiny sample instead of the wide frame, + # so the caching strategies it supersedes are turned off. + if cache_enabled or cache_estimator_input: + warnings.warn( + "fitSampleFraction is incompatible with cacheIntermediateData " + "and cacheEstimatorInput (which persist frames this option " + "avoids); the caching options are disabled.", + stacklevel=2, + ) + cache_enabled = False + cache_estimator_input = False + sampled_dataset = dataset.sample( + fraction=fit_sample_fraction, seed=self.getFitSampleSeed() + ).persist(StorageLevel.MEMORY_AND_DISK) + # Force one materialisation so the sample is computed exactly once and + # every opted-in estimator reads the cached rows instead of rescanning. + sampled_dataset.count() + # The shared sample is already drawn, so opted-in estimators must not + # sample again (fraction-of-a-fraction would leave far too few rows). + # Disable any sampleFraction they also set, restoring it in finally. + for stage in sampling_estimators: + use_sample_stage_ids.add(id(stage)) param = stage.getParam("sampleFraction") - was_set = stage.isSet(param) - original = stage.getOrDefault(param) if was_set else None - overridden_sample_fractions.append( - (stage, param, was_set, original) + if stage.isSet(param): + overridden_sample_fractions.append( + (stage, param, stage.getOrDefault(param)) + ) + stage.set(param, None) + overridden_names = sorted( + {type(s).__name__ for s, _, _ in overridden_sample_fractions} + ) + if overridden_names: + warnings.warn( + f"{overridden_names} set both useFitSample=True and their own " + "sampleFraction. The shared pipeline sample wins, so their " + "sampleFraction is ignored for this fit. Unset one to silence " + "this warning.", + stacklevel=2, ) - stage.set(param, None) + else: + opted_in_without_sample = sorted( + { + type(stage).__name__ + for stage in expanded_pipeline_stages + if isinstance(stage, BaseEstimator) + and stage.hasParam("useFitSample") + and stage.getUseFitSample() + } + ) + if opted_in_without_sample: + warnings.warn( + f"{opted_in_without_sample} set useFitSample=True but the pipeline " + "has no fitSampleFraction, so there is no shared sample to fit on " + "and useFitSample has no effect. Set fitSampleFraction on the " + "pipeline to enable sampled fitting.", + stacklevel=2, + ) # Fit each stage, appending the transformer to the list of transformers. # If the stage is a parent of an estimator, transform the dataset. @@ -599,6 +645,8 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": fitted_pipeline_model = self._run_fit_loop( expanded_pipeline_stages=expanded_pipeline_stages, dataset=dataset, + sampled_dataset=sampled_dataset, + use_sample_stage_ids=use_sample_stage_ids, estimator_parent_stages=estimator_parent_stages, transformers=transformers, checkpoint_enabled=checkpoint_enabled, @@ -607,13 +655,10 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": cache_estimator_input=cache_estimator_input, ) finally: - # Restore each estimator's original sampleFraction and release the - # shared sample, whether or not the fit succeeded. - for stage, param, was_set, original in overridden_sample_fractions: - if was_set: - stage.set(param, original) - else: - stage.clear(param) + # Restore each opted-in estimator's original sampleFraction and release + # the shared sample, whether or not the fit succeeded. + for stage, param, original in overridden_sample_fractions: + stage.set(param, original) if sampled_dataset is not None: sampled_dataset.unpersist() return fitted_pipeline_model @@ -623,6 +668,8 @@ def _run_fit_loop( *, expanded_pipeline_stages: List["KamaePipelineStage"], dataset: DataFrame, + sampled_dataset: Optional[DataFrame], + use_sample_stage_ids: Set[int], estimator_parent_stages: List["KamaePipelineStage"], transformers: List[BaseTransformer], checkpoint_enabled: bool, @@ -635,16 +682,27 @@ def _run_fit_loop( Extracted from `_fit` so the loop can run inside a try/finally that restores estimator sampling and releases the shared sample when fitSampleFraction is - used. Behaviour is identical to the previous inline loop. + used. With fitSampleFraction off (`sampled_dataset is None`) behaviour is + identical to the previous inline loop. + + When fitSampleFraction is active, two lineages are carried in parallel: the + full-data lineage and a shared-sample lineage. Every transformer that feeds an + estimator is applied to both, so an estimator can fit on whichever it opted + in to - estimators in `use_sample_stage_ids` fit on the sample, all others on + the full input. The caching options are mutually exclusive with + fitSampleFraction (disabled in `_fit`), so they only ever act on the full + lineage when no sample is present. `cacheIntermediateData` and `cacheEstimatorInput` are mutually exclusive (see `_fit`), so a single `persisted_frame` handle tracks whichever frame is - persisted. It is kept separate from `dataset` because `dataset` is reassigned - by `model.transform(...)` between boundaries; the handle is what lets us - unpersist the actual persisted frame at the end. + persisted. It is kept separate from the lineage variable because that variable + is reassigned by `model.transform(...)` between boundaries; the handle is what + lets us unpersist the actual persisted frame at the end. :param expanded_pipeline_stages: Flattened pipeline stages to fit. - :param dataset: DataFrame (possibly sampled) to fit the stages against. + :param dataset: Full-data DataFrame to fit non-sampled stages against. + :param sampled_dataset: Shared sample (or None when fitSampleFraction is off). + :param use_sample_stage_ids: `id()`s of estimators that fit on the sample. :param estimator_parent_stages: Stages whose output an estimator consumes. :param transformers: Accumulator list the fitted stages are appended to. :param checkpoint_enabled: Whether reliable checkpointing is enabled. @@ -658,11 +716,18 @@ def _run_fit_loop( # Only one of the mutually-exclusive cache strategies ever populates it. persisted_frame: Optional[DataFrame] = None estimator_input_cached = False + # Full-data lineage and, when fitSampleFraction is active, a parallel + # shared-sample lineage kept in lock-step through the estimator-feeding + # transforms. + full_dataset = dataset + sample_dataset = sampled_dataset for index, stage in enumerate(expanded_pipeline_stages): if isinstance(stage, BaseTransformer): transformers.append(stage) if stage in estimator_parent_stages: - dataset = stage.transform(dataset) + full_dataset = stage.transform(full_dataset) + if sample_dataset is not None: + sample_dataset = stage.transform(sample_dataset) else: # cacheEstimatorInput: at the first estimator boundary, project to # the columns still read downstream and persist that narrow frame @@ -674,19 +739,23 @@ def _run_fit_loop( live_columns = self.collect_required_input_columns( expanded_pipeline_stages[index:] ) - keep_columns = [c for c in dataset.columns if c in live_columns] - if keep_columns and len(keep_columns) < len(dataset.columns): - dataset = dataset.select(*keep_columns).persist( + keep_columns = [ + c for c in full_dataset.columns if c in live_columns + ] + if keep_columns and len(keep_columns) < len(full_dataset.columns): + full_dataset = full_dataset.select(*keep_columns).persist( StorageLevel.MEMORY_AND_DISK ) - persisted_frame = dataset + persisted_frame = full_dataset # Truncate accumulated lineage before the fit action to bound plan # depth. eager=True materialises now. if ( checkpoint_enabled and index - last_checkpoint_index >= checkpoint_interval ): - dataset = dataset.checkpoint(eager=True) + full_dataset = full_dataset.checkpoint(eager=True) + if sample_dataset is not None: + sample_dataset = sample_dataset.checkpoint(eager=True) last_checkpoint_index = index # cacheIntermediateData: persist so the fit action and downstream # transforms reuse a materialised frame instead of re-scanning, @@ -694,12 +763,19 @@ def _run_fit_loop( if cache_enabled: if persisted_frame is not None: persisted_frame.unpersist() - dataset = dataset.persist(StorageLevel.MEMORY_AND_DISK) - persisted_frame = dataset - model = stage.fit(dataset) + full_dataset = full_dataset.persist(StorageLevel.MEMORY_AND_DISK) + persisted_frame = full_dataset + fit_dataset = ( + sample_dataset + if sample_dataset is not None and id(stage) in use_sample_stage_ids + else full_dataset + ) + model = stage.fit(fit_dataset) transformers.append(model) if stage in estimator_parent_stages: - dataset = model.transform(dataset) + full_dataset = model.transform(full_dataset) + if sample_dataset is not None: + sample_dataset = model.transform(sample_dataset) if persisted_frame is not None: persisted_frame.unpersist() return KamaeSparkPipelineModel(transformers) diff --git a/tests/kamae/spark/pipeline/test_pipeline.py b/tests/kamae/spark/pipeline/test_pipeline.py index 8e983cfb..5d24ed41 100644 --- a/tests/kamae/spark/pipeline/test_pipeline.py +++ b/tests/kamae/spark/pipeline/test_pipeline.py @@ -932,9 +932,10 @@ def test_spark_pipeline_fit_sample_fraction_matches_full_within_tolerance( self, spark_session ): """ - fitSampleFraction fits sample-robust estimators (ConditionalStandardScale) - from a single shared sample; on enough seeded data the fitted mean/stddev - must stay within a loose statistical tolerance of a full-data fit. + fitSampleFraction fits the opted-in sample-robust estimators + (ConditionalStandardScale, with useFitSample=True) from a single + shared sample; on enough seeded data the fitted mean/stddev must stay within + a loose statistical tolerance of a full-data fit. """ from pyspark.sql import functions as F @@ -947,13 +948,16 @@ def test_spark_pipeline_fit_sample_fraction_matches_full_within_tolerance( df.count() def build(fraction): + # useFitSample=True opts each estimator in to the shared sample. The + # full-data reference leaves it off so it is a genuine full fit. + opt_in = {"useFitSample": True} if fraction is not None else {} return KamaeSparkPipeline( stages=[ ConditionalStandardScaleEstimator( - inputCol="x1", outputCol="x1_scaled" + inputCol="x1", outputCol="x1_scaled", **opt_in ), ConditionalStandardScaleEstimator( - inputCol="x2", outputCol="x2_scaled" + inputCol="x2", outputCol="x2_scaled", **opt_in ), ], fitSampleFraction=fraction, @@ -994,19 +998,22 @@ def _count(value): udf = F.udf(_count, SparkDoubleType()).asNondeterministic() return source.withColumn("x", udf(F.col("raw"))).drop("raw") - def estimators(): + def estimators(opt_in=False): + # useFitSample=True opts the estimator in to the shared sample. + kw = {"useFitSample": True} if opt_in else {} return [ - ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_a"), - ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_b"), - ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_c"), + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_a", **kw), + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_b", **kw), + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_c", **kw), ] sampled_accum = spark_session.sparkContext.accumulator(0) with pytest.warns(UserWarning): KamaeSparkPipeline( - stages=estimators(), fitSampleFraction=1.0, fitSampleSeed=1 + stages=estimators(opt_in=True), fitSampleFraction=1.0, fitSampleSeed=1 ).fit(counting_column(sampled_accum)) - # One shared-sample materialisation => exactly one pass over the source. + # One shared-sample materialisation, reused by all three opted-in estimators + # => exactly one pass over the source. assert sampled_accum.value == n_rows baseline_accum = spark_session.sparkContext.accumulator(0) @@ -1025,7 +1032,9 @@ def test_spark_pipeline_fit_sample_fraction_disables_caching(self, spark_session df = spark_session.createDataFrame([(1.0,), (2.0,), (3.0,), (4.0,)], ["x"]) pipeline = KamaeSparkPipeline( stages=[ - ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_scaled") + ConditionalStandardScaleEstimator( + inputCol="x", outputCol="x_scaled", useFitSample=True + ) ], cacheEstimatorInput=True, fitSampleFraction=1.0, @@ -1042,6 +1051,118 @@ def test_spark_pipeline_fit_sample_fraction_disables_caching(self, spark_session assert mock_collect.call_count == 0 assert model is not None + def test_spark_pipeline_fit_sample_fraction_non_opted_in_reads_full( + self, spark_session + ): + """ + Only estimators with useFitSample=True fit on the shared sample; + an estimator without it still fits on the full input. With fraction=1.0 the + sample is materialised once (n_rows) and reused by the opted-in estimator, + while the non-opted estimator triggers its own full scan => 2 * n_rows. + """ + from pyspark.sql import functions as F + from pyspark.sql.types import DoubleType as SparkDoubleType + + n_rows = 400 + source = ( + spark_session.range(0, n_rows) + .select(F.col("id").cast("double").alias("raw")) + .persist() + ) + source.count() + + accum = spark_session.sparkContext.accumulator(0) + + def _count(value): + accum.add(1) + return value + + udf = F.udf(_count, SparkDoubleType()).asNondeterministic() + counting = source.withColumn("x", udf(F.col("raw"))).drop("raw") + + with pytest.warns(UserWarning): + KamaeSparkPipeline( + stages=[ + # Opts in: reads the shared sample. + ConditionalStandardScaleEstimator( + inputCol="x", outputCol="x_sampled", useFitSample=True + ), + # No useFitSample: reads the full input. + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_full"), + ], + fitSampleFraction=1.0, + fitSampleSeed=1, + ).fit(counting) + + # One materialisation of the shared sample (n_rows), reused by the opted-in + # estimator, plus one full scan by the non-opted estimator. + assert accum.value == 2 * n_rows + source.unpersist() + + def test_spark_pipeline_fit_sample_fraction_without_opt_in_is_noop( + self, spark_session + ): + """ + fitSampleFraction with no estimator opting in (none has useFitSample=True) + warns that it has no effect and fits exactly as a full fit would. + """ + df = spark_session.createDataFrame([(1.0,), (2.0,), (3.0,), (4.0,)], ["x"]) + stage = ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_scaled") + + full_model = KamaeSparkPipeline(stages=[stage]).fit(df) + with pytest.warns(UserWarning, match="no effect"): + noop_model = KamaeSparkPipeline( + stages=[stage], fitSampleFraction=0.5, fitSampleSeed=1 + ).fit(df) + + assert noop_model.stages[0].getMean() == full_model.stages[0].getMean() + assert noop_model.stages[0].getStddev() == full_model.stages[0].getStddev() + + def test_spark_pipeline_fit_sample_fraction_overrides_estimator_sample_fraction( + self, spark_session + ): + """ + An estimator that sets both useFitSample=True and its own sampleFraction warns + that the shared pipeline sample wins and its sampleFraction is ignored. + """ + df = spark_session.createDataFrame([(1.0,), (2.0,), (3.0,), (4.0,)], ["x"]) + pipeline = KamaeSparkPipeline( + stages=[ + ConditionalStandardScaleEstimator( + inputCol="x", + outputCol="x_scaled", + useFitSample=True, + sampleFraction=0.5, + ) + ], + fitSampleFraction=1.0, + fitSampleSeed=1, + ) + with pytest.warns(UserWarning, match="sampleFraction is ignored"): + model = pipeline.fit(df) + assert model is not None + # The estimator's own sampleFraction is restored after the fit. + assert pipeline.getStages()[0].getSampleFraction() == 0.5 + + def test_spark_pipeline_fit_sample_fraction_opt_in_without_pipeline_sample_warns( + self, spark_session + ): + """ + useFitSample=True with no pipeline fitSampleFraction has no shared sample to + fit on, so it warns that useFitSample has no effect. + """ + df = spark_session.createDataFrame([(1.0,), (2.0,), (3.0,), (4.0,)], ["x"]) + pipeline = KamaeSparkPipeline( + stages=[ + ConditionalStandardScaleEstimator( + inputCol="x", outputCol="x_scaled", useFitSample=True + ) + ] + ) + with pytest.warns(UserWarning, match="useFitSample has no effect"): + model = pipeline.fit(df) + assert model is not None + @pytest.mark.parametrize( "stages, input_col, original_dtype", [ From bda7ae8149abf7147c2dda47b04a38c9844746f8 Mon Sep 17 00:00:00 2001 From: cworthington Date: Fri, 21 Aug 2026 21:46:22 +0100 Subject: [PATCH 35/36] fix: preserve null semantics in vectorized UDF and pipeline params on 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 --- src/kamae/spark/pipeline/pipeline.py | 53 ++++++++++++++++++- src/kamae/spark/utils/transform_utils.py | 7 +++ tests/kamae/spark/pipeline/test_pipeline.py | 25 +++++++++ tests/kamae/spark/utils/__init__.py | 13 +++++ .../kamae/spark/utils/test_transform_utils.py | 45 ++++++++++++++++ 5 files changed, 141 insertions(+), 2 deletions(-) create mode 100644 tests/kamae/spark/utils/__init__.py create mode 100644 tests/kamae/spark/utils/test_transform_utils.py diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index c4a97eba..538b0a43 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import warnings from typing import TYPE_CHECKING, Any, List, Optional, Set, Tuple, Type @@ -20,7 +21,7 @@ from pyspark.ml import Pipeline from pyspark.ml.param import Param, Params, TypeConverters from pyspark.ml.pipeline import PipelineReader, PipelineSharedReadWrite, PipelineWriter -from pyspark.ml.util import DefaultParamsReader, MLWriter +from pyspark.ml.util import DefaultParamsReader, DefaultParamsWriter, MLWriter from pyspark.sql import DataFrame from pyspark.storagelevel import StorageLevel @@ -830,7 +831,15 @@ def load(self, path: str) -> KamaeSparkPipeline: """ metadata = DefaultParamsReader.loadMetadata(path, self.sc) uid, stages = PipelineSharedReadWrite.load(metadata, self.sc, path) - return KamaeSparkPipeline(stages=stages)._resetUid(uid) + pipeline = KamaeSparkPipeline(stages=stages)._resetUid(uid) + # The base pipeline writer only persists stage uids, so the pipeline-level + # fit params (checkpointInterval, cache/prune flags, fitSampleFraction, ...) + # would reset to defaults. Restore any that were saved by the writer. + saved_params = metadata.get("kamaePipelineParams", {}) + for name, value in saved_params.items(): + if pipeline.hasParam(name): + pipeline.set(pipeline.getParam(name), value) + return pipeline class KamaeSparkPipelineWriter(PipelineWriter): @@ -840,3 +849,43 @@ class KamaeSparkPipelineWriter(PipelineWriter): def __init__(self, instance: KamaeSparkPipeline) -> None: super().__init__(instance=instance) + + def saveImpl(self, path: str) -> None: + """ + Saves the pipeline to the given path. + + Mirrors PipelineSharedReadWrite.saveImpl (metadata + stage uids + stages) but + additionally persists the pipeline-level fit params (which the base writer + drops) so they survive a save/load round-trip. + + :param path: Path to store the pipeline at. + :returns: None. + """ + stages = self.instance.getStages() + PipelineSharedReadWrite.validateStages(stages) + + json_params = { + "stageUids": [stage.uid for stage in stages], + "language": "Python", + } + # Only the explicitly-set, non-stages params; defaults stay implicit so + # older saves (without this metadata) still load with correct defaults. + pipeline_params = { + p.name: self.instance.getOrDefault(p) + for p in self.instance.params + if p.name != "stages" and self.instance.isSet(p) + } + DefaultParamsWriter.saveMetadata( + self.instance, + path, + self.sc, + extraMetadata={"kamaePipelineParams": pipeline_params}, + paramMap=json_params, + ) + stages_dir = os.path.join(path, "stages") + for index, stage in enumerate(stages): + stage.write().save( + PipelineSharedReadWrite.getStagePath( + stage.uid, index, len(stages), stages_dir + ) + ) diff --git a/src/kamae/spark/utils/transform_utils.py b/src/kamae/spark/utils/transform_utils.py index 87425aeb..92304b10 100644 --- a/src/kamae/spark/utils/transform_utils.py +++ b/src/kamae/spark/utils/transform_utils.py @@ -157,6 +157,13 @@ def _single_input_single_output_udf_transform( if not isinstance(input_col_datatype, ArrayType): def _vectorized_func(series: pd.Series) -> pd.Series: + # Arrow delivers Spark NULLs as NaN/pd.NA for numeric series rather than + # Python None, which would bypass the `is None` null/OOV guards in the + # element funcs (e.g. indexer/hash UDFs). Restore None so the vectorized + # path matches the plain UDF. Guarded by hasnans to keep the null-free + # fast path (the common case) untouched. + if series.hasnans: + series = series.astype(object).where(series.notna(), None) return series.map(nested_lambda_func) udf_func = F.pandas_udf(_vectorized_func, udf_return_type) diff --git a/tests/kamae/spark/pipeline/test_pipeline.py b/tests/kamae/spark/pipeline/test_pipeline.py index 5d24ed41..f910f192 100644 --- a/tests/kamae/spark/pipeline/test_pipeline.py +++ b/tests/kamae/spark/pipeline/test_pipeline.py @@ -475,6 +475,31 @@ def test_spark_read_write_pipeline(self, test_dir, stages, request): pipeline_loaded = KamaeSparkPipeline.load(f"{test_dir}/pipeline") assert pipeline.stages == pipeline_loaded.stages + def test_spark_read_write_pipeline_preserves_fit_params(self, test_dir, request): + """ + Non-default pipeline-level fit params must survive a save/load round-trip; + the base pipeline writer only persists stage uids and would drop them. + """ + stages = request.getfixturevalue("valid_stages_0") + pipeline = KamaeSparkPipeline( + stages=stages, + checkpointInterval=3, + cacheIntermediateData=True, + pruneInputColumns=False, + cacheEstimatorInput=True, + fitSampleFraction=0.25, + fitSampleSeed=7, + ) + pipeline.save(f"{test_dir}/pipeline_params") + loaded = KamaeSparkPipeline.load(f"{test_dir}/pipeline_params") + + assert loaded.getCheckpointInterval() == 3 + assert loaded.getCacheIntermediateData() is True + assert loaded.getPruneInputColumns() is False + assert loaded.getCacheEstimatorInput() is True + assert loaded.getFitSampleFraction() == 0.25 + assert loaded.getFitSampleSeed() == 7 + @pytest.mark.parametrize( "stages, expanded_stages", [ diff --git a/tests/kamae/spark/utils/__init__.py b/tests/kamae/spark/utils/__init__.py new file mode 100644 index 00000000..d47f0081 --- /dev/null +++ b/tests/kamae/spark/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright [2024] Expedia, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/kamae/spark/utils/test_transform_utils.py b/tests/kamae/spark/utils/test_transform_utils.py new file mode 100644 index 00000000..ba4a232f --- /dev/null +++ b/tests/kamae/spark/utils/test_transform_utils.py @@ -0,0 +1,45 @@ +# Copyright [2024] Expedia, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pyspark.sql.functions as F +from pyspark.sql.types import DoubleType, IntegerType, StructField, StructType + +from kamae.spark.utils.transform_utils import ( + single_input_single_output_scalar_udf_transform, +) + + +class TestTransformUtils: + def test_scalar_udf_transform_maps_numeric_null_to_none(self, spark_session): + """ + Spark NULLs in a nullable numeric column must reach the element func as Python + None, not Arrow's NaN. Otherwise `is None` null/OOV guards in the element funcs + (indexer/hash UDFs) silently misroute missing values. The func here would raise + on NaN (int(NaN)), so this fails if the vectorized path leaks NaN through. + """ + schema = StructType([StructField("x", DoubleType(), True)]) + df = spark_session.createDataFrame([(1.0,), (None,), (2.0,)], schema) + + out = df.withColumn( + "y", + single_input_single_output_scalar_udf_transform( + input_col=F.col("x"), + input_col_datatype=df.schema["x"].dataType, + func=lambda v: -1 if v is None else int(v), + udf_return_element_datatype=IntegerType(), + ), + ) + + result = sorted(row["y"] for row in out.collect()) + assert result == [-1, 1, 2] From cfde6b7b5d4ef164ad605db9720fcc72726403ad Mon Sep 17 00:00:00 2001 From: cworthington Date: Mon, 24 Aug 2026 13:50:33 +0100 Subject: [PATCH 36/36] refactor: use Kamae metadata read/write in pipeline save/load 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 --- src/kamae/spark/pipeline/pipeline.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index 538b0a43..e4e5742d 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -21,12 +21,16 @@ from pyspark.ml import Pipeline from pyspark.ml.param import Param, Params, TypeConverters from pyspark.ml.pipeline import PipelineReader, PipelineSharedReadWrite, PipelineWriter -from pyspark.ml.util import DefaultParamsReader, DefaultParamsWriter, MLWriter +from pyspark.ml.util import MLWriter from pyspark.sql import DataFrame from pyspark.storagelevel import StorageLevel from kamae.graph import PipelineGraph from kamae.spark.estimators import BaseEstimator +from kamae.spark.params.default_read_write import ( + KamaeDefaultParamsReader, + KamaeDefaultParamsWriter, +) from kamae.spark.pipeline import KamaeSparkPipelineModel from kamae.spark.transformers import BaseTransformer @@ -829,7 +833,7 @@ def load(self, path: str) -> KamaeSparkPipeline: :param path: Path to stored pipeline. :returns: KamaeSparkPipeline object. """ - metadata = DefaultParamsReader.loadMetadata(path, self.sc) + metadata = KamaeDefaultParamsReader.loadMetadata(path, self.sc) uid, stages = PipelineSharedReadWrite.load(metadata, self.sc, path) pipeline = KamaeSparkPipeline(stages=stages)._resetUid(uid) # The base pipeline writer only persists stage uids, so the pipeline-level @@ -875,7 +879,7 @@ def saveImpl(self, path: str) -> None: for p in self.instance.params if p.name != "stages" and self.instance.isSet(p) } - DefaultParamsWriter.saveMetadata( + KamaeDefaultParamsWriter.saveMetadata( self.instance, path, self.sc,