From 1a6fec6a8122717d97eeb75f3f3bfcb5ff67b0e3 Mon Sep 17 00:00:00 2001 From: Alex Archambault Date: Wed, 17 May 2023 17:40:47 +0200 Subject: [PATCH] Add keepJar method on AmmoniteSparkSessionBuilder Handy to filter out specific JARs if needed --- .../AmmoniteSparkSessionBuilder.scala | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/modules/core/src/main/scala/org/apache/spark/sql/ammonitesparkinternals/AmmoniteSparkSessionBuilder.scala b/modules/core/src/main/scala/org/apache/spark/sql/ammonitesparkinternals/AmmoniteSparkSessionBuilder.scala index d78cf00..e043361 100644 --- a/modules/core/src/main/scala/org/apache/spark/sql/ammonitesparkinternals/AmmoniteSparkSessionBuilder.scala +++ b/modules/core/src/main/scala/org/apache/spark/sql/ammonitesparkinternals/AmmoniteSparkSessionBuilder.scala @@ -203,6 +203,7 @@ class AmmoniteSparkSessionBuilder(implicit private var sendSparkYarnJars0 = true private var sendSparkJars0 = true private var ignoreJars0 = Set.empty[URI] + private var keepJars0 = Seq.empty[URI => Boolean] def sendSparkYarnJars(force: Boolean = true): this.type = { sendSparkYarnJars0 = force @@ -219,6 +220,11 @@ class AmmoniteSparkSessionBuilder(implicit this } + def keepJars(keep: URI => Boolean): this.type = { + keepJars0 = keepJars0 :+ keep + this + } + var classServerOpt = Option.empty[AmmoniteClassServer] private def isYarn(): Boolean = @@ -349,7 +355,8 @@ class AmmoniteSparkSessionBuilder(implicit .map(normalize) .toSet val nonSparkJars = jars.filter(uri => !sparkJarFileSet.contains(normalize(uri))) - config("spark.jars", nonSparkJars.map(_.toASCIIString).mkString(",")) + val finalNonSparkJars = keepJars0.foldLeft(nonSparkJars)(_.filter(_)) + config("spark.jars", finalNonSparkJars.map(_.toASCIIString).mkString(",")) } if (interpApi != null)