Do not lazy-load pyspark.ml (#2184)

* Do not lazy-load pyspark.ml Signed-off-by: Eduardo Apolinario <[email protected]> * Lint Signed-off-by: Eduardo Apolinario <[email protected]> --------- Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]>
flyteorg · Feb 14, 2024 · 9c9c1ae · 9c9c1ae
1 parent 6fc774b
commit 9c9c1ae
Showing 1 changed file with 3 additions and 4 deletions.
diff --git 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)