From 9c9c1ae486b0e0227401b3a9e5e77dce18a0bed2 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Tue, 13 Feb 2024 19:03:57 -0800 Subject: [PATCH] Do not lazy-load pyspark.ml (#2184) * Do not lazy-load pyspark.ml Signed-off-by: Eduardo Apolinario * Lint Signed-off-by: Eduardo Apolinario --------- Signed-off-by: Eduardo Apolinario Co-authored-by: Eduardo Apolinario --- .../flytekitplugins/spark/pyspark_transformers.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/pyspark_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/pyspark_transformers.py index 04574a6591..bcbdd2b632 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/pyspark_transformers.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/pyspark_transformers.py @@ -1,12 +1,11 @@ from typing import Type -from flytekit import Blob, BlobMetadata, BlobType, FlyteContext, Literal, LiteralType, Scalar, lazy_module +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 -pyspark_ml = lazy_module("pyspark.ml") -PipelineModel = pyspark_ml.PipelineModel - class PySparkPipelineModelTransformer(TypeTransformer[PipelineModel]): _TYPE_INFO = BlobType(format="binary", dimensionality=BlobType.BlobDimensionality.MULTIPART)