[SPARK-10548] [SPARK-10563] [SQL] Fix concurrent SQL executions

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -27,9 +27,10 @@ import java.util.concurrent.atomic.{AtomicReference, AtomicBoolean, AtomicIntege
 import java.util.UUID.randomUUID
 
 import scala.collection.JavaConverters._
+import scala.collection.JavaConversions._
 import scala.collection.{Map, Set}
 import scala.collection.generic.Growable
-import scala.collection.mutable.HashMap
+import scala.collection.mutable.{HashMap, HashSet}
 import scala.reflect.{ClassTag, classTag}
 import scala.util.control.NonFatal
 
@@ -348,10 +349,27 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
 
   // Thread Local variable that can be used by users to pass information down the stack
   private val localProperties = new InheritableThreadLocal[Properties] {
-    override protected def childValue(parent: Properties): Properties = new Properties(parent)
+    override protected def childValue(parent: Properties): Properties = {
+      if (nonInheritedLocalProperties.nonEmpty) {
+        // If there are properties that should not be inherited, filter them out
+        val p = new Properties
+        val filtered = parent.filter { case (k, _) =>
+          !nonInheritedLocalProperties.contains(k)
+        }
+        p.putAll(filtered)
+        p
+      } else {
+        new Properties(parent)
+      }
+    }
     override protected def initialValue(): Properties = new Properties()
   }
 
+  /**
+   * Keys of local properties that should not be inherited by children threads.
+   */
+  private[spark] val nonInheritedLocalProperties: HashSet[String] = new HashSet[String]
+
   /* ------------------------------------------------------------------------------------- *
    | Initialization. This code initializes the context in a manner that is exception-safe. |
    | All internal fields holding state are initialized here, and any error prompts the     |

core/src/test/scala/org/apache/spark/ThreadingSuite.scala

@@ -154,6 +154,7 @@ class ThreadingSuite extends SparkFunSuite with LocalSparkContext with Logging {
     val threads = (1 to 5).map { i =>
       new Thread() {
         override def run() {
+          // TODO: these assertion failures don't actually fail the test...
           sc.setLocalProperty("test", i.toString)
           assert(sc.getLocalProperty("test") === i.toString)
           sem.release()
@@ -175,6 +176,7 @@ class ThreadingSuite extends SparkFunSuite with LocalSparkContext with Logging {
     val threads = (1 to 5).map { i =>
       new Thread() {
         override def run() {
+          // TODO: these assertion failures don't actually fail the test...
           assert(sc.getLocalProperty("test") === "parent")
           sc.setLocalProperty("test", i.toString)
           assert(sc.getLocalProperty("test") === i.toString)
@@ -190,6 +192,30 @@ class ThreadingSuite extends SparkFunSuite with LocalSparkContext with Logging {
     assert(sc.getLocalProperty("Foo") === null)
   }
 
+  test("inheritance exclusions (SPARK-10548)") {
+    sc = new SparkContext("local", "test")
+    sc.nonInheritedLocalProperties.add("do-not-inherit-me")
+    sc.setLocalProperty("do-inherit-me", "parent")
+    sc.setLocalProperty("do-not-inherit-me", "parent")
+    var throwable: Option[Throwable] = None
+    val threads = (1 to 5).map { i =>
+      new Thread() {
+        override def run() {
+          // only the ones we intend to inherit will be passed to the children
+          try {
+            assert(sc.getLocalProperty("do-inherit-me") === "parent")
+            assert(sc.getLocalProperty("do-not-inherit-me") === null)
+          } catch {
+            case t: Throwable => throwable = Some(t)
+          }
+        }
+      }
+    }
+    threads.foreach(_.start())
+    threads.foreach(_.join())
+    throwable.foreach { t => throw t }
+  }
+
   test("mutations to local properties should not affect submitted jobs (SPARK-6629)") {
     val jobStarted = new Semaphore(0)
     val jobEnded = new Semaphore(0)
@@ -210,6 +236,7 @@ class ThreadingSuite extends SparkFunSuite with LocalSparkContext with Logging {
     // Create a new thread which will inherit the current thread's properties
     val thread = new Thread() {
       override def run(): Unit = {
+        // TODO: these assertion failures don't actually fail the test...
         assert(sc.getLocalProperty(SparkContext.SPARK_JOB_GROUP_ID) === "originalJobGroupId")
         // Sleeps for a total of 10 seconds, but allows cancellation to interrupt the task
         try {

sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala

@@ -78,6 +78,11 @@ class SQLContext(@transient val sparkContext: SparkContext)
   sparkContext.addSparkListener(listener)
   sparkContext.ui.foreach(new SQLTab(this, _))
 
+  // Ensure query execution IDs are not inherited across the thread hierarchy, which is
+  // the default behavior for SparkContext local properties. Otherwise, we may confuse
+  // the listener as to which query is being executed. (SPARK-10548)
+  sparkContext.nonInheritedLocalProperties.add(SQLExecution.EXECUTION_ID_KEY)
+
   /**
    * Set Spark SQL configuration properties.
    *

sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala

@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.apache.spark.sql.test.SharedSQLContext
+
+class SQLExecutionSuite extends SharedSQLContext {
+  import testImplicits._
+
+  test("query execution IDs are not inherited across threads") {
+    sparkContext.setLocalProperty(SQLExecution.EXECUTION_ID_KEY, "123")
+    sparkContext.setLocalProperty("do-inherit-me", "some-value")
+    var throwable: Option[Throwable] = None
+    val thread = new Thread {
+      override def run(): Unit = {
+        try {
+          assert(sparkContext.getLocalProperty("do-inherit-me") === "some-value")
+          assert(sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) === null)
+        } catch {
+          case t: Throwable =>
+            throwable = Some(t)
+        }
+      }
+    }
+    thread.start()
+    thread.join()
+    throwable.foreach { t => throw t }
+  }
+
+  // This is the end-to-end version of the previous test.
+  test("parallel query execution (SPARK-10548)") {
+    (1 to 5).foreach { i =>
+      // Scala's parallel collections spawns new threads as children of the existing threads.
+      // We need to run this multiple times to ensure new threads are spawned. Without the fix
+      // for SPARK-10548, this usually fails on the second try.
+      val df = sparkContext.parallelize(1 to 5).map { i => (i, i) }.toDF("a", "b")
+      (1 to 10).par.foreach { _ => df.count() }
+    }
+  }
+}