WIP

microsoft · Jun 22, 2021 · f4421b3 · f4421b3
1 parent 53698f4
commit f4421b3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/core/src/test/scala/com/microsoft/ml/spark/codegen/TestGen.scala b/core/src/test/scala/com/microsoft/ml/spark/codegen/TestGen.scala
@@ -54,11 +54,12 @@ object TestGen {
          |from pyspark.sql import SparkSession, SQLContext
          |import os
          |import mmlspark
+         |from mmlspark.core import __spark_package_version__
          |
          |spark = (SparkSession.builder
          |    .master("local[*]")
          |    .appName("PysparkTests")
-         |    .config("spark.jars.packages", "com.microsoft.ml.spark:mmlspark:" + mmlspark.core.__spark_package_version__)
+         |    .config("spark.jars.packages", "com.microsoft.ml.spark:mmlspark:" + __spark_package_version__)
          |    .config("spark.jars.repositories", "https://mmlspark.azureedge.net/maven")
          |    .config("spark.executor.heartbeatInterval", "60s")
          |    .config("spark.sql.shuffle.partitions", 10)