Skip to content

[SPARK-33273][SQL] Fix the race condition issue of subquery execution and hash aggregation. #30766

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
sarutak wants to merge 1 commit into from
Closed

[SPARK-33273][SQL] Fix the race condition issue of subquery execution and hash aggregation. #30766

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
15 changes: 10 additions & 5 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,11 @@ case class HashAggregateExec(

// The variables are used as aggregation buffers and each aggregate function has one or more
// ExprCode to initialize its buffer slots. Only used for aggregation without keys.
private var bufVars: Seq[Seq[ExprCode]] = _
private val bufVar = new ThreadLocal[Seq[Seq[ExprCode]]] {
Copy link
Contributor

@cloud-fan cloud-fan Dec 14, 2020

Choose a reason for hiding this comment

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

I thought about this too, but my concern is that we never promise the codegen to be thread-safe. We may have more places to fix, and we may break it again in the future.

override def initialValue(): Seq[Seq[ExprCode]] = {
Seq.empty
}
}

private def doProduceWithoutKeys(ctx: CodegenContext): String = {
val initAgg = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "initAgg")
Expand All @@ -169,7 +173,7 @@ case class HashAggregateExec(
// generate variables for aggregation buffer
val functions = aggregateExpressions.map(_.aggregateFunction.asInstanceOf[DeclarativeAggregate])
val initExpr = functions.map(f => f.initialValues)
bufVars = initExpr.map { exprs =>
val localBufVars = initExpr.map { exprs =>
exprs.map { e =>
val isNull = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "bufIsNull")
val value = ctx.addMutableState(CodeGenerator.javaType(e.dataType), "bufValue")
Expand All @@ -185,7 +189,8 @@ case class HashAggregateExec(
JavaCode.global(value, e.dataType))
}
}
val flatBufVars = bufVars.flatten
bufVars.set(localBufVars)
val flatBufVars = bufVars.get().flatten
val initBufVar = evaluateVariables(flatBufVars)

// generate variables for output
Expand Down Expand Up @@ -322,7 +327,7 @@ case class HashAggregateExec(
e.aggregateFunction.asInstanceOf[DeclarativeAggregate].mergeExpressions
}
}
ctx.currentVars = bufVars.flatten ++ input
ctx.currentVars = bufVars.get().flatten ++ input
val boundUpdateExprs = updateExprs.map { updateExprsForOneFunc =>
bindReferences(updateExprsForOneFunc, inputAttrs)
}
Expand All @@ -336,7 +341,7 @@ case class HashAggregateExec(

val aggNames = functions.map(_.prettyName)
val aggCodeBlocks = bufferEvals.zipWithIndex.map { case (bufferEvalsForOneFunc, i) =>
val bufVarsForOneFunc = bufVars(i)
val bufVarsForOneFunc = bufVars.get()(i)
// All the update code for aggregation buffers should be placed in the end
// of each aggregation function code.
val updates = bufferEvalsForOneFunc.zip(bufVarsForOneFunc).map { case (ev, bufVar) =>
Expand Down