Skip to content

[SPARK-26493][SQL] Allow multiple spark.sql.extensions #23398

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jamisonbennett wants to merge 7 commits into from
Closed

[SPARK-26493][SQL] Allow multiple spark.sql.extensions #23398

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
cef8eb0
[SPARK-26493][SQL] Allow multiple spark.sql.extensions
jamisonbennett Dec 28, 2018
689a4d2
Address comments from dongjoon-hyun, hvanhovell, beliefer, gatorsmile
jamisonbennett Dec 29, 2018
d89cfd9
Address comment from hvanhovell
jamisonbennett Dec 31, 2018
fb4ad34
Address comments from HyukjinKwon
jamisonbennett Jan 5, 2019
65a5f3f
Address comments from HyukjinKwon, vanzin
jamisonbennett Jan 9, 2019
9c0181d
Fix a typo
jamisonbennett Jan 9, 2019
deaf73e
Address comments from vanzin
jamisonbennett Jan 9, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,9 +99,10 @@ object StaticSQLConf {
.createWithDefault(false)

val SPARK_SESSION_EXTENSIONS = buildStaticConf("spark.sql.extensions")
.doc("Name of the class used to configure Spark Session extensions. The class should " +
.doc("List of the class names used to configure Spark Session extensions. The classes should " +
Copy link
Contributor

@hvanhovell hvanhovell Dec 28, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please document in what order rules from multiple extensions are executed, how listeners are invoked and how functions are registered (last one wins?). Be sure to cover what happens if you add duplicate listeners/rules/functions etc...

Copy link
Contributor

@beliefer beliefer Dec 29, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A suggestion of update the comment is replace 'List of the class names used to configure Spark Session extensions. The classes should implement Function1[SparkSessionExtension, Unit], and must have a no-args constructor.' to 'A comma-separated list of classes that implement Function1[SparkSessionExtension, Unit] used to configure Spark Session extensions. The classes must have a no-args constructor.'

Copy link
Contributor Author

@jamisonbennett jamisonbennett Dec 30, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I updated the documentation as suggested with respect to the ordering. I think the comment about "listeners" is related to code that is close in proximity to the code for this pull request, but it is a different configuration item. So I didn't update the spark.sql.queryExecutionListeners documentation. If you think that documentation should be updated, let me know if I should include it as a part of this pull request or as a separate pull request.

"implement Function1[SparkSessionExtension, Unit], and must have a no-args constructor.")
.stringConf
.toSequence
.createOptional

val QUERY_EXECUTION_LISTENERS = buildStaticConf("spark.sql.queryExecutionListeners")
Expand Down
11 changes: 5 additions & 6 deletions sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ class SparkSession private(
private[sql] def this(sc: SparkContext) {
this(sc, None, None,
SparkSession.applyExtensions(
sc.getConf.get(StaticSQLConf.SPARK_SESSION_EXTENSIONS),
sc.getConf.get(StaticSQLConf.SPARK_SESSION_EXTENSIONS).getOrElse(Seq.empty),
new SparkSessionExtensions))
}

Expand Down Expand Up @@ -950,7 +950,7 @@ object SparkSession extends Logging {
}

applyExtensions(
sparkContext.getConf.get(StaticSQLConf.SPARK_SESSION_EXTENSIONS),
sparkContext.getConf.get(StaticSQLConf.SPARK_SESSION_EXTENSIONS).getOrElse(Seq.empty),
extensions)

session = new SparkSession(sparkContext, None, None, extensions)
Expand Down Expand Up @@ -1138,14 +1138,13 @@ object SparkSession extends Logging {
}

/**
* Initialize extensions for given extension classname. This class will be applied to the
* Initialize extensions for given extension classnames. The classes will be applied to the
* extensions passed into this function.
*/
private def applyExtensions(
extensionOption: Option[String],
extensionConfClassNames: Seq[String],
extensions: SparkSessionExtensions): SparkSessionExtensions = {
if (extensionOption.isDefined) {
val extensionConfClassName = extensionOption.get
extensionConfClassNames.foreach { extensionConfClassName =>
try {
val extensionConfClass = Utils.classForName(extensionConfClassName)
val extensionConf = extensionConfClass.getConstructor().newInstance()
Expand Down
59 changes: 51 additions & 8 deletions sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,41 +30,43 @@ import org.apache.spark.sql.types.{DataType, IntegerType, StructType}
*/
class SparkSessionExtensionSuite extends SparkFunSuite {
type ExtensionsBuilder = SparkSessionExtensions => Unit
private def create(builder: ExtensionsBuilder): ExtensionsBuilder = builder
private def create(builder: ExtensionsBuilder): Seq[ExtensionsBuilder] = Seq(builder)

private def stop(spark: SparkSession): Unit = {
spark.stop()
SparkSession.clearActiveSession()
SparkSession.clearDefaultSession()
}

private def withSession(builder: ExtensionsBuilder)(f: SparkSession => Unit): Unit = {
val spark = SparkSession.builder().master("local[1]").withExtensions(builder).getOrCreate()
private def withSession(builders: Seq[ExtensionsBuilder])(f: SparkSession => Unit): Unit = {
val builder = SparkSession.builder().master("local[1]")
builders.foreach(builder.withExtensions)
val spark = builder.getOrCreate()
try f(spark) finally {
stop(spark)
}
}

test("inject analyzer rule") {
withSession(_.injectResolutionRule(MyRule)) { session =>
withSession(Seq(_.injectResolutionRule(MyRule))) { session =>
assert(session.sessionState.analyzer.extendedResolutionRules.contains(MyRule(session)))
}
}

test("inject check analysis rule") {
withSession(_.injectCheckRule(MyCheckRule)) { session =>
withSession(Seq(_.injectCheckRule(MyCheckRule))) { session =>
assert(session.sessionState.analyzer.extendedCheckRules.contains(MyCheckRule(session)))
}
}

test("inject optimizer rule") {
withSession(_.injectOptimizerRule(MyRule)) { session =>
withSession(Seq(_.injectOptimizerRule(MyRule))) { session =>
assert(session.sessionState.optimizer.batches.flatMap(_.rules).contains(MyRule(session)))
}
}

test("inject spark planner strategy") {
withSession(_.injectPlannerStrategy(MySparkStrategy)) { session =>
withSession(Seq(_.injectPlannerStrategy(MySparkStrategy))) { session =>
assert(session.sessionState.planner.strategies.contains(MySparkStrategy(session)))
}
}
Expand All @@ -78,6 +80,14 @@ class SparkSessionExtensionSuite extends SparkFunSuite {
}
}

test("inject multiple rules") {
withSession(Seq(_.injectOptimizerRule(MyRule),
_.injectPlannerStrategy(MySparkStrategy))) { session =>
assert(session.sessionState.optimizer.batches.flatMap(_.rules).contains(MyRule(session)))
assert(session.sessionState.planner.strategies.contains(MySparkStrategy(session)))
}
}

test("inject stacked parsers") {
val extension = create { extensions =>
extensions.injectParser((_, _) => CatalystSqlParser)
Expand Down Expand Up @@ -114,6 +124,25 @@ class SparkSessionExtensionSuite extends SparkFunSuite {
stop(session)
}
}

test("use multiple custom class for extensions") {
val session = SparkSession.builder()
.master("local[1]")
.config("spark.sql.extensions", Seq(
classOf[MyExtensions].getCanonicalName,
classOf[MyExtensions2].getCanonicalName).mkString(","))
.getOrCreate()
try {
assert(session.sessionState.planner.strategies.contains(MySparkStrategy(session)))
assert(session.sessionState.analyzer.extendedResolutionRules.contains(MyRule(session)))
assert(session.sessionState.functionRegistry
.lookupFunction(MyExtensions.myFunction._1).isDefined)
assert(session.sessionState.functionRegistry
.lookupFunction(MyExtensions2.myFunction._1).isDefined)
Copy link
Member

@dongjoon-hyun dongjoon-hyun Dec 28, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, now we have multiple extension registrations. The order of extension names might have side-effects.

  • Can we have a test case for duplicated extension names? MyExtension2 and MyExtension2?
  • Can we have a negative test case for function name conflicts? MyExtension2.myFunction and MyExtension3.myFunction?

Copy link
Member

@gatorsmile gatorsmile Dec 28, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the order matters, but we need to discuss and document the behavior when we have name conflicts.

For example, the same rule will be added twice in extendedResolutionRules. Is it desired?

class MyExtensions extends (SparkSessionExtensions => Unit) {
  def apply(e: SparkSessionExtensions): Unit = {
    e.injectResolutionRule(MyRule)
  }
}

class MyExtensions2 extends (SparkSessionExtensions => Unit) {
  def apply(e: SparkSessionExtensions): Unit = {
    e.injectResolutionRule(MyRule)
  }
}

Copy link
Member

@dongjoon-hyun dongjoon-hyun Dec 28, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep. If there is no reason to allow that, we had better disallow that by design before this PR.

Copy link
Contributor

@hvanhovell hvanhovell Dec 28, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are use cases where you want to execute rules in a certain order. So I think it is reasonable to add the same rule multiple times. If you want more control you could even create 'micro' optimizer batches by calling multiple rules from one rule.

I think this is more a matter of proper documentation than one where we should explicitly block things. Also note that this is a pretty advanced feature and by this stage users are expected to know what they are doing.

Copy link
Contributor Author

@jamisonbennett jamisonbennett Dec 28, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Prior to this change, it was possible to programmatically register multiple extensions but it was not possible to do so through the spark.sql.extensions configuration. Although it wasn't documented/tested until this pull request. E.g. The following works without this pull request:

SparkSession.builder()
  .master("..")
  .withExtensions(sparkSessionExtensions1)
  .withExtensions(sparkSessionExtensions2)
  .getOrCreate()

So I think conflicting function names are already currently possible (but not documented). In the following cases:

  1. Conflicting function names are registered by calling .withExtenions() multiple times
  2. An extension accidentally registers a function that was already registered with the builtin functions
  3. An extension accidentally registers a function multiple times by calling injectFunction(myFunction)

As for the order, it looks to me like the last function to be stored with conflicting names is the one which is retrieved:

class SimpleFunctionRegistry extends FunctionRegistry {

  @GuardedBy("this")
  private val functionBuilders =
    new mutable.HashMap[FunctionIdentifier, (ExpressionInfo, FunctionBuilder)]
  override def registerFunction(
      name: FunctionIdentifier,
      info: ExpressionInfo,
      builder: FunctionBuilder): Unit = synchronized {
    functionBuilders.put(normalizeFuncName(name), (info, builder))
  }

I will update this PR to document what happens in order of operations and conflicts. If we need to explicitly block duplicates functions from being registered, I can temporarily drop this PR and see about making those changes first.

Copy link
Member

@gatorsmile gatorsmile Dec 29, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for explaining it. We do not need to block it, but we might need to detect and throw a warning message at least.

More importantly, we need to document the current behavior and also add a test case to ensure the future changes will not break it. In the future, we can revisit the current behavior and make a change if needed.

Copy link
Contributor Author

@jamisonbennett jamisonbennett Dec 30, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we have a test case for duplicated extension names? Done
Can we have a negative test case for function name conflicts? MyExtension2.myFunction and MyExtension3.myFunction? Done

I added documentation for the behavior.
I added a warning message if a registered function is replaced.
I added a test case for the ordering.

} finally {
stop(session)
}
}
}

case class MyRule(spark: SparkSession) extends Rule[LogicalPlan] {
Expand Down Expand Up @@ -151,8 +180,9 @@ case class MyParser(spark: SparkSession, delegate: ParserInterface) extends Pars
object MyExtensions {

val myFunction = (FunctionIdentifier("myFunction"),
new ExpressionInfo("noClass", "myDb", "myFunction", "usage", "extended usage" ),
new ExpressionInfo("noClass", "myDb", "myFunction", "usage", "extended usage"),
(myArgs: Seq[Expression]) => Literal(5, IntegerType))
Copy link
Member

@HyukjinKwon HyukjinKwon Jan 9, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: (_: Seq[Expression]) => Literal(5, IntegerType))

Copy link
Contributor Author

@jamisonbennett jamisonbennett Jan 9, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done


Copy link
Member

@HyukjinKwon HyukjinKwon Jan 9, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would remove this newline. Looks unrelated.

Copy link
Contributor Author

@jamisonbennett jamisonbennett Jan 9, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

}

class MyExtensions extends (SparkSessionExtensions => Unit) {
Expand All @@ -162,3 +192,16 @@ class MyExtensions extends (SparkSessionExtensions => Unit) {
e.injectFunction(MyExtensions.myFunction)
}
}

object MyExtensions2 {

val myFunction = (FunctionIdentifier("myFunction2"),
new ExpressionInfo("noClass", "myDb", "myFunction2", "usage", "extended usage" ),
Copy link
Member

@HyukjinKwon HyukjinKwon Jan 9, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: " ) -> ")

Copy link
Contributor Author

@jamisonbennett jamisonbennett Jan 9, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

(myArgs: Seq[Expression]) => Literal(5, IntegerType))
}

class MyExtensions2 extends (SparkSessionExtensions => Unit) {
def apply(e: SparkSessionExtensions): Unit = {
e.injectFunction(MyExtensions2.myFunction)
}
}