Add pyspark pipeline model transformer (#1101)

* add pyspark pipeline model transformer

* reset global spark session in `test_spark_task`

* use pytest fixture to reset spark session

Author: esadler-hbo

plugins/flytekit-spark/flytekitplugins/spark/__init__.py

@@ -17,6 +17,7 @@
 
 from flytekit.configuration import internal as _internal
 
+from .pyspark_transformers import PySparkPipelineModelTransformer
 from .schema import SparkDataFrameSchemaReader, SparkDataFrameSchemaWriter, SparkDataFrameTransformer  # noqa
 from .sd_transformers import ParquetToSparkDecodingHandler, SparkToParquetEncodingHandler
 from .task import Spark, new_spark_session  # noqa

plugins/flytekit-spark/flytekitplugins/spark/pyspark_transformers.py

@@ -0,0 +1,45 @@
+import pathlib
+from typing import Type
+
+from pyspark.ml import PipelineModel
+
+from flytekit import Blob, BlobMetadata, BlobType, FlyteContext, Literal, LiteralType, Scalar
+from flytekit.core.type_engine import TypeEngine
+from flytekit.extend import TypeTransformer
+
+
+class PySparkPipelineModelTransformer(TypeTransformer[PipelineModel]):
+    _TYPE_INFO = BlobType(format="binary", dimensionality=BlobType.BlobDimensionality.MULTIPART)
+
+    def __init__(self):
+        super(PySparkPipelineModelTransformer, self).__init__(name="PySparkPipelineModel", t=PipelineModel)
+
+    def get_literal_type(self, t: Type[PipelineModel]) -> LiteralType:
+        return LiteralType(blob=self._TYPE_INFO)
+
+    def to_literal(
+        self,
+        ctx: FlyteContext,
+        python_val: PipelineModel,
+        python_type: Type[PipelineModel],
+        expected: LiteralType,
+    ) -> Literal:
+        local_path = ctx.file_access.get_random_local_path()
+        pathlib.Path(local_path).parent.mkdir(parents=True, exist_ok=True)
+        python_val.save(local_path)
+
+        remote_dir = ctx.file_access.get_random_remote_directory()
+        ctx.file_access.upload_directory(local_path, remote_dir)
+
+        return Literal(scalar=Scalar(blob=Blob(uri=remote_dir, metadata=BlobMetadata(type=self._TYPE_INFO))))
+
+    def to_python_value(
+        self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[PipelineModel]
+    ) -> PipelineModel:
+        local_dir = ctx.file_access.get_random_local_directory()
+        ctx.file_access.download_directory(lv.scalar.blob.uri, local_dir)
+
+        return PipelineModel.load(local_dir)
+
+
+TypeEngine.register(PySparkPipelineModelTransformer())

plugins/flytekit-spark/tests/test_pyspark_transformers.py

@@ -0,0 +1,43 @@
+import pyspark
+from flytekitplugins.spark import PySparkPipelineModelTransformer
+from flytekitplugins.spark.task import Spark
+from pyspark.ml import Pipeline, PipelineModel
+from pyspark.ml.feature import Imputer
+
+import flytekit
+from flytekit import task, workflow
+from flytekit.core.type_engine import TypeEngine
+
+
+def test_type_resolution():
+    assert type(TypeEngine.get_transformer(PipelineModel)) == PySparkPipelineModelTransformer
+
+
+def test_pipeline_model_compatibility():
+    @task(task_config=Spark())
+    def my_dataset() -> pyspark.sql.DataFrame:
+        session = flytekit.current_context().spark_session
+        df = session.createDataFrame([("Megan", 2.0), ("Wayne", float("nan")), ("Dennis", 8.0)], ["name", "age"])
+        return df
+
+    @task(task_config=Spark())
+    def my_pipleline(df: pyspark.sql.DataFrame) -> PipelineModel:
+        imputer = Imputer(inputCols=["age"], outputCols=["imputed_age"])
+        pipeline = Pipeline(stages=[imputer]).fit(df)
+        return pipeline
+
+    @task(task_config=Spark())
+    def run_pipeline(df: pyspark.sql.DataFrame, pipeline: PipelineModel) -> int:
+        imputed_df = pipeline.transform(df)
+
+        return imputed_df.filter(imputed_df["imputed_age"].isNull()).count()
+
+    @workflow
+    def my_wf() -> int:
+        df = my_dataset()
+        pipeline = my_pipleline(df=df)
+
+        return run_pipeline(df=df, pipeline=pipeline)
+
+    res = my_wf()
+    assert res == 0

plugins/flytekit-spark/tests/test_spark_task.py

@@ -1,5 +1,6 @@
 import pandas as pd
 import pyspark
+import pytest
 from flytekitplugins.spark import Spark
 from flytekitplugins.spark.task import new_spark_session
 from pyspark.sql import SparkSession
@@ -10,7 +11,14 @@
 from flytekit.core.context_manager import ExecutionParameters, FlyteContextManager
 
 
-def test_spark_task():
+@pytest.fixture(scope="function")
+def reset_spark_session() -> None:
+    pyspark.sql.SparkSession.builder.getOrCreate().stop()
+    yield
+    pyspark.sql.SparkSession.builder.getOrCreate().stop()
+
+
+def test_spark_task(reset_spark_session):
     @task(task_config=Spark(spark_conf={"spark": "1"}))
     def my_spark(a: str) -> int:
         session = flytekit.current_context().spark_session