This is an automated email from the ASF dual-hosted git repository.
gengliang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new a1aa200bdf3 [SPARK-38838][SQL] Refactor ResolveDefaultColumns.scala to
simplify helper methods
a1aa200bdf3 is described below
commit a1aa200bdf32e55ea3b1f220da882b29a7a2bf9b
Author: Daniel Tenedorio <daniel.tenedo...@databricks.com>
AuthorDate: Wed May 4 15:06:07 2022 +0800
[SPARK-38838][SQL] Refactor ResolveDefaultColumns.scala to simplify helper
methods
### What changes were proposed in this pull request?
Refactor ResolveDefaultColumns.scala to simplify helper methods:
1. Remove global variables entirely and pass all state through the method
calls instead.
2. Update method names to indicate which commands they are related to.
### Why are the changes needed?
This simplifies work for future DEFAULT column development.
### Does this PR introduce _any_ user-facing change?
No, it's just a refactoring change.
### How was this patch tested?
As a refactoring-only change, it is covered with existing test coverage.
Closes #36398 from dtenedor/refactor.
Authored-by: Daniel Tenedorio <daniel.tenedo...@databricks.com>
Signed-off-by: Gengliang Wang <gengli...@apache.org>
---
.../catalyst/analysis/ResolveDefaultColumns.scala | 340 +++++++++++++--------
1 file changed, 209 insertions(+), 131 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDefaultColumns.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDefaultColumns.scala
index 422a1e422be..ffbe18a7dfa 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDefaultColumns.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDefaultColumns.scala
@@ -17,6 +17,8 @@
package org.apache.spark.sql.catalyst.analysis
+import scala.collection.mutable
+
import org.apache.spark.sql.AnalysisException
import org.apache.spark.sql.catalyst.TableIdentifier
import org.apache.spark.sql.catalyst.catalog.{SessionCatalog,
UnresolvedCatalogRelation}
@@ -48,62 +50,112 @@ import org.apache.spark.sql.types._
* @param catalog the catalog to use for looking up the schema of INSERT INTO
table objects.
*/
case class ResolveDefaultColumns(
- analyzer: Analyzer,
- catalog: SessionCatalog) extends Rule[LogicalPlan] {
-
- // This field stores the enclosing INSERT INTO command, once we find one.
- var enclosingInsert: Option[InsertIntoStatement] = None
- // This field stores the schema of the target table of the above command.
- var insertTableSchemaWithoutPartitionColumns: Option[StructType] = None
-
+ analyzer: Analyzer,
+ catalog: SessionCatalog) extends Rule[LogicalPlan] {
override def apply(plan: LogicalPlan): LogicalPlan = {
- // Initialize by clearing our reference to the enclosing INSERT INTO
command.
- enclosingInsert = None
- insertTableSchemaWithoutPartitionColumns = None
- // Traverse the logical query plan in preorder (top-down).
plan.resolveOperatorsWithPruning(
(_ => SQLConf.get.enableDefaultColumns), ruleId) {
- case i@InsertIntoStatement(_, _, _, _, _, _)
- if i.query.collectFirst { case u: UnresolvedInlineTable
- if u.rows.nonEmpty && u.rows.forall(_.size == u.rows(0).size) => u
- }.isDefined =>
- enclosingInsert = Some(i)
- insertTableSchemaWithoutPartitionColumns =
getInsertTableSchemaWithoutPartitionColumns
- val regenerated: InsertIntoStatement = regenerateUserSpecifiedCols(i)
- regenerated
-
- case table: UnresolvedInlineTable
- if enclosingInsert.isDefined =>
- val expanded: UnresolvedInlineTable =
addMissingDefaultColumnValues(table).getOrElse(table)
- val replaced: LogicalPlan =
- replaceExplicitDefaultColumnValues(analyzer,
expanded).getOrElse(table)
- replaced
-
+ case i: InsertIntoStatement if insertsFromInlineTable(i) =>
+ resolveDefaultColumnsForInsertFromInlineTable(i)
case i@InsertIntoStatement(_, _, _, project: Project, _, _)
if !project.projectList.exists(_.isInstanceOf[Star]) =>
- enclosingInsert = Some(i)
- insertTableSchemaWithoutPartitionColumns =
getInsertTableSchemaWithoutPartitionColumns
- val expanded: Project =
addMissingDefaultColumnValues(project).getOrElse(project)
- val replaced: Option[LogicalPlan] =
replaceExplicitDefaultColumnValues(analyzer, expanded)
- val updated: InsertIntoStatement =
- if (replaced.isDefined) i.copy(query = replaced.get) else i
- val regenerated: InsertIntoStatement =
regenerateUserSpecifiedCols(updated)
- regenerated
+ resolveDefaultColumnsForInsertFromProject(i)
}
}
- // Helper method to regenerate user-specified columns of an
InsertIntoStatement based on the names
- // in the insertTableSchemaWithoutPartitionColumns field of this class.
- private def regenerateUserSpecifiedCols(i: InsertIntoStatement):
InsertIntoStatement = {
- if (i.userSpecifiedCols.nonEmpty &&
insertTableSchemaWithoutPartitionColumns.isDefined) {
+ /**
+ * Checks if a logical plan is an INSERT INTO command where the inserted
data comes from a VALUES
+ * list, with possible projection(s), aggregate(s), and/or alias(es) in
between.
+ */
+ private def insertsFromInlineTable(i: InsertIntoStatement): Boolean = {
+ var query = i.query
+ while (query.children.size == 1) {
+ query match {
+ case _: Project | _: Aggregate | _: SubqueryAlias =>
+ query = query.children(0)
+ case _ =>
+ return false
+ }
+ }
+ query match {
+ case u: UnresolvedInlineTable
+ if u.rows.nonEmpty && u.rows.forall(_.size == u.rows(0).size) =>
+ true
+ case _ =>
+ false
+ }
+ }
+
+ /**
+ * Resolves DEFAULT column references for an INSERT INTO command satisfying
the
+ * [[insertsFromInlineTable]] method.
+ */
+ private def resolveDefaultColumnsForInsertFromInlineTable(i:
InsertIntoStatement): LogicalPlan = {
+ val children = mutable.Buffer.empty[LogicalPlan]
+ var node = i.query
+ while (node.children.size == 1) {
+ children.append(node)
+ node = node.children(0)
+ }
+ val table = node.asInstanceOf[UnresolvedInlineTable]
+ val insertTableSchemaWithoutPartitionColumns: StructType =
+ getInsertTableSchemaWithoutPartitionColumns(i)
+ .getOrElse(return i)
+ val regenerated: InsertIntoStatement =
+ regenerateUserSpecifiedCols(i, insertTableSchemaWithoutPartitionColumns)
+ val expanded: UnresolvedInlineTable =
+ addMissingDefaultValuesForInsertFromInlineTable(
+ table, insertTableSchemaWithoutPartitionColumns)
+ val replaced: LogicalPlan =
+ replaceExplicitDefaultValuesForInputOfInsertInto(
+ analyzer, insertTableSchemaWithoutPartitionColumns, expanded)
+ .getOrElse(return i)
+ node = replaced
+ for (child <- children.reverse) {
+ node = child.withNewChildren(Seq(node))
+ }
+ regenerated.copy(query = node)
+ }
+
+ /**
+ * Resolves DEFAULT column references for an INSERT INTO command whose query
is a general
+ * projection.
+ */
+ private def resolveDefaultColumnsForInsertFromProject(i:
InsertIntoStatement): LogicalPlan = {
+ val insertTableSchemaWithoutPartitionColumns: StructType =
+ getInsertTableSchemaWithoutPartitionColumns(i)
+ .getOrElse(return i)
+ val regenerated: InsertIntoStatement =
+ regenerateUserSpecifiedCols(i, insertTableSchemaWithoutPartitionColumns)
+ val project: Project = i.query.asInstanceOf[Project]
+ val expanded: Project =
+ addMissingDefaultValuesForInsertFromProject(
+ project, insertTableSchemaWithoutPartitionColumns)
+ val replaced: LogicalPlan =
+ replaceExplicitDefaultValuesForInputOfInsertInto(
+ analyzer, insertTableSchemaWithoutPartitionColumns, expanded)
+ .getOrElse(return i)
+ regenerated.copy(query = replaced)
+ }
+
+ /**
+ * Regenerates user-specified columns of an InsertIntoStatement based on the
names in the
+ * insertTableSchemaWithoutPartitionColumns field of this class.
+ */
+ private def regenerateUserSpecifiedCols(
+ i: InsertIntoStatement,
+ insertTableSchemaWithoutPartitionColumns: StructType):
InsertIntoStatement = {
+ if (i.userSpecifiedCols.nonEmpty) {
i.copy(
- userSpecifiedCols =
insertTableSchemaWithoutPartitionColumns.get.fields.map(_.name))
+ userSpecifiedCols =
insertTableSchemaWithoutPartitionColumns.fields.map(_.name))
} else {
i
}
}
- // Helper method to check if an expression is an explicit DEFAULT column
reference.
+ /**
+ * Returns true if an expression is an explicit DEFAULT column reference.
+ */
private def isExplicitDefaultColumn(expr: Expression): Boolean = expr match {
case u: UnresolvedAttribute if
u.name.equalsIgnoreCase(CURRENT_DEFAULT_COLUMN_NAME) => true
case _ => false
@@ -112,51 +164,49 @@ case class ResolveDefaultColumns(
/**
* Updates an inline table to generate missing default column values.
*/
- private def addMissingDefaultColumnValues(
- table: UnresolvedInlineTable): Option[UnresolvedInlineTable] = {
- assert(enclosingInsert.isDefined)
+ private def addMissingDefaultValuesForInsertFromInlineTable(
+ table: UnresolvedInlineTable,
+ insertTableSchemaWithoutPartitionColumns: StructType):
UnresolvedInlineTable = {
val numQueryOutputs: Int = table.rows(0).size
- val schema = insertTableSchemaWithoutPartitionColumns.getOrElse(return
None)
- val newDefaultExpressions: Seq[Expression] =
getDefaultExpressions(numQueryOutputs, schema)
+ val schema = insertTableSchemaWithoutPartitionColumns
+ val newDefaultExpressions: Seq[Expression] =
+ getDefaultExpressionsForInsert(numQueryOutputs, schema)
val newNames: Seq[String] = schema.fields.drop(numQueryOutputs).map {
_.name }
- if (newDefaultExpressions.nonEmpty) {
- Some(table.copy(
- names = table.names ++ newNames,
- rows = table.rows.map { row => row ++ newDefaultExpressions }))
- } else {
- None
- }
+ table.copy(
+ names = table.names ++ newNames,
+ rows = table.rows.map { row => row ++ newDefaultExpressions })
}
/**
* Adds a new expressions to a projection to generate missing default column
values.
*/
- private def addMissingDefaultColumnValues(project: Project): Option[Project]
= {
+ private def addMissingDefaultValuesForInsertFromProject(
+ project: Project,
+ insertTableSchemaWithoutPartitionColumns: StructType): Project = {
val numQueryOutputs: Int = project.projectList.size
- val schema = insertTableSchemaWithoutPartitionColumns.getOrElse(return
None)
- val newDefaultExpressions: Seq[Expression] =
getDefaultExpressions(numQueryOutputs, schema)
+ val schema = insertTableSchemaWithoutPartitionColumns
+ val newDefaultExpressions: Seq[Expression] =
+ getDefaultExpressionsForInsert(numQueryOutputs, schema)
val newAliases: Seq[NamedExpression] =
newDefaultExpressions.zip(schema.fields).map {
case (expr, field) => Alias(expr, field.name)()
}
- if (newDefaultExpressions.nonEmpty) {
- Some(project.copy(projectList = project.projectList ++ newAliases))
- } else {
- None
- }
+ project.copy(projectList = project.projectList ++ newAliases)
}
/**
- * This is a helper for the addMissingDefaultColumnValues methods above.
+ * This is a helper for the addMissingDefaultValuesForInsertFromInlineTable
methods above.
*/
- private def getDefaultExpressions(numQueryOutputs: Int, schema: StructType):
Seq[Expression] = {
+ private def getDefaultExpressionsForInsert(
+ numQueryOutputs: Int,
+ schema: StructType): Seq[Expression] = {
val remainingFields: Seq[StructField] = schema.fields.drop(numQueryOutputs)
val numDefaultExpressionsToAdd =
getStructFieldsForDefaultExpressions(remainingFields).size
Seq.fill(numDefaultExpressionsToAdd)(UnresolvedAttribute(CURRENT_DEFAULT_COLUMN_NAME))
}
/**
- * This is a helper for the getDefaultExpressions methods above.
+ * This is a helper for the getDefaultExpressionsForInsert methods above.
*/
private def getStructFieldsForDefaultExpressions(fields: Seq[StructField]):
Seq[StructField] = {
if (SQLConf.get.useNullsForMissingDefaultColumnValues) {
@@ -167,14 +217,14 @@ case class ResolveDefaultColumns(
}
/**
- * Replaces unresolved DEFAULT column references with corresponding values
in a logical plan.
+ * Replaces unresolved DEFAULT column references with corresponding values
in an INSERT INTO
+ * command from a logical plan.
*/
- private def replaceExplicitDefaultColumnValues(
- analyzer: Analyzer,
- input: LogicalPlan): Option[LogicalPlan] = {
- assert(enclosingInsert.isDefined)
- val schema = insertTableSchemaWithoutPartitionColumns.getOrElse(return
None)
- val columnNames: Seq[String] = schema.fields.map { _.name }
+ private def replaceExplicitDefaultValuesForInputOfInsertInto(
+ analyzer: Analyzer,
+ insertTableSchemaWithoutPartitionColumns: StructType,
+ input: LogicalPlan): Option[LogicalPlan] = {
+ val schema = insertTableSchemaWithoutPartitionColumns
val defaultExpressions: Seq[Expression] = schema.fields.map {
case f if f.metadata.contains(CURRENT_DEFAULT_COLUMN_METADATA_KEY) =>
analyze(analyzer, f, "INSERT")
@@ -184,43 +234,40 @@ case class ResolveDefaultColumns(
// If necessary, return a more descriptive error message if the user tries
to nest the DEFAULT
// column reference inside some other expression, such as DEFAULT + 1
(this is not allowed).
//
- // Note that we don't need to check if
"SQLConf.get.useNullsForMissingDefaultColumnValues"
- // after this point because this method only takes responsibility to
replace *existing*
- // DEFAULT references. In contrast, the "getDefaultExpressions" method
will check that config
+ // Note that we don't need to check if
"SQLConf.get.useNullsForMissingDefaultColumnValues" after
+ // this point because this method only takes responsibility to replace
*existing* DEFAULT
+ // references. In contrast, the "getDefaultExpressionsForInsert" method
will check that config
// and add new NULLs if needed.
input match {
case table: UnresolvedInlineTable =>
- replaceExplicitDefaultColumnValues(defaultExpressions, table)
+ replaceExplicitDefaultValuesForInlineTable(defaultExpressions, table)
case project: Project =>
- replaceExplicitDefaultColumnValues(defaultExpressions, columnNames,
project)
+ replaceExplicitDefaultValuesForProject(defaultExpressions, project)
}
}
/**
* Replaces unresolved DEFAULT column references with corresponding values
in an inline table.
*/
- private def replaceExplicitDefaultColumnValues(
- defaultExpressions: Seq[Expression],
- table: UnresolvedInlineTable): Option[LogicalPlan] = {
+ private def replaceExplicitDefaultValuesForInlineTable(
+ defaultExpressions: Seq[Expression],
+ table: UnresolvedInlineTable): Option[LogicalPlan] = {
var replaced = false
- val newRows: Seq[Seq[Expression]] = {
+ val updated: Seq[Seq[Expression]] = {
table.rows.map { row: Seq[Expression] =>
for {
i <- 0 until row.size
expr = row(i)
defaultExpr = if (i < defaultExpressions.size) defaultExpressions(i)
else Literal(null)
- } yield expr match {
- case u: UnresolvedAttribute if isExplicitDefaultColumn(u) =>
- replaced = true
- defaultExpr
- case expr@_ if expr.find { isExplicitDefaultColumn }.isDefined =>
- throw new AnalysisException(DEFAULTS_IN_EXPRESSIONS_ERROR)
- case _ => expr
- }
+ } yield replaceExplicitDefaultReferenceInExpression(
+ expr, defaultExpr, DEFAULTS_IN_EXPRESSIONS_ERROR, false).map { e =>
+ replaced = true
+ e
+ }.getOrElse(expr)
}
}
if (replaced) {
- Some(table.copy(rows = newRows))
+ Some(table.copy(rows = updated))
} else {
None
}
@@ -229,28 +276,20 @@ case class ResolveDefaultColumns(
/**
* Replaces unresolved DEFAULT column references with corresponding values
in a projection.
*/
- private def replaceExplicitDefaultColumnValues(
- defaultExpressions: Seq[Expression],
- colNames: Seq[String],
- project: Project): Option[LogicalPlan] = {
+ private def replaceExplicitDefaultValuesForProject(
+ defaultExpressions: Seq[Expression],
+ project: Project): Option[LogicalPlan] = {
var replaced = false
val updated: Seq[NamedExpression] = {
for {
i <- 0 until project.projectList.size
projectExpr = project.projectList(i)
defaultExpr = if (i < defaultExpressions.size) defaultExpressions(i)
else Literal(null)
- colName = if (i < colNames.size) colNames(i) else ""
- } yield projectExpr match {
- case Alias(u: UnresolvedAttribute, _) if isExplicitDefaultColumn(u) =>
- replaced = true
- Alias(defaultExpr, colName)()
- case u: UnresolvedAttribute if isExplicitDefaultColumn(u) =>
- replaced = true
- Alias(defaultExpr, colName)()
- case expr@_ if expr.find { isExplicitDefaultColumn }.isDefined =>
- throw new AnalysisException(DEFAULTS_IN_EXPRESSIONS_ERROR)
- case _ => projectExpr
- }
+ } yield replaceExplicitDefaultReferenceInExpression(
+ projectExpr, defaultExpr, DEFAULTS_IN_EXPRESSIONS_ERROR, true).map { e
=>
+ replaced = true
+ e.asInstanceOf[NamedExpression]
+ }.getOrElse(projectExpr)
}
if (replaced) {
Some(project.copy(projectList = updated))
@@ -260,35 +299,48 @@ case class ResolveDefaultColumns(
}
/**
- * Looks up the schema for the table object of an INSERT INTO statement from
the catalog.
+ * Checks if a given input expression is an unresolved "DEFAULT" attribute
reference.
+ *
+ * @param input the input expression to examine.
+ * @param defaultExpr the default to return if [[input]] is an unresolved
"DEFAULT" reference.
+ * @param complexDefaultError error if [[input]] is a complex expression
with "DEFAULT" inside.
+ * @param addAlias if true, wraps the result with an alias of the original
default column name.
+ * @return [[defaultExpr]] if [[input]] is an unresolved "DEFAULT" attribute
reference.
*/
- private def getInsertTableSchemaWithoutPartitionColumns: Option[StructType]
= {
- assert(enclosingInsert.isDefined)
- val tableName = enclosingInsert.get.table match {
- case r: UnresolvedRelation => TableIdentifier(r.name)
- case r: UnresolvedCatalogRelation => r.tableMeta.identifier
- case _ => return None
- }
- // Lookup the relation from the catalog by name. This either succeeds or
returns some "not
- // found" error. In the latter cases, return out of this rule without
changing anything and let
- // the analyzer return a proper error message elsewhere.
- val lookup: LogicalPlan = try {
- catalog.lookupRelation(tableName)
- } catch {
- case _: AnalysisException => return None
- }
- val schema: StructType = lookup match {
- case SubqueryAlias(_, r: UnresolvedCatalogRelation) =>
- StructType(r.tableMeta.schema.fields.dropRight(
- enclosingInsert.get.partitionSpec.size))
- case SubqueryAlias(_, r: View) if r.isTempView =>
- StructType(r.schema.fields.dropRight(
- enclosingInsert.get.partitionSpec.size))
- case _ => return None
+ private def replaceExplicitDefaultReferenceInExpression(
+ input: Expression,
+ defaultExpr: Expression,
+ complexDefaultError: String,
+ addAlias: Boolean): Option[Expression] = {
+ input match {
+ case a@Alias(u: UnresolvedAttribute, _)
+ if isExplicitDefaultColumn(u) =>
+ Some(Alias(defaultExpr, a.name)())
+ case u: UnresolvedAttribute
+ if isExplicitDefaultColumn(u) =>
+ if (addAlias) {
+ Some(Alias(defaultExpr, u.name)())
+ } else {
+ Some(defaultExpr)
+ }
+ case expr@_
+ if expr.find(isExplicitDefaultColumn).isDefined =>
+ throw new AnalysisException(complexDefaultError)
+ case _ =>
+ None
}
+ }
+
+ /**
+ * Looks up the schema for the table object of an INSERT INTO statement from
the catalog.
+ */
+ private def getInsertTableSchemaWithoutPartitionColumns(
+ enclosingInsert: InsertIntoStatement): Option[StructType] = {
+ val target: StructType =
getSchemaForTargetTable(enclosingInsert.table).getOrElse(return None)
+ val schema: StructType =
StructType(target.fields.dropRight(enclosingInsert.partitionSpec.size))
// Rearrange the columns in the result schema to match the order of the
explicit column list,
// if any.
- val userSpecifiedCols: Seq[String] = enclosingInsert.get.userSpecifiedCols
+ val userSpecifiedCols: Seq[String] = enclosingInsert.userSpecifiedCols
if (userSpecifiedCols.isEmpty) {
return Some(schema)
}
@@ -311,4 +363,30 @@ case class ResolveDefaultColumns(
Some(StructType(userSpecifiedFields ++
getStructFieldsForDefaultExpressions(nonUserSpecifiedFields)))
}
+
+ /**
+ * Returns the schema for the target table of a DML command, looking into
the catalog if needed.
+ */
+ private def getSchemaForTargetTable(table: LogicalPlan): Option[StructType]
= {
+ // Lookup the relation from the catalog by name. This either succeeds or
returns some "not
+ // found" error. In the latter cases, return out of this rule without
changing anything and let
+ // the analyzer return a proper error message elsewhere.
+ val tableName: TableIdentifier = table match {
+ case r: UnresolvedRelation => TableIdentifier(r.name)
+ case r: UnresolvedCatalogRelation => r.tableMeta.identifier
+ case _ => return None
+ }
+ val lookup: LogicalPlan = try {
+ catalog.lookupRelation(tableName)
+ } catch {
+ case _: AnalysisException => return None
+ }
+ lookup match {
+ case SubqueryAlias(_, r: UnresolvedCatalogRelation) =>
+ Some(r.tableMeta.schema)
+ case SubqueryAlias(_, r: View) if r.isTempView =>
+ Some(r.schema)
+ case _ => None
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org
For additional commands, e-mail: commits-h...@spark.apache.org
