Import pyspark.core module later for Spark Classic only

HyukjinKwon · Feb 6, 2025 · 372bef2 · 372bef2
1 parent e2ef5a4
commit 372bef2
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/python/pyspark/ml/classification.py b/python/pyspark/ml/classification.py
@@ -3810,11 +3810,12 @@ def setRawPredictionCol(self, value: str) -> "OneVsRestModel":
 
     def __init__(self, models: List[ClassificationModel]):
         super(OneVsRestModel, self).__init__()
-        from pyspark.core.context import SparkContext
-
         self.models = models
         if is_remote() or not isinstance(models[0], JavaMLWritable):
             return
+
+        from pyspark.core.context import SparkContext
+
         # set java instance
         java_models = [cast(_JavaClassificationModel, model)._to_java() for model in self.models]
         sc = SparkContext._active_spark_context