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", diff --git a/src/kamae/spark/estimators/conditional_standard_scale.py b/src/kamae/spark/estimators/conditional_standard_scale.py index b6edfb2f..aa0e6098 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 ( @@ -258,6 +259,7 @@ def __init__( epsilon: float = 0, nanFillValue: Optional[float] = None, sampleFraction: Optional[float] = None, + useFitSample: bool = False, ) -> None: """ Initializes a ConditionalStandardScaleEstimator estimator. @@ -287,6 +289,8 @@ def __init__( to use it if epsilon filters out all the values. Defaults to None. :param sampleFraction: Fraction of data to sample for statistics estimation (exclusive 0.0-1.0). Default None (no sampling). + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. :returns: None - class instantiated. """ super().__init__() @@ -300,6 +304,7 @@ def __init__( epsilon=0, nanFillValue=None, sampleFraction=None, + useFitSample=False, ) kwargs = self._input_kwargs self.setParams(**kwargs) @@ -378,22 +383,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, diff --git a/src/kamae/spark/estimators/impute.py b/src/kamae/spark/estimators/impute.py index 56a4e83e..8498a57e 100644 --- a/src/kamae/spark/estimators/impute.py +++ b/src/kamae/spark/estimators/impute.py @@ -66,6 +66,7 @@ def __init__( maskValue: Optional[Union[float, int, str]] = None, imputeMethod: Optional[str] = None, sampleFraction: Optional[float] = None, + useFitSample: bool = False, ) -> None: """ Initializes a ImputeEstimator estimator. @@ -86,12 +87,15 @@ def __init__( Valid values are "mean" or "median". :param sampleFraction: Fraction of data to sample for statistics estimation (exclusive 0.0-1.0). Default None (no sampling). + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. :returns: None - class instantiated. """ super().__init__() self._setDefault( imputeMethod="mean", sampleFraction=None, + useFitSample=False, ) self.valid_impute_methods = ["mean", "median"] kwargs = self._input_kwargs diff --git a/src/kamae/spark/estimators/min_max_scale.py b/src/kamae/spark/estimators/min_max_scale.py index 36a6430f..31c878f8 100644 --- a/src/kamae/spark/estimators/min_max_scale.py +++ b/src/kamae/spark/estimators/min_max_scale.py @@ -65,6 +65,7 @@ def __init__( layerName: Optional[str] = None, maskValue: Optional[float] = None, sampleFraction: Optional[float] = None, + useFitSample: bool = False, ) -> None: """ Initializes a MinMaxScaleEstimator estimator. @@ -82,10 +83,12 @@ def __init__( during the computation of the min and max values. :param sampleFraction: Fraction of data to sample for statistics estimation (exclusive 0.0-1.0). Default None (no sampling). + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. :returns: None - class instantiated. """ super().__init__() - self._setDefault(maskValue=None, sampleFraction=None) + self._setDefault(maskValue=None, sampleFraction=None, useFitSample=False) kwargs = self._input_kwargs self.setParams(**kwargs) diff --git a/src/kamae/spark/estimators/single_feature_array_standard_scale.py b/src/kamae/spark/estimators/single_feature_array_standard_scale.py index 4c209893..121a0fbe 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 ( @@ -61,6 +62,7 @@ def __init__( layerName: Optional[str] = None, maskValue: Optional[float] = None, sampleFraction: Optional[float] = None, + useFitSample: bool = False, ) -> None: """ Initializes a SingleFeatureArrayStandardScaleEstimator estimator. @@ -76,10 +78,12 @@ def __init__( in the keras model. If not set, we use the uid of the Spark transformer. :param sampleFraction: Fraction of data to sample for statistics estimation (exclusive 0.0-1.0). Default None (no sampling). + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. :returns: None - class instantiated. """ super().__init__() - self._setDefault(maskValue=None, sampleFraction=None) + self._setDefault(maskValue=None, sampleFraction=None, useFitSample=False) kwargs = self._input_kwargs self.setParams(**kwargs) @@ -113,32 +117,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 a1c654ea..0d0b44aa 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 ( @@ -65,6 +66,7 @@ def __init__( layerName: Optional[str] = None, maskValue: Optional[float] = None, sampleFraction: Optional[float] = None, + useFitSample: bool = False, ) -> None: """ Initializes a StandardScaleEstimator estimator. @@ -80,10 +82,12 @@ def __init__( in the keras model. If not set, we use the uid of the Spark transformer. :param sampleFraction: Fraction of data to sample for statistics estimation (exclusive 0.0-1.0). Default None (no sampling). + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. :returns: None - class instantiated. """ super().__init__() - self._setDefault(maskValue=None, sampleFraction=None) + self._setDefault(maskValue=None, sampleFraction=None, useFitSample=False) kwargs = self._input_kwargs self.setParams(**kwargs) @@ -113,41 +117,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/src/kamae/spark/params/base.py b/src/kamae/spark/params/base.py index 071305ad..a96967e2 100644 --- a/src/kamae/spark/params/base.py +++ b/src/kamae/spark/params/base.py @@ -144,6 +144,15 @@ class SampleFractionParams(Params): "Default None (no sampling).", ) + useFitSample = Param( + Params._dummy(), + "useFitSample", + "If True, and the enclosing KamaeSparkPipeline has fitSampleFraction set, fit " + "this estimator on the pipeline's shared sample instead of the full input. " + "Provides opt-in per estimator without a second fraction value. Default False.", + typeConverter=TypeConverters.toBoolean, + ) + def setSampleFraction(self, value: float) -> "SampleFractionParams": """ Sets the parameter sampleFraction to the given float value. @@ -167,6 +176,24 @@ def getSampleFraction(self) -> Optional[float]: """ return self.getOrDefault(self.sampleFraction) + def setUseFitSample(self, value: bool) -> "SampleFractionParams": + """ + Sets the parameter useFitSample to the given boolean value. + + :param value: Whether to fit this estimator on the pipeline's shared sample + when fitSampleFraction is set. + :returns: Instance of class mixed in. + """ + return self._set(useFitSample=value) + + def getUseFitSample(self) -> bool: + """ + Gets the value of the useFitSample parameter. + + :returns: Whether this estimator opts in to the pipeline's shared sample. + """ + return self.getOrDefault(self.useFitSample) + class SingleInputParams(HasInputCol): """ diff --git a/src/kamae/spark/pipeline/pipeline.py b/src/kamae/spark/pipeline/pipeline.py index 1da0d7ea..e4e5742d 100644 --- a/src/kamae/spark/pipeline/pipeline.py +++ b/src/kamae/spark/pipeline/pipeline.py @@ -12,18 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import TYPE_CHECKING, List, Optional, Type +import os +import warnings +from typing import TYPE_CHECKING, Any, List, Optional, Set, Tuple, Type 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.ml.util import MLWriter from pyspark.sql import DataFrame +from pyspark.storagelevel import StorageLevel from kamae.graph import PipelineGraph from kamae.spark.estimators import BaseEstimator +from kamae.spark.params.default_read_write import ( + KamaeDefaultParamsReader, + KamaeDefaultParamsWriter, +) from kamae.spark.pipeline import KamaeSparkPipelineModel from kamae.spark.transformers import BaseTransformer @@ -38,17 +45,124 @@ 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. + + Five opt-in fit optimisations are available, all defaulting off (fit behaviour + unchanged): `checkpointInterval` reliably checkpoints every N stages to bound + logical-plan depth (requires a checkpoint dir); `cacheIntermediateData` persists + the working DataFrame at each estimator-fit boundary to avoid re-scanning the + upstream lineage; `pruneInputColumns` drops input columns no stage consumes; + `cacheEstimatorInput` projects to the columns still read downstream at the first + estimator boundary and persists that narrow frame once, so independent sibling + estimators reuse it instead of re-scanning the wide input; `fitSampleFraction` + draws a single persisted sample of the input up-front and fits the estimators + that opt in (those with `useFitSample=True`) from that shared sample, while + estimators without it keep fitting on the full input. """ + checkpointInterval = Param( + Params._dummy(), + "checkpointInterval", + "Stages between reliable checkpoint(eager=True) calls during fit, to bound " + "logical-plan depth. Requires a checkpoint dir. None (default) disables it.", + typeConverter=TypeConverters.toInt, + ) + + cacheIntermediateData = Param( + Params._dummy(), + "cacheIntermediateData", + "If True, persist the working DataFrame (MEMORY_AND_DISK) at each " + "estimator-fit boundary to avoid re-scanning the upstream lineage. False " + "(default) disables it.", + typeConverter=TypeConverters.toBoolean, + ) + + pruneInputColumns = Param( + Params._dummy(), + "pruneInputColumns", + "If True, drop input columns no stage consumes before fitting. False " + "(default) disables it.", + typeConverter=TypeConverters.toBoolean, + ) + + cacheEstimatorInput = Param( + Params._dummy(), + "cacheEstimatorInput", + "If True, at the first estimator-fit boundary project the working DataFrame " + "to the columns still read downstream and persist (MEMORY_AND_DISK) that " + "narrow frame once, reused by all subsequent estimators. Competes with " + "cacheIntermediateData; enable at most one. False (default) disables it.", + typeConverter=TypeConverters.toBoolean, + ) + + fitSampleFraction = Param( + Params._dummy(), + "fitSampleFraction", + "If set to a float in (0, 1], draw a single sample of the input up-front, " + "persist (MEMORY_AND_DISK) and materialise it once, and fit the estimators " + "that opt in from that shared sample. An estimator opts in by setting its " + "boolean useFitSample param to True. Estimators without useFitSample fit on " + "the full input, so leave it False for estimators needing exact/global " + "statistics (vocabulary builders like StringIndexer/OneHot, min/max scalers, " + "distinct counts) and set it True on sample-robust estimators (mean/std/" + "quantiles, e.g. ConditionalStandardScale). Avoids re-scanning or persisting " + "the wide source once per opted-in estimator. Disables cacheIntermediateData " + "and cacheEstimatorInput. None (default) disables it.", + typeConverter=TypeConverters.toFloat, + ) + + fitSampleSeed = Param( + Params._dummy(), + "fitSampleSeed", + "Optional integer seed passed to the fitSampleFraction sample for " + "reproducibility. None (default) leaves the sample unseeded.", + typeConverter=TypeConverters.toInt, + ) + @keyword_only - def __init__(self, *, stages: Optional[List["KamaePipelineStage"]] = None) -> None: + def __init__( + self, + *, + stages: Optional[List["KamaePipelineStage"]] = None, + checkpointInterval: Optional[int] = None, + cacheIntermediateData: bool = False, + pruneInputColumns: bool = False, + cacheEstimatorInput: bool = False, + fitSampleFraction: Optional[float] = None, + fitSampleSeed: Optional[int] = None, + ) -> 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. 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. + :param pruneInputColumns: If True, drop input columns no stage consumes + before fitting. False (default) disables it. + :param cacheEstimatorInput: If True, project to the columns still read + downstream at the first estimator boundary and persist that narrow frame + once for reuse by subsequent estimators. False (default) disables it. + :param fitSampleFraction: If set to a float in (0, 1], fit the estimators + that opt in (those with useFitSample=True) from a single up-front persisted + sample of the input; estimators without useFitSample fit on the full input. + None (default) disables it. + :param fitSampleSeed: Optional integer seed for the fitSampleFraction + sample. None (default) leaves it unseeded. :returns: None - class instantiated. """ + kwargs = self._input_kwargs super().__init__(stages=stages) + self._setDefault( + checkpointInterval=None, + cacheIntermediateData=False, + pruneInputColumns=False, + cacheEstimatorInput=False, + fitSampleFraction=None, + fitSampleSeed=None, + ) + self.setParams(**kwargs) def setStages(self, value: List["KamaePipelineStage"]) -> "KamaeSparkPipeline": """ @@ -67,18 +181,162 @@ def getStages(self) -> List["KamaePipelineStage"]: """ return self.getOrDefault("stages") + def setCheckpointInterval(self, value: Optional[int]) -> "KamaeSparkPipeline": + """ + Sets the `checkpointInterval` parameter. + + :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) -> Optional[int]: + """ + Gets the value of the `checkpointInterval` parameter. + + :returns: The checkpointInterval value. + """ + return self.getOrDefault(self.checkpointInterval) + + def setCacheIntermediateData(self, value: bool) -> "KamaeSparkPipeline": + """ + Sets the `cacheIntermediateData` parameter. + + :param value: Whether to persist the working DataFrame at each + estimator-fit boundary during fit. + :returns: KamaeSparkPipeline object with cacheIntermediateData set. + """ + 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) + + 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) + + def setCacheEstimatorInput(self, value: bool) -> "KamaeSparkPipeline": + """ + Sets the `cacheEstimatorInput` parameter. + + :param value: Whether to project and persist a narrow estimator-input + frame once at the first estimator boundary during fit. + :returns: KamaeSparkPipeline object with cacheEstimatorInput set. + """ + return self._set(cacheEstimatorInput=value) + + def getCacheEstimatorInput(self) -> bool: + """ + Gets the value of the `cacheEstimatorInput` parameter. + + :returns: The cacheEstimatorInput value. + """ + return self.getOrDefault(self.cacheEstimatorInput) + + def setFitSampleFraction(self, value: Optional[float]) -> "KamaeSparkPipeline": + """ + Sets the `fitSampleFraction` parameter. + + :param value: Fraction in (0, 1] of the input to sample once up-front and + fit every estimator from. None disables it. + :returns: KamaeSparkPipeline object with fitSampleFraction set. + :raises ValueError: If value is not None and not in the range (0, 1]. + """ + if value is not None and not 0.0 < value <= 1.0: + raise ValueError( + f"fitSampleFraction must be in the range (0, 1] or None, got {value}." + ) + return self._set(fitSampleFraction=value) + + def getFitSampleFraction(self) -> Optional[float]: + """ + Gets the value of the `fitSampleFraction` parameter. + + :returns: The fitSampleFraction value. + """ + return self.getOrDefault(self.fitSampleFraction) + + def setFitSampleSeed(self, value: Optional[int]) -> "KamaeSparkPipeline": + """ + Sets the `fitSampleSeed` parameter. + + :param value: Integer seed for the fitSampleFraction sample, or None. + :returns: KamaeSparkPipeline object with fitSampleSeed set. + """ + return self._set(fitSampleSeed=value) + + def getFitSampleSeed(self) -> Optional[int]: + """ + Gets the value of the `fitSampleSeed` parameter. + + :returns: The fitSampleSeed value. + """ + return self.getOrDefault(self.fitSampleSeed) + @keyword_only def setParams( - self, *, stages: Optional["KamaePipelineStage"] = None + self, + *, + stages: Optional["KamaePipelineStage"] = None, + checkpointInterval: Optional[int] = None, + cacheIntermediateData: bool = False, + pruneInputColumns: bool = False, + cacheEstimatorInput: bool = False, + fitSampleFraction: Optional[float] = None, + fitSampleSeed: Optional[int] = None, ) -> "KamaeSparkPipeline": """ Sets the keyword arguments of the pipeline. + Routes each supplied param through its setter so setter-level validation + (e.g. checkpointInterval) runs. + :param stages: List of pipeline stages. - :returns: KamaeSparkPipeline object with stages set. + :param checkpointInterval: Number of stages between reliable + 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. + :param cacheEstimatorInput: If True, project to the columns still read + downstream at the first estimator boundary and persist that narrow frame + once for reuse by subsequent estimators. False (default) disables it. + :param fitSampleFraction: If set to a float in (0, 1], fit every estimator + from a single up-front persisted sample of the input. None (default) + disables it. + :param fitSampleSeed: Optional integer seed for the fitSampleFraction + sample. None (default) leaves it unseeded. + :returns: KamaeSparkPipeline object with params set. """ - 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"]: """ @@ -132,6 +390,99 @@ def collect_estimator_parents( ] return estimator_parent_stages + @staticmethod + def collect_required_input_columns( + stages: List["KamaePipelineStage"], + ) -> Set[str]: + """ + 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 (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. + """ + 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( + self, + dataset: DataFrame, + stages: List["KamaePipelineStage"], + ) -> DataFrame: + """ + Projects the input DataFrame down to only the columns the pipeline reads. + + Returned unchanged if there are no unused columns to drop. + + :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: + 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. @@ -139,18 +490,36 @@ 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. + Optionally applies the opt-in fit optimisations (`pruneInputColumns`, + `checkpointInterval`, `cacheIntermediateData`, `cacheEstimatorInput`, + `fitSampleFraction`); see the class docstring. With the exception of + `fitSampleFraction`, all preserve data exactly, so fitted results match the + defaults-off behaviour. + + If both `cacheIntermediateData` and `cacheEstimatorInput` are enabled, + `cacheEstimatorInput` takes precedence (a warning is emitted) and + `cacheIntermediateData` is ignored, since the narrow frame is a strictly + smaller cache and the intermediate cache would evict it. + + If `fitSampleFraction` is set, the input is sampled once, persisted and + materialised, and the estimators that opt in (those with `useFitSample=True`) + are fit from that shared sample with any `sampleFraction` they also set + temporarily disabled; every other estimator still fits on the full input. + `cacheIntermediateData` and `cacheEstimatorInput` are disabled (with a + warning) as they persist frames this option avoids; a warning also names the + opted-in estimators. See the `fitSampleFraction` param docstring. + :param dataset: PySpark DataFrame to fit the pipeline to. :returns: KamaeSparkPipelineModel object. + :raises ValueError: If checkpointing is enabled but no checkpoint directory + 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) - ) + # 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) # Native Spark checks for the last estimator and executes all transformers # before it, regardless whether there is a dependency between them. See here: @@ -162,19 +531,258 @@ def _fit(self, dataset: DataFrame) -> "KamaeSparkPipelineModel": estimator_parent_stages = self.collect_estimator_parents( expanded_pipeline_stages ) - # Fit each stage, appending the transformer to the list of transformers + # Opt-in plan-depth bounding. 0 (or None) = no change. + checkpoint_interval = self.getCheckpointInterval() + checkpoint_enabled = self._resolve_checkpoint_enabled( + dataset, checkpoint_interval + ) + cache_enabled = self.getCacheIntermediateData() + cache_estimator_input = self.getCacheEstimatorInput() + # Competing caching strategies - both would persist the wide frame, and the + # intermediate cache's per-boundary unpersist would evict the narrow frame. + # cacheEstimatorInput wins: it persists a strictly narrower frame once. + if cache_enabled and cache_estimator_input: + warnings.warn( + "cacheIntermediateData and cacheEstimatorInput are competing " + "caching strategies; cacheEstimatorInput takes precedence and " + "cacheIntermediateData is ignored.", + stacklevel=2, + ) + cache_enabled = False + + # Opt-in: draw one shared sample up-front and fit only the estimators that + # opt in against it, instead of each re-scanning or persisting the wide + # source. An estimator opts in via its boolean `useFitSample` param. + # Estimators without it set still fit on the full input, so exact/global- + # statistic estimators (vocabulary builders, min/max) stay correct. + fit_sample_fraction = self.getFitSampleFraction() + sampled_dataset: Optional[DataFrame] = None + # (stage, param, original_value) for restoration of overridden sampleFractions. + overridden_sample_fractions: List[Tuple[Any, Any, Any]] = [] + # Identities of the estimators that fit on the shared sample. + use_sample_stage_ids: Set[int] = set() + if fit_sample_fraction is not None: + sampling_estimators = [ + stage + for stage in expanded_pipeline_stages + if isinstance(stage, BaseEstimator) + and stage.hasParam("useFitSample") + and stage.getUseFitSample() + ] + if not sampling_estimators: + warnings.warn( + "fitSampleFraction is set but no estimator has useFitSample=True, " + "so nothing opts in to the shared sample and fitSampleFraction " + "has no effect. Set useFitSample=True on the estimators that " + "should fit on the sample.", + stacklevel=2, + ) + else: + warnings.warn( + "fitSampleFraction fits " + f"{sorted({type(s).__name__ for s in sampling_estimators})} on " + "one shared pipeline sample (useFitSample=True). Estimators " + "without useFitSample set fit on the full input.", + stacklevel=2, + ) + # fitSampleFraction persists a tiny sample instead of the wide frame, + # so the caching strategies it supersedes are turned off. + if cache_enabled or cache_estimator_input: + warnings.warn( + "fitSampleFraction is incompatible with cacheIntermediateData " + "and cacheEstimatorInput (which persist frames this option " + "avoids); the caching options are disabled.", + stacklevel=2, + ) + cache_enabled = False + cache_estimator_input = False + sampled_dataset = dataset.sample( + fraction=fit_sample_fraction, seed=self.getFitSampleSeed() + ).persist(StorageLevel.MEMORY_AND_DISK) + # Force one materialisation so the sample is computed exactly once and + # every opted-in estimator reads the cached rows instead of rescanning. + sampled_dataset.count() + # The shared sample is already drawn, so opted-in estimators must not + # sample again (fraction-of-a-fraction would leave far too few rows). + # Disable any sampleFraction they also set, restoring it in finally. + for stage in sampling_estimators: + use_sample_stage_ids.add(id(stage)) + param = stage.getParam("sampleFraction") + if stage.isSet(param): + overridden_sample_fractions.append( + (stage, param, stage.getOrDefault(param)) + ) + stage.set(param, None) + overridden_names = sorted( + {type(s).__name__ for s, _, _ in overridden_sample_fractions} + ) + if overridden_names: + warnings.warn( + f"{overridden_names} set both useFitSample=True and their own " + "sampleFraction. The shared pipeline sample wins, so their " + "sampleFraction is ignored for this fit. Unset one to silence " + "this warning.", + stacklevel=2, + ) + else: + opted_in_without_sample = sorted( + { + type(stage).__name__ + for stage in expanded_pipeline_stages + if isinstance(stage, BaseEstimator) + and stage.hasParam("useFitSample") + and stage.getUseFitSample() + } + ) + if opted_in_without_sample: + warnings.warn( + f"{opted_in_without_sample} set useFitSample=True but the pipeline " + "has no fitSampleFraction, so there is no shared sample to fit on " + "and useFitSample has no effect. Set fitSampleFraction on the " + "pipeline to enable sampled fitting.", + stacklevel=2, + ) + + # Fit each stage, appending the transformer to the list of transformers. # If the stage is a parent of an estimator, transform the dataset. transformers: List[BaseTransformer] = [] - for stage in expanded_pipeline_stages: + try: + fitted_pipeline_model = self._run_fit_loop( + expanded_pipeline_stages=expanded_pipeline_stages, + dataset=dataset, + sampled_dataset=sampled_dataset, + use_sample_stage_ids=use_sample_stage_ids, + estimator_parent_stages=estimator_parent_stages, + transformers=transformers, + checkpoint_enabled=checkpoint_enabled, + checkpoint_interval=checkpoint_interval, + cache_enabled=cache_enabled, + cache_estimator_input=cache_estimator_input, + ) + finally: + # Restore each opted-in estimator's original sampleFraction and release + # the shared sample, whether or not the fit succeeded. + for stage, param, original in overridden_sample_fractions: + stage.set(param, original) + if sampled_dataset is not None: + sampled_dataset.unpersist() + return fitted_pipeline_model + + def _run_fit_loop( + self, + *, + expanded_pipeline_stages: List["KamaePipelineStage"], + dataset: DataFrame, + sampled_dataset: Optional[DataFrame], + use_sample_stage_ids: Set[int], + estimator_parent_stages: List["KamaePipelineStage"], + transformers: List[BaseTransformer], + checkpoint_enabled: bool, + checkpoint_interval: Optional[int], + cache_enabled: bool, + cache_estimator_input: bool, + ) -> "KamaeSparkPipelineModel": + """ + Runs the stage-by-stage fit loop, applying the checkpoint/cache optimisations. + + Extracted from `_fit` so the loop can run inside a try/finally that restores + estimator sampling and releases the shared sample when fitSampleFraction is + used. With fitSampleFraction off (`sampled_dataset is None`) behaviour is + identical to the previous inline loop. + + When fitSampleFraction is active, two lineages are carried in parallel: the + full-data lineage and a shared-sample lineage. Every transformer that feeds an + estimator is applied to both, so an estimator can fit on whichever it opted + in to - estimators in `use_sample_stage_ids` fit on the sample, all others on + the full input. The caching options are mutually exclusive with + fitSampleFraction (disabled in `_fit`), so they only ever act on the full + lineage when no sample is present. + + `cacheIntermediateData` and `cacheEstimatorInput` are mutually exclusive (see + `_fit`), so a single `persisted_frame` handle tracks whichever frame is + persisted. It is kept separate from the lineage variable because that variable + is reassigned by `model.transform(...)` between boundaries; the handle is what + lets us unpersist the actual persisted frame at the end. + + :param expanded_pipeline_stages: Flattened pipeline stages to fit. + :param dataset: Full-data DataFrame to fit non-sampled stages against. + :param sampled_dataset: Shared sample (or None when fitSampleFraction is off). + :param use_sample_stage_ids: `id()`s of estimators that fit on the sample. + :param estimator_parent_stages: Stages whose output an estimator consumes. + :param transformers: Accumulator list the fitted stages are appended to. + :param checkpoint_enabled: Whether reliable checkpointing is enabled. + :param checkpoint_interval: Stages between checkpoints (when enabled). + :param cache_enabled: Whether to persist at each estimator boundary. + :param cache_estimator_input: Whether to persist a narrow frame once. + :returns: KamaeSparkPipelineModel object. + """ + last_checkpoint_index = 0 + # The single persisted frame (if any), unpersisted once superseded or done. + # Only one of the mutually-exclusive cache strategies ever populates it. + persisted_frame: Optional[DataFrame] = None + estimator_input_cached = False + # Full-data lineage and, when fitSampleFraction is active, a parallel + # shared-sample lineage kept in lock-step through the estimator-feeding + # transforms. + full_dataset = dataset + sample_dataset = sampled_dataset + for index, stage in enumerate(expanded_pipeline_stages): if isinstance(stage, BaseTransformer): transformers.append(stage) if stage in estimator_parent_stages: - dataset = stage.transform(dataset) + full_dataset = stage.transform(full_dataset) + if sample_dataset is not None: + sample_dataset = stage.transform(sample_dataset) else: - model = stage.fit(dataset) + # cacheEstimatorInput: at the first estimator boundary, project to + # the columns still read downstream and persist that narrow frame + # once. Independent sibling estimators then fit against it instead of + # re-scanning the wide input. The persist sits below each estimator's + # in-fit sample, so sampling is unchanged. + if cache_estimator_input and not estimator_input_cached: + estimator_input_cached = True + live_columns = self.collect_required_input_columns( + expanded_pipeline_stages[index:] + ) + keep_columns = [ + c for c in full_dataset.columns if c in live_columns + ] + if keep_columns and len(keep_columns) < len(full_dataset.columns): + full_dataset = full_dataset.select(*keep_columns).persist( + StorageLevel.MEMORY_AND_DISK + ) + persisted_frame = full_dataset + # Truncate accumulated lineage before the fit action to bound plan + # depth. eager=True materialises now. + if ( + checkpoint_enabled + and index - last_checkpoint_index >= checkpoint_interval + ): + full_dataset = full_dataset.checkpoint(eager=True) + if sample_dataset is not None: + sample_dataset = sample_dataset.checkpoint(eager=True) + last_checkpoint_index = index + # cacheIntermediateData: persist so the fit action and downstream + # transforms reuse a materialised frame instead of re-scanning, + # releasing the previous frame first. One frame held at a time. + if cache_enabled: + if persisted_frame is not None: + persisted_frame.unpersist() + full_dataset = full_dataset.persist(StorageLevel.MEMORY_AND_DISK) + persisted_frame = full_dataset + fit_dataset = ( + sample_dataset + if sample_dataset is not None and id(stage) in use_sample_stage_ids + else full_dataset + ) + model = stage.fit(fit_dataset) transformers.append(model) if stage in estimator_parent_stages: - dataset = model.transform(dataset) + full_dataset = model.transform(full_dataset) + if sample_dataset is not None: + sample_dataset = model.transform(sample_dataset) + if persisted_frame is not None: + persisted_frame.unpersist() return KamaeSparkPipelineModel(transformers) def copy(self, extra: Optional["ParamMap"] = None) -> "KamaeSparkPipeline": @@ -225,9 +833,17 @@ def load(self, path: str) -> KamaeSparkPipeline: :param path: Path to stored pipeline. :returns: KamaeSparkPipeline object. """ - metadata = DefaultParamsReader.loadMetadata(path, self.sc) + metadata = KamaeDefaultParamsReader.loadMetadata(path, self.sc) uid, stages = PipelineSharedReadWrite.load(metadata, self.sc, path) - return KamaeSparkPipeline(stages=stages)._resetUid(uid) + pipeline = KamaeSparkPipeline(stages=stages)._resetUid(uid) + # The base pipeline writer only persists stage uids, so the pipeline-level + # fit params (checkpointInterval, cache/prune flags, fitSampleFraction, ...) + # would reset to defaults. Restore any that were saved by the writer. + saved_params = metadata.get("kamaePipelineParams", {}) + for name, value in saved_params.items(): + if pipeline.hasParam(name): + pipeline.set(pipeline.getParam(name), value) + return pipeline class KamaeSparkPipelineWriter(PipelineWriter): @@ -237,3 +853,43 @@ class KamaeSparkPipelineWriter(PipelineWriter): def __init__(self, instance: KamaeSparkPipeline) -> None: super().__init__(instance=instance) + + def saveImpl(self, path: str) -> None: + """ + Saves the pipeline to the given path. + + Mirrors PipelineSharedReadWrite.saveImpl (metadata + stage uids + stages) but + additionally persists the pipeline-level fit params (which the base writer + drops) so they survive a save/load round-trip. + + :param path: Path to store the pipeline at. + :returns: None. + """ + stages = self.instance.getStages() + PipelineSharedReadWrite.validateStages(stages) + + json_params = { + "stageUids": [stage.uid for stage in stages], + "language": "Python", + } + # Only the explicitly-set, non-stages params; defaults stay implicit so + # older saves (without this metadata) still load with correct defaults. + pipeline_params = { + p.name: self.instance.getOrDefault(p) + for p in self.instance.params + if p.name != "stages" and self.instance.isSet(p) + } + KamaeDefaultParamsWriter.saveMetadata( + self.instance, + path, + self.sc, + extraMetadata={"kamaePipelineParams": pipeline_params}, + paramMap=json_params, + ) + stages_dir = os.path.join(path, "stages") + for index, stage in enumerate(stages): + stage.write().save( + PipelineSharedReadWrite.getStagePath( + stage.uid, index, len(stages), stages_dir + ) + ) diff --git a/src/kamae/spark/transformers/bucketize.py b/src/kamae/spark/transformers/bucketize.py index b639f0cc..e9055cf7 100644 --- a/src/kamae/spark/transformers/bucketize.py +++ b/src/kamae/spark/transformers/bucketize.py @@ -16,21 +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 .base import BaseTransformer @@ -90,9 +90,8 @@ class BucketizeTransformer( The 0 index is reserved for masking/padding. """ - jit_compatible = True - supported_backends = TENSORFLOW_ONLY + jit_compatible = True @keyword_only def __init__( @@ -113,7 +112,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 +140,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( diff --git a/src/kamae/spark/utils/transform_utils.py b/src/kamae/spark/utils/transform_utils.py index c67d8773..92304b10 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. - 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 @@ -150,6 +150,25 @@ 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: + # Arrow delivers Spark NULLs as NaN/pd.NA for numeric series rather than + # Python None, which would bypass the `is None` null/OOV guards in the + # element funcs (e.g. indexer/hash UDFs). Restore None so the vectorized + # path matches the plain UDF. Guarded by hasnans to keep the null-free + # fast path (the common case) untouched. + if series.hasnans: + series = series.astype(object).where(series.notna(), None) + return series.map(nested_lambda_func) + + udf_func = F.pandas_udf(_vectorized_func, udf_return_type) + return udf_func(input_col) + udf_func = F.udf(nested_lambda_func, udf_return_type) return udf_func(input_col) 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..f910f192 100644 --- a/tests/kamae/spark/pipeline/test_pipeline.py +++ b/tests/kamae/spark/pipeline/test_pipeline.py @@ -12,14 +12,20 @@ # 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 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, ) @@ -39,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) @@ -469,6 +475,31 @@ def test_spark_read_write_pipeline(self, test_dir, stages, request): pipeline_loaded = KamaeSparkPipeline.load(f"{test_dir}/pipeline") assert pipeline.stages == pipeline_loaded.stages + def test_spark_read_write_pipeline_preserves_fit_params(self, test_dir, request): + """ + Non-default pipeline-level fit params must survive a save/load round-trip; + the base pipeline writer only persists stage uids and would drop them. + """ + stages = request.getfixturevalue("valid_stages_0") + pipeline = KamaeSparkPipeline( + stages=stages, + checkpointInterval=3, + cacheIntermediateData=True, + pruneInputColumns=False, + cacheEstimatorInput=True, + fitSampleFraction=0.25, + fitSampleSeed=7, + ) + pipeline.save(f"{test_dir}/pipeline_params") + loaded = KamaeSparkPipeline.load(f"{test_dir}/pipeline_params") + + assert loaded.getCheckpointInterval() == 3 + assert loaded.getCacheIntermediateData() is True + assert loaded.getPruneInputColumns() is False + assert loaded.getCacheEstimatorInput() is True + assert loaded.getFitSampleFraction() == 0.25 + assert loaded.getFitSampleSeed() == 7 + @pytest.mark.parametrize( "stages, expanded_stages", [ @@ -552,6 +583,611 @@ 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 (None). + """ + stages = request.getfixturevalue(stages) + + baseline_model = KamaeSparkPipeline(stages=stages, checkpointInterval=None).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=None).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 + + @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 + 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(None) + 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() + + def test_spark_pipeline_cache_estimator_input_is_transparent_to_fit( + self, spark_session + ): + """ + cacheEstimatorInput projects to still-needed columns and persists that + narrow frame once; independent sibling estimators must fit to byte-identical + params whether it is on or off, with the genuinely-unused column dropped. + """ + df = spark_session.createDataFrame( + [ + (1.0, 2.0, 3.0, 99.0), + (2.0, 4.0, 6.0, 99.0), + (3.0, 6.0, 9.0, 99.0), + (4.0, 8.0, 12.0, 99.0), + ], + ["x1", "x2", "x3", "junk"], + ) + + def build_pipeline(cache): + return KamaeSparkPipeline( + stages=[ + StandardScaleEstimator(inputCol="x1", outputCol="x1_scaled"), + StandardScaleEstimator(inputCol="x2", outputCol="x2_scaled"), + StandardScaleEstimator(inputCol="x3", outputCol="x3_scaled"), + ], + cacheEstimatorInput=cache, + ) + + cached_model = build_pipeline(cache=True).fit(df) + baseline_model = build_pipeline(cache=False).fit(df) + + for cached_scaler, baseline_scaler in zip( + cached_model.stages, baseline_model.stages + ): + assert cached_scaler.getMean() == baseline_scaler.getMean() + assert cached_scaler.getStddev() == baseline_scaler.getStddev() + + def test_spark_pipeline_cache_estimator_input_is_opt_in(self, spark_session): + """ + The narrow-cache projection (which computes the live keep-set via + collect_required_input_columns) must only run when cacheEstimatorInput is + True. Pruning is left off so the collector is not called for that reason. + """ + df = spark_session.createDataFrame( + [(1.0, 2.0, 99.0), (2.0, 4.0, 99.0), (3.0, 6.0, 99.0)], + ["x1", "x2", "junk"], + ) + stages = [ + StandardScaleEstimator(inputCol="x1", outputCol="x1_scaled"), + StandardScaleEstimator(inputCol="x2", outputCol="x2_scaled"), + ] + + with patch.object( + KamaeSparkPipeline, + "collect_required_input_columns", + wraps=KamaeSparkPipeline.collect_required_input_columns, + ) as mock_collect: + KamaeSparkPipeline(stages=stages).fit(df) + assert mock_collect.call_count == 0 + + mock_collect.reset_mock() + KamaeSparkPipeline(stages=stages, cacheEstimatorInput=True).fit(df) + assert mock_collect.call_count == 1 + + def test_spark_pipeline_cache_estimator_input_mutually_exclusive( + self, valid_stages_1, example_dataframe + ): + """ + cacheEstimatorInput and cacheIntermediateData are competing strategies; + enabling both warns, prefers cacheEstimatorInput, and still fits. + """ + pipeline = KamaeSparkPipeline( + stages=valid_stages_1, + cacheIntermediateData=True, + cacheEstimatorInput=True, + ) + with pytest.warns(UserWarning, match="takes precedence"): + pipeline_model = pipeline.fit(example_dataframe) + assert pipeline_model is not None + + def test_spark_pipeline_fit_sample_fraction_none_is_unchanged( + self, valid_stages_1, example_dataframe + ): + """ + With fitSampleFraction=None (the default), fit is unchanged: a smoke fit + succeeds and the resulting model transforms the full dataset. + """ + model = KamaeSparkPipeline(stages=valid_stages_1, fitSampleFraction=None).fit( + example_dataframe + ) + assert isinstance(model, KamaeSparkPipelineModel) + # The model still applies to the full dataset at transform time. + assert model.transform(example_dataframe).count() == example_dataframe.count() + + def test_spark_pipeline_fit_sample_fraction_matches_full_within_tolerance( + self, spark_session + ): + """ + fitSampleFraction fits the opted-in sample-robust estimators + (ConditionalStandardScale, with useFitSample=True) from a single + shared sample; on enough seeded data the fitted mean/stddev must stay within + a loose statistical tolerance of a full-data fit. + """ + from pyspark.sql import functions as F + + df = ( + spark_session.range(0, 40000) + .withColumn("x1", F.randn(seed=42)) + .withColumn("x2", 5.0 + 2.0 * F.randn(seed=7)) + .select("x1", "x2") + ).persist() + df.count() + + def build(fraction): + # useFitSample=True opts each estimator in to the shared sample. The + # full-data reference leaves it off so it is a genuine full fit. + opt_in = {"useFitSample": True} if fraction is not None else {} + return KamaeSparkPipeline( + stages=[ + ConditionalStandardScaleEstimator( + inputCol="x1", outputCol="x1_scaled", **opt_in + ), + ConditionalStandardScaleEstimator( + inputCol="x2", outputCol="x2_scaled", **opt_in + ), + ], + fitSampleFraction=fraction, + fitSampleSeed=13, + ) + + full_model = build(None).fit(df) + with pytest.warns(UserWarning): + sampled_model = build(0.2).fit(df) + + for full_stage, sampled_stage in zip(full_model.stages, sampled_model.stages): + assert abs(full_stage.getMean()[0] - sampled_stage.getMean()[0]) < 0.15 + assert abs(full_stage.getStddev()[0] - sampled_stage.getStddev()[0]) < 0.15 + df.unpersist() + + def test_spark_pipeline_fit_sample_fraction_scans_source_once(self, spark_session): + """ + fitSampleFraction materialises one shared sample, so the source is scanned + exactly once regardless of how many estimators fit from it - unlike the + default, where each independent estimator rescans the source. + """ + from pyspark.sql import functions as F + from pyspark.sql.types import DoubleType as SparkDoubleType + + n_rows = 500 + source = ( + spark_session.range(0, n_rows) + .select(F.col("id").cast("double").alias("raw")) + .persist() + ) + source.count() + + def counting_column(accumulator): + def _count(value): + accumulator.add(1) + return value + + udf = F.udf(_count, SparkDoubleType()).asNondeterministic() + return source.withColumn("x", udf(F.col("raw"))).drop("raw") + + def estimators(opt_in=False): + # useFitSample=True opts the estimator in to the shared sample. + kw = {"useFitSample": True} if opt_in else {} + return [ + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_a", **kw), + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_b", **kw), + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_c", **kw), + ] + + sampled_accum = spark_session.sparkContext.accumulator(0) + with pytest.warns(UserWarning): + KamaeSparkPipeline( + stages=estimators(opt_in=True), fitSampleFraction=1.0, fitSampleSeed=1 + ).fit(counting_column(sampled_accum)) + # One shared-sample materialisation, reused by all three opted-in estimators + # => exactly one pass over the source. + assert sampled_accum.value == n_rows + + baseline_accum = spark_session.sparkContext.accumulator(0) + KamaeSparkPipeline(stages=estimators()).fit(counting_column(baseline_accum)) + # Default: each of the three estimators rescans the source. + assert baseline_accum.value > n_rows + + source.unpersist() + + def test_spark_pipeline_fit_sample_fraction_disables_caching(self, spark_session): + """ + fitSampleFraction persists a tiny sample instead of the wide frame, so it is + incompatible with cacheEstimatorInput: enabling both must warn and disable + the cache (its narrow-projection path never runs). + """ + df = spark_session.createDataFrame([(1.0,), (2.0,), (3.0,), (4.0,)], ["x"]) + pipeline = KamaeSparkPipeline( + stages=[ + ConditionalStandardScaleEstimator( + inputCol="x", outputCol="x_scaled", useFitSample=True + ) + ], + cacheEstimatorInput=True, + fitSampleFraction=1.0, + fitSampleSeed=1, + ) + with patch.object( + KamaeSparkPipeline, + "collect_required_input_columns", + wraps=KamaeSparkPipeline.collect_required_input_columns, + ) as mock_collect: + with pytest.warns(UserWarning, match="incompatible"): + model = pipeline.fit(df) + # cacheEstimatorInput disabled => its keep-set collector is never called. + assert mock_collect.call_count == 0 + assert model is not None + + def test_spark_pipeline_fit_sample_fraction_non_opted_in_reads_full( + self, spark_session + ): + """ + Only estimators with useFitSample=True fit on the shared sample; + an estimator without it still fits on the full input. With fraction=1.0 the + sample is materialised once (n_rows) and reused by the opted-in estimator, + while the non-opted estimator triggers its own full scan => 2 * n_rows. + """ + from pyspark.sql import functions as F + from pyspark.sql.types import DoubleType as SparkDoubleType + + n_rows = 400 + source = ( + spark_session.range(0, n_rows) + .select(F.col("id").cast("double").alias("raw")) + .persist() + ) + source.count() + + accum = spark_session.sparkContext.accumulator(0) + + def _count(value): + accum.add(1) + return value + + udf = F.udf(_count, SparkDoubleType()).asNondeterministic() + counting = source.withColumn("x", udf(F.col("raw"))).drop("raw") + + with pytest.warns(UserWarning): + KamaeSparkPipeline( + stages=[ + # Opts in: reads the shared sample. + ConditionalStandardScaleEstimator( + inputCol="x", outputCol="x_sampled", useFitSample=True + ), + # No useFitSample: reads the full input. + ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_full"), + ], + fitSampleFraction=1.0, + fitSampleSeed=1, + ).fit(counting) + + # One materialisation of the shared sample (n_rows), reused by the opted-in + # estimator, plus one full scan by the non-opted estimator. + assert accum.value == 2 * n_rows + source.unpersist() + + def test_spark_pipeline_fit_sample_fraction_without_opt_in_is_noop( + self, spark_session + ): + """ + fitSampleFraction with no estimator opting in (none has useFitSample=True) + warns that it has no effect and fits exactly as a full fit would. + """ + df = spark_session.createDataFrame([(1.0,), (2.0,), (3.0,), (4.0,)], ["x"]) + stage = ConditionalStandardScaleEstimator(inputCol="x", outputCol="x_scaled") + + full_model = KamaeSparkPipeline(stages=[stage]).fit(df) + with pytest.warns(UserWarning, match="no effect"): + noop_model = KamaeSparkPipeline( + stages=[stage], fitSampleFraction=0.5, fitSampleSeed=1 + ).fit(df) + + assert noop_model.stages[0].getMean() == full_model.stages[0].getMean() + assert noop_model.stages[0].getStddev() == full_model.stages[0].getStddev() + + def test_spark_pipeline_fit_sample_fraction_overrides_estimator_sample_fraction( + self, spark_session + ): + """ + An estimator that sets both useFitSample=True and its own sampleFraction warns + that the shared pipeline sample wins and its sampleFraction is ignored. + """ + df = spark_session.createDataFrame([(1.0,), (2.0,), (3.0,), (4.0,)], ["x"]) + pipeline = KamaeSparkPipeline( + stages=[ + ConditionalStandardScaleEstimator( + inputCol="x", + outputCol="x_scaled", + useFitSample=True, + sampleFraction=0.5, + ) + ], + fitSampleFraction=1.0, + fitSampleSeed=1, + ) + with pytest.warns(UserWarning, match="sampleFraction is ignored"): + model = pipeline.fit(df) + assert model is not None + # The estimator's own sampleFraction is restored after the fit. + assert pipeline.getStages()[0].getSampleFraction() == 0.5 + + def test_spark_pipeline_fit_sample_fraction_opt_in_without_pipeline_sample_warns( + self, spark_session + ): + """ + useFitSample=True with no pipeline fitSampleFraction has no shared sample to + fit on, so it warns that useFitSample has no effect. + """ + df = spark_session.createDataFrame([(1.0,), (2.0,), (3.0,), (4.0,)], ["x"]) + pipeline = KamaeSparkPipeline( + stages=[ + ConditionalStandardScaleEstimator( + inputCol="x", outputCol="x_scaled", useFitSample=True + ) + ] + ) + with pytest.warns(UserWarning, match="useFitSample has no effect"): + model = pipeline.fit(df) + assert model is not None + @pytest.mark.parametrize( "stages, input_col, original_dtype", [ diff --git a/tests/kamae/spark/utils/__init__.py b/tests/kamae/spark/utils/__init__.py new file mode 100644 index 00000000..d47f0081 --- /dev/null +++ b/tests/kamae/spark/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright [2024] Expedia, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/kamae/spark/utils/test_transform_utils.py b/tests/kamae/spark/utils/test_transform_utils.py new file mode 100644 index 00000000..ba4a232f --- /dev/null +++ b/tests/kamae/spark/utils/test_transform_utils.py @@ -0,0 +1,45 @@ +# Copyright [2024] Expedia, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pyspark.sql.functions as F +from pyspark.sql.types import DoubleType, IntegerType, StructField, StructType + +from kamae.spark.utils.transform_utils import ( + single_input_single_output_scalar_udf_transform, +) + + +class TestTransformUtils: + def test_scalar_udf_transform_maps_numeric_null_to_none(self, spark_session): + """ + Spark NULLs in a nullable numeric column must reach the element func as Python + None, not Arrow's NaN. Otherwise `is None` null/OOV guards in the element funcs + (indexer/hash UDFs) silently misroute missing values. The func here would raise + on NaN (int(NaN)), so this fails if the vectorized path leaks NaN through. + """ + schema = StructType([StructField("x", DoubleType(), True)]) + df = spark_session.createDataFrame([(1.0,), (None,), (2.0,)], schema) + + out = df.withColumn( + "y", + single_input_single_output_scalar_udf_transform( + input_col=F.col("x"), + input_col_datatype=df.schema["x"].dataType, + func=lambda v: -1 if v is None else int(v), + udf_return_element_datatype=IntegerType(), + ), + ) + + result = sorted(row["y"] for row in out.collect()) + assert result == [-1, 1, 2] diff --git a/uv.lock b/uv.lock index 281ec5ba..1159c334 100644 --- a/uv.lock +++ b/uv.lock @@ -880,6 +880,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" }, @@ -928,6 +929,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" }, @@ -1843,6 +1845,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, upload-time = "2022-08-12T22:49:07.05Z" }, ] +[[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"