now installing venv as part of sbt test

Author: minskya

spark-plugin/build.sbt

@@ -1,7 +1,12 @@
 import xerial.sbt.Sonatype._
 import sbtassembly.AssemblyPlugin.autoImport._
+import scala.sys.process._
 
 lazy val versionNum: String = "0.8.8"
+lazy val spark3Version: String = "3.5.1"
+lazy val pythonVenvDir: String = System.getProperty("java.io.tmpdir") + "/dataflint-pyspark-venv"
+lazy val pythonExec: String = "python3.11"
+val createPythonVenv = taskKey[Unit]("Create Python venv and install pyspark test dependencies")
 lazy val scala212 = "2.12.20"
 lazy val scala213 = "2.13.16"
 lazy val supportedScalaVersions = List(scala212, scala213)
@@ -35,8 +40,8 @@ lazy val plugin = (project in file("plugin"))
     } else {
       versionNum + "-SNAPSHOT"
     }),
-    libraryDependencies += "org.apache.spark" %% "spark-core" % "3.5.1" % "provided",
-    libraryDependencies += "org.apache.spark" %% "spark-sql" % "3.5.1"  % "provided",
+    libraryDependencies += "org.apache.spark" %% "spark-core" % spark3Version % "provided",
+    libraryDependencies += "org.apache.spark" %% "spark-sql" % spark3Version  % "provided",
     libraryDependencies +=  "com.amazonaws" % "aws-java-sdk-s3" % "1.12.470" % "provided",
     libraryDependencies += "org.apache.iceberg" %% "iceberg-spark-runtime-3.5" % "1.5.0" % "provided",
     libraryDependencies += "io.delta" %% "delta-spark" % "3.2.0" % "provided",
@@ -61,12 +66,12 @@ lazy val pluginspark3 = (project in file("pluginspark3"))
     } else {
       versionNum + "-SNAPSHOT"
     }),
-    libraryDependencies += "org.apache.spark" %% "spark-core" % "3.5.1" % "provided",
-    libraryDependencies += "org.apache.spark" %% "spark-sql" % "3.5.1"  % "provided",
+    libraryDependencies += "org.apache.spark" %% "spark-core" % spark3Version % "provided",
+    libraryDependencies += "org.apache.spark" %% "spark-sql" % spark3Version  % "provided",
     libraryDependencies +=  "com.amazonaws" % "aws-java-sdk-s3" % "1.12.470" % "provided",
     libraryDependencies += "org.apache.iceberg" %% "iceberg-spark-runtime-3.5" % "1.5.0" % "provided",
     libraryDependencies += "io.delta" %% "delta-spark" % "3.2.0" % "provided",
-    
+
     // Assembly configuration to create fat JAR with common code
     assembly / assemblyJarName := s"${name.value}_${scalaBinaryVersion.value}-${version.value}.jar",
     // Exclude Scala library from assembly - Spark provides its own Scala runtime
@@ -96,8 +101,8 @@ lazy val pluginspark3 = (project in file("pluginspark3"))
     Compile / unmanagedResourceDirectories += (plugin / Compile / resourceDirectory).value,
     libraryDependencies += "org.scalatest" %% "scalatest-funsuite"      % "3.2.17" % Test,
     libraryDependencies += "org.scalatest" %% "scalatest-shouldmatchers" % "3.2.17" % Test,
-    libraryDependencies += "org.apache.spark" %% "spark-core" % "3.5.1" % Test,
-    libraryDependencies += "org.apache.spark" %% "spark-sql"  % "3.5.1" % Test,
+    libraryDependencies += "org.apache.spark" %% "spark-core" % spark3Version % Test,
+    libraryDependencies += "org.apache.spark" %% "spark-sql"  % spark3Version % Test,
 
     // Include source and resources from plugin directory for tests
     Test / unmanagedSourceDirectories += (plugin / Compile / sourceDirectory).value / "scala",
@@ -107,6 +112,24 @@ lazy val pluginspark3 = (project in file("pluginspark3"))
     // Run test suites sequentially — parallel suites share the SparkSession via getOrCreate()
     // and one suite stopping the session causes NPEs in concurrently-running suites
     Test / parallelExecution := false,
