feat: opt-in pipeline fit optimisations, native bucketize, and sampled fitting support with caching

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",

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 (
@@ -378,22 +379,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,

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,32 +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)

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,41 +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)]