Skip to content
Closed
Show file tree
Hide file tree
Changes from 31 commits
Commits
Show all changes
62 commits
Select commit Hold shift + click to select a range
dad709c
Support LeafRunnableCommand as sub query
beliefer May 12, 2021
0338a23
Support LeafRunnableCommand as sub query
beliefer May 12, 2021
40cea6b
Support LeafRunnableCommand as sub query
beliefer May 12, 2021
a82ed76
Unify the behavior eagerly execute the commands
beliefer May 14, 2021
bf296fb
Update code
beliefer May 17, 2021
babe0d0
Merge branch 'master' into SPARK-35378
beliefer May 17, 2021
803a12a
Update code
beliefer May 17, 2021
4c9b3cf
Merge branch 'SPARK-35378' of github.com:beliefer/spark into SPARK-35378
beliefer May 17, 2021
e3b8454
Update code
beliefer May 17, 2021
6516cc4
Update code
beliefer May 17, 2021
ddbb5bb
Optimize code
beliefer May 18, 2021
bde1062
Optimize code
beliefer May 18, 2021
7a7bc55
Update code
beliefer May 18, 2021
2ad5391
Update code
beliefer May 19, 2021
33f0297
Optimize code
beliefer May 19, 2021
309fb0f
Update code
beliefer May 20, 2021
8e9277d
Update code
beliefer May 20, 2021
d006b2a
Update code
beliefer May 20, 2021
fc3afe3
Update code
beliefer May 20, 2021
fde3c31
Update code
beliefer May 21, 2021
b58c0ea
Update code
beliefer May 21, 2021
4dcc759
Update code
beliefer May 21, 2021
f47dda6
Update code
beliefer May 22, 2021
2a61f23
Update code
beliefer May 25, 2021
cd7d39c
Update code
beliefer May 25, 2021
6c74474
Merge branch 'master' into SPARK-35378
beliefer May 25, 2021
26c2341
Update code
beliefer May 25, 2021
d11b00d
Update code
beliefer May 26, 2021
a61c267
Update code
beliefer May 26, 2021
8fe6c06
Update code
beliefer May 26, 2021
b20b000
Update code
beliefer May 26, 2021
b60e04e
Update code
beliefer May 26, 2021
94ba930
Update code
beliefer May 26, 2021
0ed9485
Update code
beliefer May 26, 2021
911e081
Update code
beliefer May 26, 2021
ecfe9ba
update code
beliefer May 26, 2021
cfeadd1
Update code
beliefer May 27, 2021
6a80674
Update code
beliefer May 29, 2021
c81b082
Update code
beliefer May 29, 2021
ee1e84a
Update code
beliefer May 29, 2021
9c95570
Update code
beliefer May 31, 2021
219abb4
Merge branch 'SPARK-35378' of github.com:beliefer/spark into SPARK-35378
beliefer May 31, 2021
f39f920
Update code
beliefer May 31, 2021
1d10b61
Update code
beliefer May 31, 2021
2bcdddd
Update code
beliefer May 31, 2021
5d9f7ee
Update code
beliefer Jun 1, 2021
6011bbe
Update code
beliefer Jun 2, 2021
0905d84
Update code
beliefer Jun 2, 2021
3f6cb85
Update code
beliefer Jun 2, 2021
1d821e0
Update code
beliefer Jun 4, 2021
ddbc5c4
Update code
beliefer Jun 4, 2021
d545d9b
Update code
beliefer Jun 4, 2021
4b730d4
Update code
beliefer Jun 4, 2021
de55034
Update code
beliefer Jun 4, 2021
1a3ce51
Update code
beliefer Jun 5, 2021
ccf8ba3
Update code
beliefer Jun 7, 2021
35ea747
Merge branch 'master' into SPARK-35378
beliefer Jun 7, 2021
2b2caf7
Update code
beliefer Jun 7, 2021
1db52b9
Update code
beliefer Jun 7, 2021
83d2710
Update QueryExecution.scala
cloud-fan Jun 8, 2021
8054799
Update code
beliefer Jun 8, 2021
d15e166
Update code
beliefer Jun 8, 2021
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
11 changes: 1 addition & 10 deletions sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
Original file line number Diff line number Diff line change
Expand Up @@ -221,16 +221,7 @@ class Dataset[T] private[sql](
}

