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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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"