You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "Kimahriman (via GitHub)" <gi...@apache.org> on 2023/06/23 00:33:49 UTC
[GitHub] [spark] Kimahriman opened a new pull request, #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Kimahriman opened a new pull request, #34558:
URL: https://github.com/apache/spark/pull/34558
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'core/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
This PR adds codegen support to array based higher order functions except ArraySort. This is my first time playing around with codegen, so definitely looking for any feedback.
A few notes:
- Disabled subexpression elimination for lambda functions (this already was the case because it was CodegenFallback). I plan to explore supprting subexpression elimination inside lambda functions later on, as it will require special handling.
- I set the AtomicReference for all lambda values as well in case a child expression reverts to interpreted evaluation for any reason (CodegenFallback or otherwise)
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
To improve performance of array higher-order function operations, letting the children be codegen'd and participate in WholeStageCodegen
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as the documentation fix.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
No, only performance improvements.
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
-->
Existing unit tests, let me know if there's other codegen-specific unit tests I should add.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3364225454
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -886,6 +1179,115 @@ case class ArrayAggregate(
}
}
+ protected def nullSafeCodeGen(
+ ctx: CodegenContext,
+ ev: ExprCode,
+ f: String => String): ExprCode = {
+ val argumentGen = argument.genCode(ctx)
+ val resultCode = f(argumentGen.value)
+
+ if (nullable) {
+ val nullSafeEval = ctx.nullSafeExec(argument.nullable, argumentGen.isNull)(resultCode)
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |boolean ${ev.isNull} = ${argumentGen.isNull};
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$nullSafeEval
+ """)
+ } else {
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$resultCode
+ """, isNull = FalseLiteral)
+ }
+ }
+
+ protected def assignVar(
+ varCode: ExprCode,
+ atomicVar: String,
+ value: String,
+ isNull: String,
+ nullable: Boolean): String = {
+ val atomicAssign = assignAtomic(atomicVar, value, isNull, nullable)
+ if (nullable) {
+ s"""
+ ${varCode.value} = $value;
+ ${varCode.isNull} = $isNull;
+ $atomicAssign
+ """
+ } else {
+ s"""
+ ${varCode.value} = $value;
+ $atomicAssign
+ """
+ }
+ }
+
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar, accForMergeVar, accForFinishVar), varCodes => {
+ val Seq(elementCode, accForMergeCode, accForFinishCode) = varCodes
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val i = ctx.freshName("i")
+
+ val zeroCode = zero.genCode(ctx)
+ val Seq(mergeForCodegen, finishForCodegen) = functionsForEval
+ val mergeCode = mergeForCodegen.genCode(ctx)
+ val finishCode = finishForCodegen.genCode(ctx)
Review Comment:
[P1] Widen accumulator nested nullability before codegen
The analyzer accepts `zero` and `merge` types that differ only in nested nullability, but this generated merge/finish path uses accumulator variables bound to `zero.dataType`. For example, an `ARRAY<INT NOT NULL>` zero can legally become an `ARRAY<INT>` containing nulls in `merge`; the generated finish expression then trusts the stale `containsNull = false`, skips `isNullAt`, and returns the primitive default `0` instead of `NULL`. This reproduces with ordinary SQL under the default codegen configuration, while the exact base commit and `NO_CODEGEN` return `NULL`. Please bind the accumulator with `zero.dataType.asNullable -> true` and add array/map regression coverage.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] jaceklaskowski commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "jaceklaskowski (via GitHub)" <gi...@apache.org>.
jaceklaskowski commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1238227638
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -172,6 +172,39 @@ class CodegenContext extends Logging {
*/
var currentVars: Seq[ExprCode] = null
+ /**
+ * Holding a map of current lambda variables.
+ */
+ var currentLambdaVars: mutable.Map[Long, ExprCode] = mutable.HashMap.empty
+
+ def withLambdaVars(namedLambdas: Seq[NamedLambdaVariable],
Review Comment:
Can you put `namedLambdas...` on a separate line?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -172,6 +172,39 @@ class CodegenContext extends Logging {
*/
var currentVars: Seq[ExprCode] = null
+ /**
+ * Holding a map of current lambda variables.
+ */
+ var currentLambdaVars: mutable.Map[Long, ExprCode] = mutable.HashMap.empty
+
+ def withLambdaVars(namedLambdas: Seq[NamedLambdaVariable],
+ f: Seq[ExprCode] => ExprCode): ExprCode = {
+ val lambdaVars = namedLambdas.map { namedLambda =>
Review Comment:
nit: Replace `namedLambda` to `lambda`?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -280,6 +335,29 @@ trait SimpleHigherOrderFunction extends HigherOrderFunction with BinaryLike[Expr
}
}
+ protected def nullSafeCodeGen(
+ ctx: CodegenContext,
+ ev: ExprCode,
+ f: String => String): ExprCode = {
Review Comment:
I'd be very happy if you use this parameter formatting style in the other places in your PR. Makes reading so much easier, esp. with functions with 5+ params. Can you make such change? 🙏
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -781,6 +1007,49 @@ case class ArrayForAll(
}
}
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar), { case Seq(elementCode) =>
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val forall = ctx.freshName("forall")
+ val foundNull = ctx.freshName("foundNull")
+ val i = ctx.freshName("i")
+
+ val functionCode = function.genCode(ctx)
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+
+ val nullCheck = if (nullable) {
+ s"""
+ if ($forall && $foundNull) {
+ ${ev.isNull} = true;
+ }
+ """
+ } else {
+ ""
+ }
+
+ s"""
+ |final int $numElements = ${arg}.numElements();
+ |boolean $forall = true;
+ |boolean $foundNull = false;
+ |int $i = 0;
+ |while ($i < $numElements && $forall) {
+ | $elementAssignment
+ | ${functionCode.code}
+ | if (${functionCode.isNull}) {
+ | $foundNull = true;
+ | } else if (!${functionCode.value}) {
+ | $forall = false;
+ | }
+ | $i++;
+ |}
+ |$nullCheck
+ |${ev.value} = $forall;
+ """.stripMargin
+ })
+ })
+ }
+
Review Comment:
Looks like a copy and paste of `exists`, doesn't it? Can we have a parent class for some sharing? Unless I'm mistaken, the generated code block is the exact copy except `$forall` (vs `$exists`).
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -172,6 +172,39 @@ class CodegenContext extends Logging {
*/
var currentVars: Seq[ExprCode] = null
+ /**
+ * Holding a map of current lambda variables.
+ */
+ var currentLambdaVars: mutable.Map[Long, ExprCode] = mutable.HashMap.empty
+
+ def withLambdaVars(namedLambdas: Seq[NamedLambdaVariable],
+ f: Seq[ExprCode] => ExprCode): ExprCode = {
+ val lambdaVars = namedLambdas.map { namedLambda =>
+ val id = namedLambda.exprId.id
+ if (currentLambdaVars.get(id).nonEmpty) {
+ throw QueryExecutionErrors.lambdaVariableAlreadyDefinedError(id)
+ }
+ val isNull = if (namedLambda.nullable) {
+ JavaCode.isNullGlobal(addMutableState(JAVA_BOOLEAN, "lambdaIsNull"))
+ } else {
+ FalseLiteral
+ }
+ val value = addMutableState(javaType(namedLambda.dataType), "lambdaValue")
+ val lambdaVar = ExprCode(isNull, JavaCode.global(value, namedLambda.dataType))
+ currentLambdaVars.put(id, lambdaVar)
+ lambdaVar
+ }
+
+ val result = f(lambdaVars)
+ namedLambdas.map(_.exprId.id).foreach(currentLambdaVars.remove)
+ result
+ }
+
+ def getLambdaVar(id: Long): ExprCode = {
+ currentLambdaVars.getOrElse(id,
Review Comment:
Can you move `id` to its own new line? 🙏
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-4556766626
Rebased to master to re-trigger CI
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1199788038
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -130,6 +134,23 @@ case class LambdaFunction(
override def eval(input: InternalRow): Any = function.eval(input)
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ val functionCode = function.genCode(ctx)
Review Comment:
Okay nevermind I see what you mean now, it's just assigning the same values to new variables for no reason. Changed to just return `function.genCode(ctx)`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] jaceklaskowski commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "jaceklaskowski (via GitHub)" <gi...@apache.org>.
jaceklaskowski commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1156370630
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -172,6 +172,40 @@ class CodegenContext extends Logging {
*/
var currentVars: Seq[ExprCode] = null
+ /**
+ * Holding a map of current lambda variables.
+ */
+ var currentLambdaVars: mutable.Map[String, ExprCode] = mutable.HashMap.empty
+
+ def withLambdaVars(namedLambdas: Seq[NamedLambdaVariable],
+ f: Seq[ExprCode] => ExprCode): ExprCode = {
+ val lambdaVars = namedLambdas.map { namedLambda =>
+ val name = namedLambda.variableName
+ if (currentLambdaVars.get(name).nonEmpty) {
+ throw QueryExecutionErrors.lambdaVariableAlreadyDefinedError(name)
+ }
+ val isNull = if (namedLambda.nullable) {
+ JavaCode.isNullGlobal(addMutableState(JAVA_BOOLEAN, "lambdaIsNull"))
+ } else {
+ FalseLiteral
+ }
+ val value = addMutableState(javaType(namedLambda.dataType), "lambdaValue")
+ val lambdaVar = ExprCode(isNull, JavaCode.global(value, namedLambda.dataType))
+ currentLambdaVars.put(name, lambdaVar)
+ lambdaVar
+ }
+
+ val result = f(lambdaVars)
+ namedLambdas.foreach(v => currentLambdaVars.remove(v.variableName))
Review Comment:
nit: s/v/lamba? Don't want to ask for `namedLambda` as used earlier, but `v` does not fit really.
I'd also consider `namedLambdas.map(_.variableName).foreach(currentLambdaVars.remove)`
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -101,6 +101,14 @@ case class NamedLambdaVariable(
override def simpleString(maxFields: Int): String = {
s"lambda $name#${exprId.id}: ${dataType.simpleString(maxFields)}"
}
+
+ // We need to include the Expr ID in the Codegen variable name since several tests bypass
+ // `UnresolvedNamedLambdaVariable.freshVarName`
+ lazy val variableName = s"${name}_${exprId.id}"
Review Comment:
It does not seem consistent with `simpleString`. Is this intentional?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -172,6 +172,40 @@ class CodegenContext extends Logging {
*/
var currentVars: Seq[ExprCode] = null
+ /**
+ * Holding a map of current lambda variables.
+ */
+ var currentLambdaVars: mutable.Map[String, ExprCode] = mutable.HashMap.empty
+
+ def withLambdaVars(namedLambdas: Seq[NamedLambdaVariable],
+ f: Seq[ExprCode] => ExprCode): ExprCode = {
+ val lambdaVars = namedLambdas.map { namedLambda =>
+ val name = namedLambda.variableName
+ if (currentLambdaVars.get(name).nonEmpty) {
+ throw QueryExecutionErrors.lambdaVariableAlreadyDefinedError(name)
+ }
+ val isNull = if (namedLambda.nullable) {
+ JavaCode.isNullGlobal(addMutableState(JAVA_BOOLEAN, "lambdaIsNull"))
+ } else {
+ FalseLiteral
+ }
+ val value = addMutableState(javaType(namedLambda.dataType), "lambdaValue")
+ val lambdaVar = ExprCode(isNull, JavaCode.global(value, namedLambda.dataType))
+ currentLambdaVars.put(name, lambdaVar)
+ lambdaVar
+ }
+
+ val result = f(lambdaVars)
+ namedLambdas.foreach(v => currentLambdaVars.remove(v.variableName))
+ result
+ }
+
+ def getLambdaVar(name: String): ExprCode = {
+ currentLambdaVars.getOrElse(name, {
Review Comment:
nit: Are these curly brackets required?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -350,6 +445,49 @@ case class ArrayTransform(
result
}
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar) ++ indexVar, { lambdaExprs =>
Review Comment:
nit: Be consistent with `{`s; here, `{` is before an input argument to a function literal while a few lines below it's after `=>`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1198880920
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala:
##########
@@ -149,9 +149,13 @@ class EquivalentExpressions(
// There are some special expressions that we should not recurse into all of its children.
// 1. CodegenFallback: it's children will not be used to generate code (call eval() instead)
// 2. ConditionalExpression: use its children that will always be evaluated.
+ // 3. HigherOrderFunction: lambda functions operate in the context of local lambdas and can't
+ // be called outside of that scope, only the arguments can be evaluated ahead of
+ // time.
private def childrenToRecurse(expr: Expression): Seq[Expression] = expr match {
case _: CodegenFallback => Nil
case c: ConditionalExpression => c.alwaysEvaluatedInputs.map(skipForShortcut)
+ case h: HigherOrderFunction => h.arguments
Review Comment:
I don't think so. That only cares if the current expression is a `ConditionalExpression`. The default is `Nil` if it's not that so I don't think it needs any special handling for HOFs
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] cloud-fan commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1198526911
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -130,6 +134,23 @@ case class LambdaFunction(
override def eval(input: InternalRow): Any = function.eval(input)
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ val functionCode = function.genCode(ctx)
Review Comment:
seems we can just return `function.genCode(ctx)`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1239152129
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -781,6 +1007,49 @@ case class ArrayForAll(
}
}
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar), { case Seq(elementCode) =>
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val forall = ctx.freshName("forall")
+ val foundNull = ctx.freshName("foundNull")
+ val i = ctx.freshName("i")
+
+ val functionCode = function.genCode(ctx)
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+
+ val nullCheck = if (nullable) {
+ s"""
+ if ($forall && $foundNull) {
+ ${ev.isNull} = true;
+ }
+ """
+ } else {
+ ""
+ }
+
+ s"""
+ |final int $numElements = ${arg}.numElements();
+ |boolean $forall = true;
+ |boolean $foundNull = false;
+ |int $i = 0;
+ |while ($i < $numElements && $forall) {
+ | $elementAssignment
+ | ${functionCode.code}
+ | if (${functionCode.isNull}) {
+ | $foundNull = true;
+ | } else if (!${functionCode.value}) {
+ | $forall = false;
+ | }
+ | $i++;
+ |}
+ |$nullCheck
+ |${ev.value} = $forall;
+ """.stripMargin
+ })
+ })
+ }
+
Review Comment:
I'll have to look at this a little bit. It's tricky because there's a few places they are different (default true vs false, `!`'d for one and not the other, the `followThreeValuedLogic` flag for exists)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1993969482
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -235,6 +256,53 @@ trait HigherOrderFunction extends Expression with ExpectsInputTypes {
val canonicalizedChildren = cleaned.children.map(_.canonicalized)
withNewChildren(canonicalizedChildren)
}
+
+
+ protected def assignAtomic(atomicRef: String, value: String, isNull: String = FalseLiteral,
+ nullable: Boolean = false) = {
+ if (nullable) {
+ s"""
+ if ($isNull) {
+ $atomicRef.set(null);
+ } else {
+ $atomicRef.set($value);
+ }
+ """
+ } else {
+ s"$atomicRef.set($value);"
+ }
+ }
+
+ protected def assignArrayElement(ctx: CodegenContext, arrayName: String, elementCode: ExprCode,
+ elementVar: NamedLambdaVariable, index: String): String = {
+ val elementType = elementVar.dataType
+ val elementAtomic = ctx.addReferenceObj(elementVar.name, elementVar.value)
+ val extractElement = CodeGenerator.getValue(arrayName, elementType, index)
Review Comment:
Do you just mean the built-in concurrency of `AtomicReference` isn't needed except for something related to the the `transform` operations, and otherwise you could just use a regular variable to store the result?
Agree that might be preferable, but probably out of scope for this initial codegen implementation since that would change how the interpreted path works as well
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1198915674
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -130,6 +134,23 @@ case class LambdaFunction(
override def eval(input: InternalRow): Any = function.eval(input)
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ val functionCode = function.genCode(ctx)
Review Comment:
But maybe this already handles that:
```
val resultNull = if (function.nullable) Some(functionCode.isNull.toString) else None
val resultAssignment = CodeGenerator.setArrayElement(arrayData, dataType.elementType,
i, copy, isNull = resultNull)
```
I'll need to think about this a little more
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1157743684
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -172,6 +172,40 @@ class CodegenContext extends Logging {
*/
var currentVars: Seq[ExprCode] = null
+ /**
+ * Holding a map of current lambda variables.
+ */
+ var currentLambdaVars: mutable.Map[String, ExprCode] = mutable.HashMap.empty
+
+ def withLambdaVars(namedLambdas: Seq[NamedLambdaVariable],
+ f: Seq[ExprCode] => ExprCode): ExprCode = {
+ val lambdaVars = namedLambdas.map { namedLambda =>
+ val name = namedLambda.variableName
+ if (currentLambdaVars.get(name).nonEmpty) {
+ throw QueryExecutionErrors.lambdaVariableAlreadyDefinedError(name)
+ }
+ val isNull = if (namedLambda.nullable) {
+ JavaCode.isNullGlobal(addMutableState(JAVA_BOOLEAN, "lambdaIsNull"))
+ } else {
+ FalseLiteral
+ }
+ val value = addMutableState(javaType(namedLambda.dataType), "lambdaValue")
+ val lambdaVar = ExprCode(isNull, JavaCode.global(value, namedLambda.dataType))
+ currentLambdaVars.put(name, lambdaVar)
+ lambdaVar
+ }
+
+ val result = f(lambdaVars)
+ namedLambdas.foreach(v => currentLambdaVars.remove(v.variableName))
Review Comment:
Yeah I never know all the exact ways that using `_` does and doesn't work when you use a function call, struggle with the cleanest way to write things like this. I like the second one I think
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -172,6 +172,40 @@ class CodegenContext extends Logging {
*/
var currentVars: Seq[ExprCode] = null
+ /**
+ * Holding a map of current lambda variables.
+ */
+ var currentLambdaVars: mutable.Map[String, ExprCode] = mutable.HashMap.empty
+
+ def withLambdaVars(namedLambdas: Seq[NamedLambdaVariable],
+ f: Seq[ExprCode] => ExprCode): ExprCode = {
+ val lambdaVars = namedLambdas.map { namedLambda =>
+ val name = namedLambda.variableName
+ if (currentLambdaVars.get(name).nonEmpty) {
+ throw QueryExecutionErrors.lambdaVariableAlreadyDefinedError(name)
+ }
+ val isNull = if (namedLambda.nullable) {
+ JavaCode.isNullGlobal(addMutableState(JAVA_BOOLEAN, "lambdaIsNull"))
+ } else {
+ FalseLiteral
+ }
+ val value = addMutableState(javaType(namedLambda.dataType), "lambdaValue")
+ val lambdaVar = ExprCode(isNull, JavaCode.global(value, namedLambda.dataType))
+ currentLambdaVars.put(name, lambdaVar)
+ lambdaVar
+ }
+
+ val result = f(lambdaVars)
+ namedLambdas.foreach(v => currentLambdaVars.remove(v.variableName))
+ result
+ }
+
+ def getLambdaVar(name: String): ExprCode = {
+ currentLambdaVars.getOrElse(name, {
Review Comment:
I saw this format in other parts of the code, but I also have seen it without, so can remove
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3293711737
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -886,6 +1179,114 @@ case class ArrayAggregate(
}
}
+ protected def nullSafeCodeGen(
+ ctx: CodegenContext,
+ ev: ExprCode,
+ f: String => String): ExprCode = {
+ val argumentGen = argument.genCode(ctx)
+ val resultCode = f(argumentGen.value)
+
+ if (nullable) {
+ val nullSafeEval = ctx.nullSafeExec(argument.nullable, argumentGen.isNull)(resultCode)
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |boolean ${ev.isNull} = ${argumentGen.isNull};
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$nullSafeEval
+ """)
+ } else {
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$resultCode
+ """, isNull = FalseLiteral)
+ }
+ }
+
+ protected def assignVar(
+ varCode: ExprCode,
+ atomicVar: String,
+ value: String,
+ isNull: String,
+ nullable: Boolean): String = {
+ val atomicAssign = assignAtomic(atomicVar, value, isNull, nullable)
+ if (nullable) {
+ s"""
+ ${varCode.value} = $value;
+ ${varCode.isNull} = $isNull;
+ $atomicAssign
+ """
+ } else {
+ s"""
+ ${varCode.value} = $value;
+ $atomicAssign
+ """
+ }
+ }
+
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar, accForMergeVar, accForFinishVar), varCodes => {
+ val Seq(elementCode, accForMergeCode, accForFinishCode) = varCodes
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val i = ctx.freshName("i")
+
+ val zeroCode = zero.genCode(ctx)
+ val mergeCode = merge.genCode(ctx)
+ val finishCode = finish.genCode(ctx)
+
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+ val mergeAtomic = ctx.addReferenceObj(accForMergeVar.name,
+ accForMergeVar.value)
+ val finishAtomic = ctx.addReferenceObj(accForFinishVar.name,
+ accForFinishVar.value)
+
+ val mergeJavaType = CodeGenerator.javaType(accForMergeVar.dataType)
+ val finishJavaType = CodeGenerator.javaType(accForFinishVar.dataType)
+
+ // Some expressions return internal buffers that we have to copy
+ val mergeCopy = if (CodeGenerator.isPrimitiveType(merge.dataType)) {
+ s"${mergeCode.value}"
+ } else {
+ s"($mergeJavaType)InternalRow.copyValue(${mergeCode.value})"
+ }
+
+ val nullCheck = if (nullable) {
+ s"${ev.isNull} = ${finishCode.isNull};"
+ } else {
+ ""
+ }
+
+ val initialAssignment = assignVar(accForMergeCode, mergeAtomic, zeroCode.value,
+ zeroCode.isNull, zero.nullable)
+
+ val mergeAssignment = assignVar(accForMergeCode, mergeAtomic, mergeCopy,
+ mergeCode.isNull, merge.nullable)
Review Comment:
[P1] Clear `ArrayAggregate` accumulator null state on every generated assignment
`accForMergeVar` is always bound nullable in `bindInternal`, and `withLambdaVars` stores its `isNull` flag as mutable generated state. These assignments only update that flag when the source expression is nullable, which produces wrong results in both cases below:
- Within one row, `aggregate(array(CAST(id AS INT) + 1, CAST(id AS INT) + 2), CAST(NULL AS INT), (acc, x) -> coalesce(acc, 0) + x, acc -> coalesce(acc, -1))` returns `-1` with `CODEGEN_ONLY` instead of `3`; the same query returns `3` with `NO_CODEGEN`.
- Across rows in one generated partition, after a row whose merge becomes null, an empty-array row with zero `0` returns `-1` instead of `0` with `CODEGEN_ONLY`; the interpreted path returns `0`.
Please pass `accForMergeVar.nullable` to both `initialAssignment` and `mergeAssignment` and add generated-code regressions for these paths.
_[ :robot: posted by Codex on behalf of sunchao using the code-review-for-me skill :robot: ]_
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3293007181
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -622,6 +753,72 @@ case class ArrayFilter(
new GenericArrayData(buffer)
}
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar) ++ indexVar, varCodes => {
+ val elementCode = varCodes.head
+ val indexCode = varCodes.tail.headOption
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val count = ctx.freshName("count")
+ val arrayTracker = ctx.freshName("arrayTracker")
+ val arrayData = ctx.freshName("arrayData")
+ val i = ctx.freshName("i")
+ val j = ctx.freshName("j")
+
+ val arrayType = dataType.asInstanceOf[ArrayType]
+
+ val trackerInit = CodeGenerator.createArrayData(
+ arrayTracker, BooleanType, numElements, s" $prettyName failed.")
+ val resultInit = CodeGenerator.createArrayData(
+ arrayData, arrayType.elementType, count, s" $prettyName failed.")
+
+ val functionCode = function.genCode(ctx)
+
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+ val indexAssignment = indexCode.map(c => assignIndex(ctx, c, indexVar.get, i))
+ val varAssignments = (Seq(elementAssignment) ++ indexAssignment).mkString("\n")
+
+ val resultAssignment = CodeGenerator.setArrayElement(arrayTracker, BooleanType,
+ i, functionCode.value, isNull = None)
+
+ val getTrackerValue = CodeGenerator.getValue(arrayTracker, BooleanType, i)
+ val copy = CodeGenerator.createArrayAssignment(arrayData, arrayType.elementType, arg,
+ j, i, arrayType.containsNull)
+
+ // This takes a two passes to avoid evaluating the predicate multiple times
+ // The first pass evaluates each element in the array, tracks how many elements
+ // returned true, and tracks the result of each element in a boolean array `arrayTracker`.
+ // The second pass copies elements from the original array to the new array created
+ // based on the number of elements matching the first pass.
+
+ s"""
+ |final int $numElements = $arg.numElements();
+ |$trackerInit
+ |int $count = 0;
+ |for (int $i = 0; $i < $numElements; $i++) {
+ | $varAssignments
+ | ${functionCode.code}
+ | $resultAssignment
+ | if ((boolean)${functionCode.value}) {
Review Comment:
[P1] Preserve null predicate semantics in generated `ArrayFilter`
`filter` treats a null predicate result as false, but this generated loop reads `${functionCode.value}` without checking `${functionCode.isNull}`. I reproduced this on the current head with a vectorized Parquet array column containing `[true]` and `[null]`: with `CODEGEN_ONLY`, `filter(a, x -> x)` returns `[true]` and `[null]` instead of `[true]` and `[]`; with `NO_CODEGEN`, it returns the expected results. Please use a predicate equivalent to `!functionCode.isNull && functionCode.value` for both the tracker and count, and add a generated-code regression test.
_\[ :robot: posted by Codex on behalf of sunchao using the [`code-review-for-me`](https://github.com/openai/openai/tree/master/skills/skills/code-review-for-me) skill :robot: \]_
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala:
##########
@@ -146,9 +146,13 @@ class EquivalentExpressions(
// There are some special expressions that we should not recurse into all of its children.
// 1. CodegenFallback: it's children will not be used to generate code (call eval() instead)
// 2. ConditionalExpression: use its children that will always be evaluated.
+ // 3. HigherOrderFunction: lambda functions operate in the context of local lambdas and can't
+ // be called outside of that scope, only the arguments can be evaluated ahead of
+ // time.
private def childrenToRecurse(expr: Expression): Seq[Expression] = expr match {
case _: CodegenFallback => Nil
case c: ConditionalExpression => c.alwaysEvaluatedInputs.map(skipForShortcut)
+ case h: HigherOrderFunction => h.arguments
Review Comment:
[P1] Do not eagerly evaluate `ArrayAggregate.zero` for null input arrays
For `ArrayAggregate`, `zero` is not evaluated when the array argument is null, but recursing into all `h.arguments` allows subexpression elimination to generate expressions from `zero` outside that null guard. On the current head, with ANSI mode, CSE, and generated code enabled, `aggregate(CAST(NULL AS ARRAY<INT>), (cast(id as int) / 0) + (cast(id as int) / 0), (acc, x) -> acc + x)` throws `DIVIDE_BY_ZERO`; with CSE disabled, the same generated query returns null. Please recurse only into the always-evaluated input array for `ArrayAggregate` and add this regression. This is separate from the empty-array nullable-accumulator fix in `9f330488952`.
_\[ :robot: posted by Codex on behalf of sunchao using the [`code-review-for-me`](https://github.com/openai/openai/tree/master/skills/skills/code-review-for-me) skill :robot: \]_
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -354,6 +442,49 @@ case class ArrayTransform(
result
}
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar) ++ indexVar, varCodes => {
+ val elementCode = varCodes.head
+ val indexCode = varCodes.tail.headOption
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val arrayData = ctx.freshName("arrayData")
+ val i = ctx.freshName("i")
+
+ val initialization = CodeGenerator.createArrayData(
+ arrayData, dataType.elementType, numElements, s" $prettyName failed.")
+
+ val functionCode = function.genCode(ctx)
+
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+ val indexAssignment = indexCode.map(c => assignIndex(ctx, c, indexVar.get, i))
+ val varAssignments = (Seq(elementAssignment) ++ indexAssignment).mkString("\n")
+
+ // Some expressions return internal buffers that we have to copy
+ val copy = if (CodeGenerator.isPrimitiveType(function.dataType)) {
+ s"${functionCode.value}"
+ } else {
+ s"InternalRow.copyValue(${functionCode.value})"
+ }
+ val resultNull = if (function.nullable) Some(functionCode.isNull.toString) else None
+ val resultAssignment = CodeGenerator.setArrayElement(arrayData, dataType.elementType,
Review Comment:
[P1] Preserve null complex elements in generated array HOF output
This new output-writing path passes a null bit for nullable lambda results into `CodeGenerator.setArrayElement`, but that helper applies `setNullAt` only for primitive element types. I reproduced this on the current head with native vectorized ORC input containing `array(CAST(NULL AS ARRAY<INT>), array(1))`: the input reads back with `a[0] IS NULL = true`, and interpreted `transform`/`filter` return `[true, true]`, but generated `transform(a, x -> x)[0] IS NULL` and `filter(a, x -> x IS NULL)[0] IS NULL` return `[false, false]`. Please make nullable non-primitive array writes honor `isNull` and add a generated nested-null regression test.
_\[ :robot: posted by Codex on behalf of sunchao using the [`code-review-for-me`](https://github.com/openai/openai/tree/master/skills/skills/code-review-for-me) skill :robot: \]_
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-4633886403
@Kimahriman sorry just found another issue. Could you address it? I'm planning to merge this afterwards.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "chris-twiner (via GitHub)" <gi...@apache.org>.
chris-twiner commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-2722268943
> > > @Kimahriman just out of curiosity, how much did the performance improve?
> >
> >
> > I just wanted to add to the above response that I've implemented a compilation scheme [here](https://sparkutils.github.io/quality/latest/advanced/userFunctions/#controlling-compilation-tweaking-the-quality-optimisations), as part of Quality, and we saw perf boosts of up to 40%, after that adding further lambdas triggered the cost of code generation being higher than the saving. It's definitely usage dependant though, the more work done in the function the higher the cost (and therefore potential saving by compilation), a small boost is noticeable on removal of the atomic under similar ideal circumstances.
> > edit - [the source](https://github.com/sparkutils/quality/blob/main/src/main/scala/org/apache/spark/sql/qualityFunctions/LambdaCompilation.scala)
>
> Not sure how I missed this comment. we haven't done extensive performance comparisons with and without this, we've just been using it for a few years now. It's hard to quantify the impact since it's completely dependent on the expressions run inside the functions. But that's also the whole point, by enabling codegen for HOFs you enable codegen for expressions inside the lambda functions, which are assumed to be more performant since that's the whole point of codegen.
>
> Additionally this enables a follow on I'm currently working on which is enabling subexpression elimination inside of lambda functions, which we've recently identified as a major performance killer for us, as it's very easy to generate a lot of duplicate expression evaluations in certain cases
The subexpression elimination option is huge! Very exciting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1198873238
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -235,6 +256,53 @@ trait HigherOrderFunction extends Expression with ExpectsInputTypes {
val canonicalizedChildren = cleaned.children.map(_.canonicalized)
withNewChildren(canonicalizedChildren)
}
+
+
+ protected def assignAtomic(atomicRef: String, value: String, isNull: String = FalseLiteral,
+ nullable: Boolean = false) = {
+ if (nullable) {
+ s"""
+ if ($isNull) {
+ $atomicRef.set(null);
+ } else {
+ $atomicRef.set($value);
+ }
+ """
+ } else {
+ s"$atomicRef.set($value);"
+ }
+ }
+
+ protected def assignArrayElement(ctx: CodegenContext, arrayName: String, elementCode: ExprCode,
+ elementVar: NamedLambdaVariable, index: String): String = {
+ val elementType = elementVar.dataType
+ val elementAtomic = ctx.addReferenceObj(elementVar.name, elementVar.value)
+ val extractElement = CodeGenerator.getValue(arrayName, elementType, index)
Review Comment:
The atomic references are for cases where you might have a `CodegenFallback` expression inside your lambda function. I think the only way an `expression.eval(input)` will work is if this atomic reference is set to the lambda variable value as well. I could potentially try to detect if the function has a `CodegenFallback` expression, but I wasn't sure if there were other reasons some expression would fall back to interpreted inside the lambda function that I couldn't detect.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-1603478055
> More nitting 😉
Yeah I struggle with the right formatting for multiline things in scala, tried to update all the suggestions, thanks for the tips!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3293466426
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -354,6 +442,49 @@ case class ArrayTransform(
result
}
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar) ++ indexVar, varCodes => {
+ val elementCode = varCodes.head
+ val indexCode = varCodes.tail.headOption
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val arrayData = ctx.freshName("arrayData")
+ val i = ctx.freshName("i")
+
+ val initialization = CodeGenerator.createArrayData(
+ arrayData, dataType.elementType, numElements, s" $prettyName failed.")
+
+ val functionCode = function.genCode(ctx)
+
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+ val indexAssignment = indexCode.map(c => assignIndex(ctx, c, indexVar.get, i))
+ val varAssignments = (Seq(elementAssignment) ++ indexAssignment).mkString("\n")
+
+ // Some expressions return internal buffers that we have to copy
+ val copy = if (CodeGenerator.isPrimitiveType(function.dataType)) {
+ s"${functionCode.value}"
+ } else {
+ s"InternalRow.copyValue(${functionCode.value})"
+ }
+ val resultNull = if (function.nullable) Some(functionCode.isNull.toString) else None
+ val resultAssignment = CodeGenerator.setArrayElement(arrayData, dataType.elementType,
Review Comment:
Fixed in `eecc672cb44`.
`CodeGenerator.setArrayElement` now honors the provided `isNull` bit for non-primitive element types as well, so nullable complex array writes call `setNullAt` instead of writing a stale complex value.
I moved the regression to Catalyst with custom `GenericArrayData` that simulates the vectorized-reader shape: `isNullAt(0)` reports null while the typed complex getter can still expose a stale nested value. The test covers both generated `transform` output and `filter` output.
Verified with:
```bash
build/sbt 'catalyst/testOnly org.apache.spark.sql.catalyst.expressions.HigherOrderFunctionsSuite -- -z "generated code"'
```
Result: passed (`2` tests, `0` failures).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman closed pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman closed pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
URL: https://github.com/apache/spark/pull/34558
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-1496583248
> There seems to be a lot of repetition. Wish it could be avoided somehow but can't help though (beside nit-picking).
Thanks for the review! I tried to get as much common code in the parent classes as I could, can take another pass to see if anything jumps out for deduping
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "chris-twiner (via GitHub)" <gi...@apache.org>.
chris-twiner commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1994100057
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -235,6 +256,53 @@ trait HigherOrderFunction extends Expression with ExpectsInputTypes {
val canonicalizedChildren = cleaned.children.map(_.canonicalized)
withNewChildren(canonicalizedChildren)
}
+
+
+ protected def assignAtomic(atomicRef: String, value: String, isNull: String = FalseLiteral,
+ nullable: Boolean = false) = {
+ if (nullable) {
+ s"""
+ if ($isNull) {
+ $atomicRef.set(null);
+ } else {
+ $atomicRef.set($value);
+ }
+ """
+ } else {
+ s"$atomicRef.set($value);"
+ }
+ }
+
+ protected def assignArrayElement(ctx: CodegenContext, arrayName: String, elementCode: ExprCode,
+ elementVar: NamedLambdaVariable, index: String): String = {
+ val elementType = elementVar.dataType
+ val elementAtomic = ctx.addReferenceObj(elementVar.name, elementVar.value)
+ val extractElement = CodeGenerator.getValue(arrayName, elementType, index)
Review Comment:
Yes, that's definitely all it is, I may spend time to find a reproducible case of this but it's not in the hof test suite.
It makes sense to keep the two things separate.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3297942369
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -886,6 +1179,114 @@ case class ArrayAggregate(
}
}
+ protected def nullSafeCodeGen(
+ ctx: CodegenContext,
+ ev: ExprCode,
+ f: String => String): ExprCode = {
+ val argumentGen = argument.genCode(ctx)
+ val resultCode = f(argumentGen.value)
+
+ if (nullable) {
+ val nullSafeEval = ctx.nullSafeExec(argument.nullable, argumentGen.isNull)(resultCode)
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |boolean ${ev.isNull} = ${argumentGen.isNull};
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$nullSafeEval
+ """)
+ } else {
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$resultCode
+ """, isNull = FalseLiteral)
+ }
+ }
+
+ protected def assignVar(
+ varCode: ExprCode,
+ atomicVar: String,
+ value: String,
+ isNull: String,
+ nullable: Boolean): String = {
+ val atomicAssign = assignAtomic(atomicVar, value, isNull, nullable)
+ if (nullable) {
+ s"""
+ ${varCode.value} = $value;
+ ${varCode.isNull} = $isNull;
+ $atomicAssign
+ """
+ } else {
+ s"""
+ ${varCode.value} = $value;
+ $atomicAssign
+ """
+ }
+ }
+
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar, accForMergeVar, accForFinishVar), varCodes => {
+ val Seq(elementCode, accForMergeCode, accForFinishCode) = varCodes
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val i = ctx.freshName("i")
+
+ val zeroCode = zero.genCode(ctx)
+ val mergeCode = merge.genCode(ctx)
+ val finishCode = finish.genCode(ctx)
+
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+ val mergeAtomic = ctx.addReferenceObj(accForMergeVar.name,
+ accForMergeVar.value)
+ val finishAtomic = ctx.addReferenceObj(accForFinishVar.name,
+ accForFinishVar.value)
+
+ val mergeJavaType = CodeGenerator.javaType(accForMergeVar.dataType)
+ val finishJavaType = CodeGenerator.javaType(accForFinishVar.dataType)
+
+ // Some expressions return internal buffers that we have to copy
+ val mergeCopy = if (CodeGenerator.isPrimitiveType(merge.dataType)) {
+ s"${mergeCode.value}"
+ } else {
+ s"($mergeJavaType)InternalRow.copyValue(${mergeCode.value})"
+ }
+
+ val nullCheck = if (nullable) {
+ s"${ev.isNull} = ${finishCode.isNull};"
+ } else {
+ ""
+ }
+
+ val initialAssignment = assignVar(accForMergeCode, mergeAtomic, zeroCode.value,
+ zeroCode.isNull, zero.nullable)
+
+ val mergeAssignment = assignVar(accForMergeCode, mergeAtomic, mergeCopy,
+ mergeCode.isNull, merge.nullable)
Review Comment:
Fixed in `261147f456d`.
`initialAssignment` and `mergeAssignment` now both use `accForMergeVar.nullable`, so the generated accumulator lambda variable null flag is updated on every assignment. That clears the stale-null state both after assigning a non-null `zero` and after a non-null merge result.
I added generated-code regressions for both paths from the comment:
- stale null state within a single aggregate loop
- stale null state carried across rows in the same generated partition
I verified both tests fail before the fix and pass after it.
Verified with:
```bash
build/sbt 'sql/testOnly org.apache.spark.sql.DataFrameFunctionsSuite -- -z "aggregate function - generated code clears accumulator null state"'
```
Result: passed (`2` tests, `0` failures).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3364790060
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -886,6 +1179,115 @@ case class ArrayAggregate(
}
}
+ protected def nullSafeCodeGen(
+ ctx: CodegenContext,
+ ev: ExprCode,
+ f: String => String): ExprCode = {
+ val argumentGen = argument.genCode(ctx)
+ val resultCode = f(argumentGen.value)
+
+ if (nullable) {
+ val nullSafeEval = ctx.nullSafeExec(argument.nullable, argumentGen.isNull)(resultCode)
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |boolean ${ev.isNull} = ${argumentGen.isNull};
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$nullSafeEval
+ """)
+ } else {
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$resultCode
+ """, isNull = FalseLiteral)
+ }
+ }
+
+ protected def assignVar(
+ varCode: ExprCode,
+ atomicVar: String,
+ value: String,
+ isNull: String,
+ nullable: Boolean): String = {
+ val atomicAssign = assignAtomic(atomicVar, value, isNull, nullable)
+ if (nullable) {
+ s"""
+ ${varCode.value} = $value;
+ ${varCode.isNull} = $isNull;
+ $atomicAssign
+ """
+ } else {
+ s"""
+ ${varCode.value} = $value;
+ $atomicAssign
+ """
+ }
+ }
+
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar, accForMergeVar, accForFinishVar), varCodes => {
+ val Seq(elementCode, accForMergeCode, accForFinishCode) = varCodes
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val i = ctx.freshName("i")
+
+ val zeroCode = zero.genCode(ctx)
+ val Seq(mergeForCodegen, finishForCodegen) = functionsForEval
+ val mergeCode = mergeForCodegen.genCode(ctx)
+ val finishCode = finishForCodegen.genCode(ctx)
Review Comment:
Thanks, fixed in the latest push by binding the aggregate accumulator lambda variables with `zero.dataType.asNullable -> true`, so generated finish/merge code sees widened nested nullability. I also added regression coverage for both array element nullability and map value nullability across `CODEGEN_ONLY` and `NO_CODEGEN`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-4556451061
@Kimahriman can you also fix the CI failures? cc @cloud-fan @peter-toth @viirya also for reviews
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] cloud-fan commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1198529114
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -235,6 +256,53 @@ trait HigherOrderFunction extends Expression with ExpectsInputTypes {
val canonicalizedChildren = cleaned.children.map(_.canonicalized)
withNewChildren(canonicalizedChildren)
}
+
+
+ protected def assignAtomic(atomicRef: String, value: String, isNull: String = FalseLiteral,
+ nullable: Boolean = false) = {
+ if (nullable) {
+ s"""
+ if ($isNull) {
+ $atomicRef.set(null);
+ } else {
+ $atomicRef.set($value);
+ }
+ """
+ } else {
+ s"$atomicRef.set($value);"
+ }
+ }
+
+ protected def assignArrayElement(ctx: CodegenContext, arrayName: String, elementCode: ExprCode,
+ elementVar: NamedLambdaVariable, index: String): String = {
+ val elementType = elementVar.dataType
+ val elementAtomic = ctx.addReferenceObj(elementVar.name, elementVar.value)
+ val extractElement = CodeGenerator.getValue(arrayName, elementType, index)
Review Comment:
With codegen, I don't think we need a global atomic reference to represent the lambda variable. We can just generate local variables in the java code.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-1603476095
>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-2722033268
> > @Kimahriman just out of curiosity, how much did the performance improve?
>
> I just wanted to add to the above response that I've implemented a compilation scheme [here](https://sparkutils.github.io/quality/latest/advanced/userFunctions/#controlling-compilation-tweaking-the-quality-optimisations), as part of Quality, and we saw perf boosts of up to 40%, after that adding further lambdas triggered the cost of code generation being higher than the saving. It's definitely usage dependant though, the more work done in the function the higher the cost (and therefore potential saving by compilation), a small boost is noticeable on removal of the atomic under similar ideal circumstances.
>
> edit - [the source](https://github.com/sparkutils/quality/blob/main/src/main/scala/org/apache/spark/sql/qualityFunctions/LambdaCompilation.scala)
Not sure how I missed this comment. we haven't done extensive performance comparisons with and without this, we've just been using it for a few years now. It's hard to quantify the impact since it's completely dependent on the expressions run inside the functions. But that's also the whole point, by enabling codegen for HOFs you enable codegen for expressions inside the lambda functions, which are assumed to be more performant since that's the whole point of codegen.
Additionally this enables a follow on I'm currently working on which is enabling subexpression elimination inside of lambda functions, which we've recently identified as a major performance killer for us, as it's very easy to generate a lot of duplicate expression evaluations in certain cases
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3290456907
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -886,6 +1170,114 @@ case class ArrayAggregate(
}
}
+ protected def nullSafeCodeGen(
+ ctx: CodegenContext,
+ ev: ExprCode,
+ f: String => String): ExprCode = {
+ val argumentGen = argument.genCode(ctx)
+ val resultCode = f(argumentGen.value)
+
+ if (nullable) {
+ val nullSafeEval = ctx.nullSafeExec(argument.nullable, argumentGen.isNull)(resultCode)
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |boolean ${ev.isNull} = ${argumentGen.isNull};
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$nullSafeEval
+ """)
+ } else {
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$resultCode
+ """, isNull = FalseLiteral)
+ }
+ }
+
+ protected def assignVar(
+ varCode: ExprCode,
+ atomicVar: String,
+ value: String,
+ isNull: String,
+ nullable: Boolean): String = {
+ val atomicAssign = assignAtomic(atomicVar, value, isNull, nullable)
+ if (nullable) {
+ s"""
+ ${varCode.value} = $value;
+ ${varCode.isNull} = $isNull;
+ $atomicAssign
+ """
+ } else {
+ s"""
+ ${varCode.value} = $value;
+ $atomicAssign
+ """
+ }
+ }
+
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar, accForMergeVar, accForFinishVar), varCodes => {
+ val Seq(elementCode, accForMergeCode, accForFinishCode) = varCodes
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val i = ctx.freshName("i")
+
+ val zeroCode = zero.genCode(ctx)
+ val mergeCode = merge.genCode(ctx)
+ val finishCode = finish.genCode(ctx)
+
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+ val mergeAtomic = ctx.addReferenceObj(accForMergeVar.name,
+ accForMergeVar.value)
+ val finishAtomic = ctx.addReferenceObj(accForFinishVar.name,
+ accForFinishVar.value)
+
+ val mergeJavaType = CodeGenerator.javaType(accForMergeVar.dataType)
+ val finishJavaType = CodeGenerator.javaType(accForFinishVar.dataType)
+
+ // Some expressions return internal buffers that we have to copy
+ val mergeCopy = if (CodeGenerator.isPrimitiveType(merge.dataType)) {
+ s"${mergeCode.value}"
+ } else {
+ s"($mergeJavaType)InternalRow.copyValue(${mergeCode.value})"
+ }
+
+ val nullCheck = if (nullable) {
+ s"${ev.isNull} = ${finishCode.isNull};"
+ } else {
+ ""
+ }
+
+ val initialAssignment = assignVar(accForMergeCode, mergeAtomic, zeroCode.value,
+ zeroCode.isNull, zero.nullable)
+
+ val mergeAssignment = assignVar(accForMergeCode, mergeAtomic, mergeCopy,
+ mergeCode.isNull, merge.nullable)
+
+ val finishAssignment = assignVar(accForFinishCode, finishAtomic, accForMergeCode.value,
+ accForMergeCode.isNull, merge.nullable)
Review Comment:
Thanks, good catch. Fixed in `9f330488952` by making `finishAssignment` follow the accumulator lambda variable nullability instead of `merge.nullable`, so a nullable `zero` is still propagated into `finish` when the array is empty.
I also added the regression case from the comment and verified it with:
```bash
build/sbt 'catalyst/testOnly org.apache.spark.sql.catalyst.expressions.HigherOrderFunctionsSuite -- -z ArrayAggregate'
```
Result: passed (`2` tests, `0` failures).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-4651855021
Merged to `master` / `branch-4.x`, thanks @Kimahriman !! especially for your patience on this PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao closed pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
URL: https://github.com/apache/spark/pull/34558
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "chris-twiner (via GitHub)" <gi...@apache.org>.
chris-twiner commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1993925656
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -235,6 +256,53 @@ trait HigherOrderFunction extends Expression with ExpectsInputTypes {
val canonicalizedChildren = cleaned.children.map(_.canonicalized)
withNewChildren(canonicalizedChildren)
}
+
+
+ protected def assignAtomic(atomicRef: String, value: String, isNull: String = FalseLiteral,
+ nullable: Boolean = false) = {
+ if (nullable) {
+ s"""
+ if ($isNull) {
+ $atomicRef.set(null);
+ } else {
+ $atomicRef.set($value);
+ }
+ """
+ } else {
+ s"$atomicRef.set($value);"
+ }
+ }
+
+ protected def assignArrayElement(ctx: CodegenContext, arrayName: String, elementCode: ExprCode,
+ elementVar: NamedLambdaVariable, index: String): String = {
+ val elementType = elementVar.dataType
+ val elementAtomic = ctx.addReferenceObj(elementVar.name, elementVar.value)
+ val extractElement = CodeGenerator.getValue(arrayName, elementType, index)
Review Comment:
> Finally got back around to this, added tests showing why the atomic reference is needed in case of `CodegenFallback` inside a lambda expression
FYI - In [my testing](https://github.com/sparkutils/quality/blob/temp/fabric_0.1.3.1_rc6_map_with_issues/src/main/scala/org/apache/spark/sql/qualityFunctions/LambdaCompilation.scala#L21) you can seemingly, at codegen time, swap out the AtomicReference/NamedLambdaVariable in all cases to a simple set/get wrapper around a field for the eval CodegenFallback case.
Per [the original NamedLambdaVariable PR](https://github.com/apache/spark/pull/21954#discussion_r207162167) the AtomicReference only seems to be needed for proper tree transformations, so you can swap out at compilation time to squeeze a little bit [more perf out](https://brooker.co.za/blog/2012/09/10/volatile.html) (lower memory usage, no need for volatile read/memory barriers on write etc. this only really makes a difference if the code is HOF heavy of course, like Quality tends to be).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] cloud-fan commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1198524171
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala:
##########
@@ -149,9 +149,13 @@ class EquivalentExpressions(
// There are some special expressions that we should not recurse into all of its children.
// 1. CodegenFallback: it's children will not be used to generate code (call eval() instead)
// 2. ConditionalExpression: use its children that will always be evaluated.
+ // 3. HigherOrderFunction: lambda functions operate in the context of local lambdas and can't
+ // be called outside of that scope, only the arguments can be evaluated ahead of
+ // time.
private def childrenToRecurse(expr: Expression): Seq[Expression] = expr match {
case _: CodegenFallback => Nil
case c: ConditionalExpression => c.alwaysEvaluatedInputs.map(skipForShortcut)
+ case h: HigherOrderFunction => h.arguments
Review Comment:
do we need to do the same for `commonChildrenToRecurse`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] cloud-fan commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1198528181
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -130,6 +134,23 @@ case class LambdaFunction(
override def eval(input: InternalRow): Any = function.eval(input)
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ val functionCode = function.genCode(ctx)
Review Comment:
The code here is only to optimize the non-nullable case, which should be handled by the function body codegen already.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1198914316
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -130,6 +134,23 @@ case class LambdaFunction(
override def eval(input: InternalRow): Any = function.eval(input)
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ val functionCode = function.genCode(ctx)
Review Comment:
I'm trying to wrap my ahead around this again. I _think_ they might be doing different things? Like for `ArrayTransform`, the `nullSafeCodeGen` is only optimizing the non-nullable case of the result of `ArrayTransform`, whereas this is optimizing the non-nullable case of the result of the lambda function being called for each element in the `ArrayTransform`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "chris-twiner (via GitHub)" <gi...@apache.org>.
chris-twiner commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-2404981056
> @Kimahriman just out of curiosity, how much did the performance improve?
I just wanted to add to the above response that I've implemented a compilation scheme [here](https://sparkutils.github.io/quality/latest/advanced/userFunctions/#controlling-compilation-tweaking-the-quality-optimisations), as part of Quality, and we saw perf boosts of up to 40%, after that adding further lambdas triggered the cost of code generation being higher than the saving. It's definitely usage dependant though, the more work done in the function the higher the cost, a small boost is noticeable on removal of the atomic under similar ideal circumstances.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3285649179
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -886,6 +1170,114 @@ case class ArrayAggregate(
}
}
+ protected def nullSafeCodeGen(
+ ctx: CodegenContext,
+ ev: ExprCode,
+ f: String => String): ExprCode = {
+ val argumentGen = argument.genCode(ctx)
+ val resultCode = f(argumentGen.value)
+
+ if (nullable) {
+ val nullSafeEval = ctx.nullSafeExec(argument.nullable, argumentGen.isNull)(resultCode)
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |boolean ${ev.isNull} = ${argumentGen.isNull};
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$nullSafeEval
+ """)
+ } else {
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$resultCode
+ """, isNull = FalseLiteral)
+ }
+ }
+
+ protected def assignVar(
+ varCode: ExprCode,
+ atomicVar: String,
+ value: String,
+ isNull: String,
+ nullable: Boolean): String = {
+ val atomicAssign = assignAtomic(atomicVar, value, isNull, nullable)
+ if (nullable) {
+ s"""
+ ${varCode.value} = $value;
+ ${varCode.isNull} = $isNull;
+ $atomicAssign
+ """
+ } else {
+ s"""
+ ${varCode.value} = $value;
+ $atomicAssign
+ """
+ }
+ }
+
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar, accForMergeVar, accForFinishVar), varCodes => {
+ val Seq(elementCode, accForMergeCode, accForFinishCode) = varCodes
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val i = ctx.freshName("i")
+
+ val zeroCode = zero.genCode(ctx)
+ val mergeCode = merge.genCode(ctx)
+ val finishCode = finish.genCode(ctx)
+
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+ val mergeAtomic = ctx.addReferenceObj(accForMergeVar.name,
+ accForMergeVar.value)
+ val finishAtomic = ctx.addReferenceObj(accForFinishVar.name,
+ accForFinishVar.value)
+
+ val mergeJavaType = CodeGenerator.javaType(accForMergeVar.dataType)
+ val finishJavaType = CodeGenerator.javaType(accForFinishVar.dataType)
+
+ // Some expressions return internal buffers that we have to copy
+ val mergeCopy = if (CodeGenerator.isPrimitiveType(merge.dataType)) {
+ s"${mergeCode.value}"
+ } else {
+ s"($mergeJavaType)InternalRow.copyValue(${mergeCode.value})"
+ }
+
+ val nullCheck = if (nullable) {
+ s"${ev.isNull} = ${finishCode.isNull};"
+ } else {
+ ""
+ }
+
+ val initialAssignment = assignVar(accForMergeCode, mergeAtomic, zeroCode.value,
+ zeroCode.isNull, zero.nullable)
+
+ val mergeAssignment = assignVar(accForMergeCode, mergeAtomic, mergeCopy,
+ mergeCode.isNull, merge.nullable)
+
+ val finishAssignment = assignVar(accForFinishCode, finishAtomic, accForMergeCode.value,
+ accForMergeCode.isNull, merge.nullable)
Review Comment:
`finishAssignment` is using `merge.nullable` to decide whether to propagate the accumulator null bit into the finish lambda. That loses a nullable `zero` on empty arrays when `merge` itself is non-nullable. For example, with `aggregate(array(), CAST(NULL AS INT), (acc, x) -> coalesce(acc, 0) + x, acc -> acc IS NULL)`, interpreted eval passes `NULL` into `finish` and returns `true`, but the generated path leaves `accForFinishCode.isNull` at its default `false` and can return `false`. This should follow the accumulator nullability here, not `merge.nullable`.
_[ :robot: posted by Codex on behalf of sunchao using the code-review-for-me skill :robot: ]_
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3293465912
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -622,6 +753,72 @@ case class ArrayFilter(
new GenericArrayData(buffer)
}
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar) ++ indexVar, varCodes => {
+ val elementCode = varCodes.head
+ val indexCode = varCodes.tail.headOption
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val count = ctx.freshName("count")
+ val arrayTracker = ctx.freshName("arrayTracker")
+ val arrayData = ctx.freshName("arrayData")
+ val i = ctx.freshName("i")
+ val j = ctx.freshName("j")
+
+ val arrayType = dataType.asInstanceOf[ArrayType]
+
+ val trackerInit = CodeGenerator.createArrayData(
+ arrayTracker, BooleanType, numElements, s" $prettyName failed.")
+ val resultInit = CodeGenerator.createArrayData(
+ arrayData, arrayType.elementType, count, s" $prettyName failed.")
+
+ val functionCode = function.genCode(ctx)
+
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+ val indexAssignment = indexCode.map(c => assignIndex(ctx, c, indexVar.get, i))
+ val varAssignments = (Seq(elementAssignment) ++ indexAssignment).mkString("\n")
+
+ val resultAssignment = CodeGenerator.setArrayElement(arrayTracker, BooleanType,
+ i, functionCode.value, isNull = None)
+
+ val getTrackerValue = CodeGenerator.getValue(arrayTracker, BooleanType, i)
+ val copy = CodeGenerator.createArrayAssignment(arrayData, arrayType.elementType, arg,
+ j, i, arrayType.containsNull)
+
+ // This takes a two passes to avoid evaluating the predicate multiple times
+ // The first pass evaluates each element in the array, tracks how many elements
+ // returned true, and tracks the result of each element in a boolean array `arrayTracker`.
+ // The second pass copies elements from the original array to the new array created
+ // based on the number of elements matching the first pass.
+
+ s"""
+ |final int $numElements = $arg.numElements();
+ |$trackerInit
+ |int $count = 0;
+ |for (int $i = 0; $i < $numElements; $i++) {
+ | $varAssignments
+ | ${functionCode.code}
+ | $resultAssignment
+ | if ((boolean)${functionCode.value}) {
Review Comment:
Fixed in `eecc672cb44`.
Generated `ArrayFilter` now computes a local `keep` predicate as `!functionCode.isNull && functionCode.value` and uses that value for both the tracker array and the retained element count. I added a Catalyst regression with custom `GenericArrayData` that keeps `isNullAt=true` while exposing a stale typed boolean value, which fails against the old generated path.
Verified with:
```bash
build/sbt 'catalyst/testOnly org.apache.spark.sql.catalyst.expressions.HigherOrderFunctionsSuite -- -z "generated code"'
```
Result: passed (`2` tests, `0` failures).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1157743929
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -101,6 +101,14 @@ case class NamedLambdaVariable(
override def simpleString(maxFields: Int): String = {
s"lambda $name#${exprId.id}: ${dataType.simpleString(maxFields)}"
}
+
+ // We need to include the Expr ID in the Codegen variable name since several tests bypass
+ // `UnresolvedNamedLambdaVariable.freshVarName`
+ lazy val variableName = s"${name}_${exprId.id}"
Review Comment:
Hmm must have thought it was used in the actual generated code or something but doesn't look like it is. In fact I don't think `variableName` is really needed at all, can just use `exprId.id` as the map key, and `name` for the code comment
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -350,6 +445,49 @@ case class ArrayTransform(
result
}
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar) ++ indexVar, { lambdaExprs =>
Review Comment:
I think I did that because of the instances where I used the case expression to pull out the single element list, will update these for consistency
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1239152129
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -781,6 +1007,49 @@ case class ArrayForAll(
}
}
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar), { case Seq(elementCode) =>
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val forall = ctx.freshName("forall")
+ val foundNull = ctx.freshName("foundNull")
+ val i = ctx.freshName("i")
+
+ val functionCode = function.genCode(ctx)
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+
+ val nullCheck = if (nullable) {
+ s"""
+ if ($forall && $foundNull) {
+ ${ev.isNull} = true;
+ }
+ """
+ } else {
+ ""
+ }
+
+ s"""
+ |final int $numElements = ${arg}.numElements();
+ |boolean $forall = true;
+ |boolean $foundNull = false;
+ |int $i = 0;
+ |while ($i < $numElements && $forall) {
+ | $elementAssignment
+ | ${functionCode.code}
+ | if (${functionCode.isNull}) {
+ | $foundNull = true;
+ | } else if (!${functionCode.value}) {
+ | $forall = false;
+ | }
+ | $i++;
+ |}
+ |$nullCheck
+ |${ev.value} = $forall;
+ """.stripMargin
+ })
+ })
+ }
+
Review Comment:
> Looks like a copy and paste of `exists`, doesn't it? Can we have a parent class for some sharing? Unless I'm mistaken, the generated code block is the exact copy except `$forall` (vs `$exists`).
I'll have to look at this a little bit. It's tricky because there's a few places they are different (default true vs false, !'d for one and not the other, the `followThreeValuedLogic` flag for exists)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] Kimahriman commented on a diff in pull request #34558: [SPARK-37019][SQL] Add codegen support to array higher-order functions
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r1239152129
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -781,6 +1007,49 @@ case class ArrayForAll(
}
}
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar), { case Seq(elementCode) =>
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val forall = ctx.freshName("forall")
+ val foundNull = ctx.freshName("foundNull")
+ val i = ctx.freshName("i")
+
+ val functionCode = function.genCode(ctx)
+ val elementAssignment = assignArrayElement(ctx, arg, elementCode, elementVar, i)
+
+ val nullCheck = if (nullable) {
+ s"""
+ if ($forall && $foundNull) {
+ ${ev.isNull} = true;
+ }
+ """
+ } else {
+ ""
+ }
+
+ s"""
+ |final int $numElements = ${arg}.numElements();
+ |boolean $forall = true;
+ |boolean $foundNull = false;
+ |int $i = 0;
+ |while ($i < $numElements && $forall) {
+ | $elementAssignment
+ | ${functionCode.code}
+ | if (${functionCode.isNull}) {
+ | $foundNull = true;
+ | } else if (!${functionCode.value}) {
+ | $forall = false;
+ | }
+ | $i++;
+ |}
+ |$nullCheck
+ |${ev.value} = $forall;
+ """.stripMargin
+ })
+ })
+ }
+
Review Comment:
I'll have to look at this a little bit. It's tricky because there's a few places they are different (default true vs false, !'d for one and not the other, the `followThreeValuedLogic` flag for exists)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-2736265492
I added a simple benchmark, local results:
```
[info] Running benchmark: transform
[info] Running case: codegen
[info] Stopped after 10 iterations, 4605 ms
[info] Running case: interpreted
[info] Stopped after 10 iterations, 16947 ms
[info] OpenJDK 64-Bit Server VM 17.0.12+0 on Mac OS X 14.7.4
[info] Apple M1 Max
[info] transform: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
[info] ------------------------------------------------------------------------------------------------------------------------
[info] codegen 371 461 54 26.9 37.1 1.0X
[info] interpreted 1615 1695 86 6.2 161.5 0.2X
[info] Running benchmark: filter
[info] Running case: codegen
[info] Stopped after 10 iterations, 4042 ms
[info] Running case: interpreted
[info] Stopped after 10 iterations, 16382 ms
[info] OpenJDK 64-Bit Server VM 17.0.12+0 on Mac OS X 14.7.4
[info] Apple M1 Max
[info] filter: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
[info] ------------------------------------------------------------------------------------------------------------------------
[info] codegen 360 404 45 27.8 36.0 1.0X
[info] interpreted 1589 1638 53 6.3 158.9 0.2X
[info] Running benchmark: forall - fast
[info] Running case: codegen
[info] Stopped after 10 iterations, 2129 ms
[info] Running case: interpreted
[info] Stopped after 10 iterations, 7804 ms
[info] OpenJDK 64-Bit Server VM 17.0.12+0 on Mac OS X 14.7.4
[info] Apple M1 Max
[info] forall - fast: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
[info] ------------------------------------------------------------------------------------------------------------------------
[info] codegen 199 213 9 50.4 19.9 1.0X
[info] interpreted 749 780 24 13.4 74.9 0.3X
[info] Running benchmark: forall - slow
[info] Running case: codegen
[info] Stopped after 10 iterations, 2570 ms
[info] Running case: interpreted
[info] Stopped after 10 iterations, 8253 ms
[info] OpenJDK 64-Bit Server VM 17.0.12+0 on Mac OS X 14.7.4
[info] Apple M1 Max
[info] forall - slow: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
[info] ------------------------------------------------------------------------------------------------------------------------
[info] codegen 240 257 17 41.7 24.0 1.0X
[info] interpreted 801 825 23 12.5 80.1 0.3X
[info] Running benchmark: exists - fast
[info] Running case: codegen
[info] Stopped after 10 iterations, 2061 ms
[info] Running case: interpreted
[info] Stopped after 10 iterations, 8190 ms
[info] OpenJDK 64-Bit Server VM 17.0.12+0 on Mac OS X 14.7.4
[info] Apple M1 Max
[info] exists - fast: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
[info] ------------------------------------------------------------------------------------------------------------------------
[info] codegen 188 206 21 53.2 18.8 1.0X
[info] interpreted 794 819 25 12.6 79.4 0.2X
[info] Running benchmark: exists - slow
[info] Running case: codegen
[info] Stopped after 10 iterations, 2441 ms
[info] Running case: interpreted
[info] Stopped after 10 iterations, 8876 ms
[info] OpenJDK 64-Bit Server VM 17.0.12+0 on Mac OS X 14.7.4
[info] Apple M1 Max
[info] exists - slow: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
[info] ------------------------------------------------------------------------------------------------------------------------
[info] codegen 232 244 14 43.1 23.2 1.0X
[info] interpreted 857 888 32 11.7 85.7 0.3X
[info] Running benchmark: aggregate
[info] Running case: codegen
[info] Stopped after 10 iterations, 4029 ms
[info] Running case: interpreted
[info] Stopped after 10 iterations, 12054 ms
[info] OpenJDK 64-Bit Server VM 17.0.12+0 on Mac OS X 14.7.4
[info] Apple M1 Max
[info] aggregate: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
[info] ------------------------------------------------------------------------------------------------------------------------
[info] codegen 381 403 18 26.3 38.1 1.0X
[info] interpreted 1157 1205 34 8.6 115.7 0.3X
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #34558:
URL: https://github.com/apache/spark/pull/34558#issuecomment-3001086535
> The subexpression elimination option is huge! Very exciting
https://github.com/apache/spark/pull/51272
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3293466169
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala:
##########
@@ -146,9 +146,13 @@ class EquivalentExpressions(
// There are some special expressions that we should not recurse into all of its children.
// 1. CodegenFallback: it's children will not be used to generate code (call eval() instead)
// 2. ConditionalExpression: use its children that will always be evaluated.
+ // 3. HigherOrderFunction: lambda functions operate in the context of local lambdas and can't
+ // be called outside of that scope, only the arguments can be evaluated ahead of
+ // time.
private def childrenToRecurse(expr: Expression): Seq[Expression] = expr match {
case _: CodegenFallback => Nil
case c: ConditionalExpression => c.alwaysEvaluatedInputs.map(skipForShortcut)
+ case h: HigherOrderFunction => h.arguments
Review Comment:
Fixed in `eecc672cb44`.
I changed HOF CSE recursion to use an `alwaysEvaluatedArguments` hook. By default it returns `arguments`, but `ArrayAggregate` overrides it to return only the input array argument, so `zero` is no longer made CSE-visible before the null-array guard.
I added the ANSI/CSE regression from this comment and verified it fails with the old `h.arguments` recursion, then passes with the fix.
Verified with:
```bash
build/sbt 'sql/testOnly org.apache.spark.sql.DataFrameFunctionsSuite -- -z "aggregate function - null array does not evaluate zero expression through CSE"'
```
Result: passed (`1` test, `0` failures).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3303960462
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -354,6 +447,49 @@ case class ArrayTransform(
result
}
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar) ++ indexVar, varCodes => {
+ val elementCode = varCodes.head
+ val indexCode = varCodes.tail.headOption
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val arrayData = ctx.freshName("arrayData")
+ val i = ctx.freshName("i")
+
+ val initialization = CodeGenerator.createArrayData(
+ arrayData, dataType.elementType, numElements, s" $prettyName failed.")
+
+ val functionCode = function.genCode(ctx)
Review Comment:
Fixed in `fe004a5435e`.
The generated HOF paths now mirror interpreted evaluation by generating from the rebound lambda body (`functionForEval` / `functionsForEval`) instead of the original lambda tree. This preserves the existing exprId-based lambda variable rebinding semantics while keeping the change scoped to HOF codegen.
I added a Catalyst regression that constructs a resolved `ArrayTransform` lambda where a `CodegenFallbackExpr` references a detached `NamedLambdaVariable` with the same `exprId` as the bound argument. The test failed before the fix with the generated-path NPE described in this comment, and passes after the fix.
Verified with:
```bash
build/sbt 'catalyst/testOnly org.apache.spark.sql.catalyst.expressions.HigherOrderFunctionsSuite -- -z "ArrayTransform codegen rebinds fallback lambda variables by expression ID"'
build/sbt 'catalyst/testOnly org.apache.spark.sql.catalyst.expressions.HigherOrderFunctionsSuite -- -z "ArrayTransform"'
```
Results: focused regression passed (`1` test), ArrayTransform slice passed (`3` tests).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3299655506
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -354,6 +447,49 @@ case class ArrayTransform(
result
}
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar) ++ indexVar, varCodes => {
+ val elementCode = varCodes.head
+ val indexCode = varCodes.tail.headOption
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val arrayData = ctx.freshName("arrayData")
+ val i = ctx.freshName("i")
+
+ val initialization = CodeGenerator.createArrayData(
+ arrayData, dataType.elementType, numElements, s" $prettyName failed.")
+
+ val functionCode = function.genCode(ctx)
Review Comment:
[P2] Rebind fallback lambda variables before code generation
The generated path uses the original lambda tree here, while interpreted evaluation deliberately uses `functionsForEval` to replace separately instantiated `NamedLambdaVariable`s with the bound argument instance by `exprId`. As a result, a valid resolved lambda containing a `CodegenFallback` expression can succeed interpreted but fail under codegen: I reproduced `ArrayTransform(array(1, 2, 3), LambdaFunction(CodegenFallbackExpr(detachedArg + 1), Seq(arg)))`, where `detachedArg` has the same `exprId` as `arg` but a different `AtomicReference`; `checkEvaluation` passes interpreted evaluation and then throws `NullPointerException` in `GeneratedClass$SpecificMutableProjection.apply`. Please generate the rebound function tree (and keep fallback atomic state synchronized) so codegen preserves the existing lambda-binding semantics.
_[ :robot: posted by Codex on behalf of sunchao using the code-review-for-me skill :robot: ]_
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-37019][SQL] Add codegen support to array higher-order functions [spark]
Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #34558:
URL: https://github.com/apache/spark/pull/34558#discussion_r3364790060
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala:
##########
@@ -886,6 +1179,115 @@ case class ArrayAggregate(
}
}
+ protected def nullSafeCodeGen(
+ ctx: CodegenContext,
+ ev: ExprCode,
+ f: String => String): ExprCode = {
+ val argumentGen = argument.genCode(ctx)
+ val resultCode = f(argumentGen.value)
+
+ if (nullable) {
+ val nullSafeEval = ctx.nullSafeExec(argument.nullable, argumentGen.isNull)(resultCode)
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |boolean ${ev.isNull} = ${argumentGen.isNull};
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$nullSafeEval
+ """)
+ } else {
+ ev.copy(code = code"""
+ |${argumentGen.code}
+ |${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+ |$resultCode
+ """, isNull = FalseLiteral)
+ }
+ }
+
+ protected def assignVar(
+ varCode: ExprCode,
+ atomicVar: String,
+ value: String,
+ isNull: String,
+ nullable: Boolean): String = {
+ val atomicAssign = assignAtomic(atomicVar, value, isNull, nullable)
+ if (nullable) {
+ s"""
+ ${varCode.value} = $value;
+ ${varCode.isNull} = $isNull;
+ $atomicAssign
+ """
+ } else {
+ s"""
+ ${varCode.value} = $value;
+ $atomicAssign
+ """
+ }
+ }
+
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ ctx.withLambdaVars(Seq(elementVar, accForMergeVar, accForFinishVar), varCodes => {
+ val Seq(elementCode, accForMergeCode, accForFinishCode) = varCodes
+
+ nullSafeCodeGen(ctx, ev, arg => {
+ val numElements = ctx.freshName("numElements")
+ val i = ctx.freshName("i")
+
+ val zeroCode = zero.genCode(ctx)
+ val Seq(mergeForCodegen, finishForCodegen) = functionsForEval
+ val mergeCode = mergeForCodegen.genCode(ctx)
+ val finishCode = finishForCodegen.genCode(ctx)
Review Comment:
Thanks, fixed in the latest push by binding the aggregate accumulator lambda variables with `zero.dataType.asNullable -> true`, so generated finish/merge code sees widened nested nullability. I also added regression coverage for both array element nullability and map value nullability across `CODEGEN_ONLY` and `NO_CODEGEN`.
Validation: `git diff --check`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org