@transient private[sql] val logicalPlan: LogicalPlan = {
// For various commands (like DDL) and queries with side effects, we force query execution
// to happen right away to let these side effects take place eagerly.
val plan = queryExecution.analyzed match {
case c: Command =>
LocalRelation(c.output, withAction("command", queryExecution)(_.executeCollect()))
case u @ Union(children, _, _) if children.forall(_.isInstanceOf[Command]) =>
LocalRelation(u.output, withAction("command", queryExecution)(_.executeCollect()))
case _ =>
queryExecution.analyzed
}
val plan = queryExecution.commandExecuted
if (sparkSession.sessionState.conf.getConf(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED)) {
val dsIds = plan.getTagValue(Dataset.DATASET_ID_TAG).getOrElse(new HashSet[Long])
dsIds.add(id)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
/*
* 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.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute, UnsafeProjection}
import org.apache.spark.sql.catalyst.plans.QueryPlan
import org.apache.spark.sql.execution.metric.SQLMetrics

/**
* Physical plan node for holding data from a command.
*
* `rows` may not be serializable and ideally we should not send `rows` to the executors.
* Thus marking it as transient.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

it -> them

*/
case class CommandResultExec(
Comment thread
cloud-fan marked this conversation as resolved.
output: Seq[Attribute],

@yaooqinn yaooqinn May 26, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

override def output: Seq[Attribute] = commandPhysicalPlan.output?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Because we obtain the output in QueryExecution first, we not need to define it duplicate.

commandPhysicalPlan: SparkPlan,
@transient rows: Seq[InternalRow]) extends LeafExecNode {

override lazy val metrics = Map(
"numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"))

override def innerChildren: Seq[QueryPlan[_]] = Seq(commandPhysicalPlan)

@transient private lazy val unsafeRows: Array[InternalRow] = {
Comment thread
cloud-fan marked this conversation as resolved.
if (rows.isEmpty) {
Array.empty
} else {
val proj = UnsafeProjection.create(output, output)
rows.map(r => proj(r).copy()).toArray
}
}

@transient private lazy val rdd: RDD[InternalRow] = {
if (rows.isEmpty) {
sqlContext.sparkContext.emptyRDD
} else {
val numSlices = math.min(
unsafeRows.length, sqlContext.sparkSession.leafNodeDefaultParallelism)
sqlContext.sparkContext.parallelize(unsafeRows, numSlices)
}
}

override def doExecute(): RDD[InternalRow] = {
val numOutputRows = longMetric("numOutputRows")
rdd.map { r =>
numOutputRows += 1
r
}
}

override protected def stringArgs: Iterator[Any] = {
if (unsafeRows.isEmpty) {
Iterator("<empty>", output)
} else {
Iterator(output)
}
}

override def executeCollect(): Array[InternalRow] = {
longMetric("numOutputRows").add(unsafeRows.size)
unsafeRows
}

override def executeTake(limit: Int): Array[InternalRow] = {
val taken = unsafeRows.take(limit)
longMetric("numOutputRows").add(taken.size)
taken
}

override def executeTail(limit: Int): Array[InternalRow] = {
val taken: Seq[InternalRow] = unsafeRows.takeRight(limit)
longMetric("numOutputRows").add(taken.size)
taken.toArray
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ import org.apache.spark.sql.catalyst.analysis.UnsupportedOperationChecker
import org.apache.spark.sql.catalyst.expressions.SubqueryExpression
import org.apache.spark.sql.catalyst.expressions.codegen.ByteCodeStats
import org.apache.spark.sql.catalyst.plans.QueryPlan
import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, ReturnAnswer}
import org.apache.spark.sql.catalyst.plans.logical.{Command, LogicalPlan, ReturnAnswer}
import org.apache.spark.sql.catalyst.rules.{PlanChangeLogger, Rule}
import org.apache.spark.sql.catalyst.util.StringUtils.PlanStringConcat
import org.apache.spark.sql.catalyst.util.truncatedString
Expand All @@ -40,6 +40,7 @@ import org.apache.spark.sql.execution.bucketing.{CoalesceBucketsInJoin, DisableU
import org.apache.spark.sql.execution.dynamicpruning.PlanDynamicPruningFilters
import org.apache.spark.sql.execution.exchange.{EnsureRequirements, ReuseExchange}
import org.apache.spark.sql.execution.streaming.{IncrementalExecution, OffsetSeqMetadata}
import org.apache.spark.sql.expressions.CommandResult
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.streaming.OutputMode
import org.apache.spark.util.Utils
Expand All @@ -54,7 +55,8 @@ import org.apache.spark.util.Utils
class QueryExecution(
val sparkSession: SparkSession,
val logical: LogicalPlan,
val tracker: QueryPlanningTracker = new QueryPlanningTracker) extends Logging {
val tracker: QueryPlanningTracker = new QueryPlanningTracker,
val isExecutingCommand: Boolean = false) extends Logging {

val id: Long = QueryExecution.nextExecutionId

Expand All @@ -74,12 +76,33 @@ class QueryExecution(
sparkSession.sessionState.analyzer.executeAndCheck(logical, tracker)
}

// SPARK-35378: Commands should be executed eagerly so that `sql("INSERT ...")` can trigger the
// table insertion immediately without a `.collect()`. We also need to eagerly execute non-root
// commands, because many commands return `GenericInternalRow` and can't be put in a query plan
// directly, otherwise the query engine may cast `GenericInternalRow` to `UnsafeRow` and fail.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

now it seems better to put the doc in CommandExecutionMode

// SPARK-35378: Commands should be executed eagerly so that something like `sql("INSERT ...")`
// can trigger the table insertion immediately without a `.collect()`. To avoid end-less recursion we
// should use `NON_ROOT` when recursively executing commands. Note that we can't execute
// a query plan with leaf command nodes, because many commands return `GenericInternalRow`
// and can't be put in a query plan directly, otherwise the query engine may cast `GenericInternalRow`
// to `UnsafeRow` and fail. When running EXPLAIN, we should use `SKIP` to not trigger any execution.
object CommandExecutionMode extends Enumeration {
  val SKIP, NON_ROOT, ALL = Value
}

lazy val commandExecuted: LogicalPlan = if (isExecutingCommand) {
analyzed.mapChildren(eagerlyExecuteCommands)
} else {
eagerlyExecuteCommands(analyzed)
}

private def eagerlyExecuteCommands(p: LogicalPlan) = p transformDown {
case c: Command =>
val qe = sparkSession.sessionState.executePlan(c, true)
CommandResult(
qe.analyzed.output,
qe.commandExecuted,
qe.executedPlan,
SQLExecution.withNewExecutionId(qe, Some("command"))(qe.executedPlan.executeCollect()))
case other => other
}

lazy val withCachedData: LogicalPlan = sparkSession.withActive {
assertAnalyzed()
assertSupported()
// clone the plan to avoid sharing the plan instance between different stages like analyzing,
// optimizing and planning.
sparkSession.sharedState.cacheManager.useCachedData(analyzed.clone())
sparkSession.sharedState.cacheManager.useCachedData(commandExecuted.clone())
}

lazy val optimizedPlan: LogicalPlan = executePhase(QueryPlanningTracker.OPTIMIZATION) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

To make sure we don't include command execution time in the optimization phase, it's better to follow lazy val sparkPlan

lazy val optimizedPlan: LogicalPlan = {
  assertCommandExecuted()
  executePhase(QueryPlanningTracker.OPTIMIZATION) {
    ...
  }
}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

OK

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ import org.apache.spark.sql.execution.exchange.{REPARTITION, REPARTITION_WITH_NU
import org.apache.spark.sql.execution.python._
import org.apache.spark.sql.execution.streaming._
import org.apache.spark.sql.execution.streaming.sources.MemoryPlan
import org.apache.spark.sql.expressions.CommandResult
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.streaming.OutputMode
import org.apache.spark.sql.types.StructType
Expand Down Expand Up @@ -703,6 +704,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
execution.SampleExec(lb, ub, withReplacement, seed, planLater(child)) :: Nil
case logical.LocalRelation(output, data, _) =>
LocalTableScanExec(output, data) :: Nil
case CommandResult(output, _, plan, data) => CommandResultExec(output, plan, data) :: Nil
case logical.LocalLimit(IntegerLiteral(limit), child) =>
execution.LocalLimitExec(limit, planLater(child)) :: Nil
case logical.GlobalLimit(IntegerLiteral(limit), child) =>
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
/*
* 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.expressions

@cloud-fan cloud-fan Jun 15, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

This is a public package which makes CommandResult a public API. This is unexpected. We should move this class to org.apache.spark.sql.catalyst.plans.logical. @beliefer can you help to make this change?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

OK


import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.plans.QueryPlan
import org.apache.spark.sql.catalyst.plans.logical.{LeafNode, LogicalPlan}
import org.apache.spark.sql.execution.SparkPlan

/**
* Logical plan node for holding data from a command.
*
* `rows` may not be serializable and ideally we should not send `rows` to the executors.
* Thus marking it as transient.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

ditto, them

*/
case class CommandResult(
output: Seq[Attribute],
commandLogicalPlan: LogicalPlan,
commandPhysicalPlan: SparkPlan,
@transient rows: Seq[InternalRow]) extends LeafNode {
override def innerChildren: Seq[QueryPlan[_]] = Seq(commandLogicalPlan)
}
Original file line number Diff line number Diff line change
Expand Up @@ -310,8 +310,9 @@ abstract class BaseSessionStateBuilder(
/**
* Create a query execution object.
*/
protected def createQueryExecution: LogicalPlan => QueryExecution = { plan =>
new QueryExecution(session, plan)
protected def createQueryExecution: (LogicalPlan, Boolean) => QueryExecution = {
(plan, isExecutingCommand) =>
new QueryExecution(session, plan, isExecutingCommand = isExecutingCommand)
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ private[sql] class SessionState(
val streamingQueryManagerBuilder: () => StreamingQueryManager,
val listenerManager: ExecutionListenerManager,
resourceLoaderBuilder: () => SessionResourceLoader,
createQueryExecution: LogicalPlan => QueryExecution,
createQueryExecution: (LogicalPlan, Boolean) => QueryExecution,
createClone: (SparkSession, SessionState) => SessionState,
val columnarRules: Seq[ColumnarRule],
val queryStagePrepRules: Seq[Rule[SparkPlan]]) {
Expand Down Expand Up @@ -119,7 +119,8 @@ private[sql] class SessionState(
// Helper methods, partially leftover from pre-2.0 days
// ------------------------------------------------------

def executePlan(plan: LogicalPlan): QueryExecution = createQueryExecution(plan)
def executePlan(plan: LogicalPlan, isExecutingCommand: Boolean = false): QueryExecution =
createQueryExecution(plan, isExecutingCommand)
}

private[sql] object SessionState {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,12 @@ package org.apache.spark.sql.execution
import scala.io.Source

import org.apache.spark.sql.{AnalysisException, FastOperator}
import org.apache.spark.sql.catalyst.analysis.UnresolvedNamespace
import org.apache.spark.sql.catalyst.plans.QueryPlan
import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, OneRowRelation}
import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.trees.TreeNodeTag
import org.apache.spark.sql.execution.command.{ExecutedCommandExec, ShowTablesCommand}
import org.apache.spark.sql.expressions.CommandResult
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.util.Utils
Expand Down Expand Up @@ -236,4 +239,30 @@ class QueryExecutionSuite extends SharedSparkSession {
assert(df.queryExecution.optimizedPlan.toString.startsWith("Relation default.spark_34129["))
}
}

test("SPARK-35378: Eagerly execute non-root Command") {
def qe(logicalPlan: LogicalPlan): QueryExecution = new QueryExecution(spark, logicalPlan)

val showTables = ShowTables(UnresolvedNamespace(Seq.empty[String]), None)
val showTablesQe = qe(showTables)
assert(showTablesQe.commandExecuted.isInstanceOf[CommandResult])
assert(showTablesQe.executedPlan.isInstanceOf[CommandResultExec])
val showTablesResultExec = showTablesQe.executedPlan.asInstanceOf[CommandResultExec]
assert(showTablesResultExec.commandPhysicalPlan.isInstanceOf[ExecutedCommandExec])
assert(showTablesResultExec.commandPhysicalPlan.asInstanceOf[ExecutedCommandExec]
.cmd.isInstanceOf[ShowTablesCommand])

val project = Project(showTables.output, SubqueryAlias("s", showTables))
val projectQe = qe(project)
assert(projectQe.commandExecuted.isInstanceOf[Project])
assert(projectQe.commandExecuted.children.length == 1)
assert(projectQe.commandExecuted.children(0).isInstanceOf[SubqueryAlias])
assert(projectQe.commandExecuted.children(0).children.length == 1)
assert(projectQe.commandExecuted.children(0).children(0).isInstanceOf[CommandResult])
assert(projectQe.executedPlan.isInstanceOf[CommandResultExec])
val cmdResultExec = projectQe.executedPlan.asInstanceOf[CommandResultExec]
assert(cmdResultExec.commandPhysicalPlan.isInstanceOf[ExecutedCommandExec])
assert(cmdResultExec.commandPhysicalPlan.asInstanceOf[ExecutedCommandExec]
.cmd.isInstanceOf[ShowTablesCommand])
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -584,8 +584,10 @@ private[hive] class TestHiveSparkSession(

private[hive] class TestHiveQueryExecution(
sparkSession: TestHiveSparkSession,
logicalPlan: LogicalPlan)
extends QueryExecution(sparkSession, logicalPlan) with Logging {
logicalPlan: LogicalPlan,
isExecutingCommand: Boolean = false)
extends QueryExecution(sparkSession, logicalPlan, isExecutingCommand = isExecutingCommand)
with Logging {

def this(sparkSession: TestHiveSparkSession, sql: String) = {
this(sparkSession, sparkSession.sessionState.sqlParser.parsePlan(sql))
Expand Down Expand Up @@ -661,8 +663,10 @@ private[sql] class TestHiveSessionStateBuilder(

override def overrideConfs: Map[String, String] = TestHiveContext.overrideConfs

override def createQueryExecution: (LogicalPlan) => QueryExecution = { plan =>
new TestHiveQueryExecution(session.asInstanceOf[TestHiveSparkSession], plan)
override def createQueryExecution: (LogicalPlan, Boolean) => QueryExecution = {
(plan, isExecutingCommand) =>
new TestHiveQueryExecution(
session.asInstanceOf[TestHiveSparkSession], plan, isExecutingCommand)
}

override protected def newBuilder: NewBuilder = new TestHiveSessionStateBuilder(_, _)
Expand Down