+    createPythonVenv := {
+      val venvDir = new java.io.File(pythonVenvDir)
+      val log = streams.value.log
+      if (!venvDir.exists()) {
+        log.info(s"Creating Python venv at $pythonVenvDir ...")
+        val rc = Process(Seq(pythonExec, "-m", "venv", pythonVenvDir)).!
+        if (rc != 0) sys.error(s"Failed to create Python venv at $pythonVenvDir")
+      }
+      val pip = s"$pythonVenvDir/bin/pip"
+      log.info(s"Installing pyspark==$spark3Version pandas pyarrow into venv...")
+      val rc = Process(Seq(pip, "install", "--quiet", s"pyspark==$spark3Version", "pandas", "pyarrow")).!
+      if (rc != 0) sys.error("pip install failed")
+    },
+    Test / compile := (Test / compile).dependsOn(createPythonVenv).value,
+    Test / testOnly := (Test / testOnly).dependsOn(createPythonVenv).evaluated,
+    Test / javaOptions ++= Seq(
+      s"-Ddataflint.projectRoot=${baseDirectory.value.getParentFile.toString}",
+    ),
     Test / javaOptions ++= {
       // --add-opens is not supported on Java 8 (spec version starts with "1.")
       if (sys.props("java.specification.version").startsWith("1.")) Seq.empty

spark-plugin/pluginspark3/src/test/scala/org/apache/spark/dataflint/DataFlintPythonIntegrationSpec.scala

@@ -16,22 +16,13 @@ import java.nio.file.Paths
  * for the subprocess — no manual gateway setup needed. The Python script then connects
  * to this JVM via launch_gateway() and accesses the session through DataFlintStaticSession.
  *
- * Requires: .venv with pyspark, pandas, pyarrow installed.
- *   python3 -m venv .venv && .venv/bin/pip install pyspark pandas pyarrow
+ * Python dependencies (pyspark, pandas, pyarrow) are installed automatically into
+ * .venv at test startup if not already present.
  */
 class DataFlintPythonIntegrationSpec extends AnyFunSuite with Matchers with BeforeAndAfterAll with DataFlintTestHelper {
 
-  // pluginspark3 tests run with CWD = spark-plugin/pluginspark3/, so go up one level
-  // to reach the project root where .venv and pyspark-testing live.
-  private val projectRoot = Paths.get("").toAbsolutePath//.getParent
-
-  private val venvPython: String = {
-    val p = projectRoot.resolve(Paths.get(".venv", "bin", "python3"))
-    require(p.toFile.exists(),
-      s"Python venv not found at $p\n" +
-      "Run: python3 -m venv .venv && .venv/bin/pip install pyspark pandas pyarrow")
-    p.toString
-  }
+  private val projectRoot = Paths.get(sys.props.getOrElse("dataflint.projectRoot", ""))
+  private val venvPython: String = System.getProperty("java.io.tmpdir") + "/dataflint-pyspark-venv/bin/python3"
 
   private val scriptPath: String =
     projectRoot.resolve(

spark-plugin/pyspark-testing/dataflint_python_exec_integration_test.py

@@ -0,0 +1,85 @@
+"""
+Integration test helper: registers one DataFlint Python exec node scenario as a temp view.
+
+Usage: pass the test name as the first argument (sys.argv[1]):
+  batch_eval       — @udf               → registers batch_eval_view
+  arrow_eval       — @pandas_udf scalar  → registers arrow_eval_view
+  flat_map_groups  — applyInPandas       → registers flat_map_groups_view
+  flat_map_cogroups— cogroup.applyIn…    → registers flat_map_cogroups_view
+
+The Scala test (DataFlintPythonIntegrationSpec) calls PythonRunner.main with one of
+the above names, then checks the corresponding view's executedPlan.
+"""
+import sys
+import pyspark
+import pyspark.java_gateway
+from pyspark.sql import SparkSession
+from pyspark.sql.functions import udf, pandas_udf
+from pyspark.sql.types import LongType
+import pandas as pd
+
+test_name = sys.argv[1] if len(sys.argv) > 1 else ""
+
+gateway   = pyspark.java_gateway.launch_gateway()
+static    = gateway.jvm.org.apache.spark.dataflint.DataFlintStaticSession
+jsc       = static.javaSparkContext()
+spark_jvm = static.session()
+
+conf  = pyspark.conf.SparkConf(True, gateway.jvm, jsc.getConf())
+sc    = pyspark.SparkContext(gateway=gateway, jsc=jsc, conf=conf)
+spark = SparkSession(sc, jsparkSession=spark_jvm)
+
+df = spark.createDataFrame(
+    [(1, "a"), (2, "b"), (3, "a"), (4, "b")],
+    ["id", "cat"]
+)
+
+
+def register_batch_eval():
+    @udf(returnType=LongType())
+    def double_udf(x):
+        return x * 2
+    df.select(double_udf("id")).createOrReplaceTempView("batch_eval_view")
+
+
+def register_arrow_eval():
+    @pandas_udf(LongType())
+    def double_pandas_udf(s: pd.Series) -> pd.Series:
+        return s * 2
+    df.select(double_pandas_udf("id")).createOrReplaceTempView("arrow_eval_view")
+
+
+def register_flat_map_groups():
+    def identity_group(key, pdf):
+        return pdf
+    df.groupby("cat").applyInPandas(
+        identity_group, schema="id long, cat string"
+    ).createOrReplaceTempView("flat_map_groups_view")
+
+
+def register_flat_map_cogroups():
+    df2 = spark.createDataFrame(
+        [(1, "x"), (2, "y"), (3, "z"), (4, "w")],
+        ["id", "label"]
+    )
+    def cogroup_fn(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame:
+        left = left.copy()
+        left["label"] = right["label"].values[0] if len(right) > 0 else None
+        return left[["id", "cat", "label"]]
+    df.groupby("id").cogroup(df2.groupby("id")).applyInPandas(
+        cogroup_fn, schema="id long, cat string, label string"
+    ).createOrReplaceTempView("flat_map_cogroups_view")
+
+
+_tests = {
+    "batch_eval":        register_batch_eval,
+    "arrow_eval":        register_arrow_eval,
+    "flat_map_groups":   register_flat_map_groups,
+    "flat_map_cogroups": register_flat_map_cogroups,
+}
+
+if test_name not in _tests:
+    print(f"Unknown test '{test_name}'. Valid options: {list(_tests)}", file=sys.stderr)
+    sys.exit(1)
+
+_tests[test_name]()