[SPARK-24361][SQL] Polish code block manipulation API #21405

viirya · 2018-05-23T03:26:37Z

What changes were proposed in this pull request?

Current code block manipulation API is immature and hacky. We need a formal API to manipulate code blocks.

The basic idea is making JavaCode as TreeNode. So we can use familiar transform API to manipulate code blocks and expressions in code blocks.

For example, we can replace SimpleExprValue in a code block like this:

code.transformExprValues {
  case SimpleExprValue("1 + 1", _) => aliasedParam
}

The example use case is splitting code to methods.

For example, we have an ExprCode containing generated code. But it is too long and we need to split it as method. Because statement-based expressions can't be directly passed into. We need to transform them as variables first:

def getExprValues(block: Block): Set[ExprValue] = block match {
  case c: CodeBlock =>
    c.blockInputs.collect {
      case e: ExprValue => e
    }.toSet
  case _ => Set.empty
}

def currentCodegenInputs(ctx: CodegenContext): Set[ExprValue] = {
  // Collects current variables in ctx.currentVars and ctx.INPUT_ROW.
  // It looks roughly like...
  ctx.currentVars.flatMap { v =>
    getExprValues(v.code) ++ Set(v.value, v.isNull)
  }.toSet + ctx.INPUT_ROW
}

// A code block of an expression contains too long code, making it as method
if (eval.code.length > 1024) {
  val setIsNull = if (!eval.isNull.isInstanceOf[LiteralValue]) {
    ...
  } else {
    ""
  }

  // Pick up variables and statements necessary to pass in.
  val currentVars = currentCodegenInputs(ctx)
  val varsPassIn = getExprValues(eval.code).intersect(currentVars)
  val aliasedExprs = HashMap.empty[SimpleExprValue, VariableValue]

  // Replace statement-based expressions which can't be directly passed in the method.
  val newCode = eval.code.transform {
    case block =>
      block.transformExprValues {
        case s: SimpleExprValue(_, javaType) if varsPassIn.contains(s) =>
          if (aliasedExprs.contains(s)) {
            aliasedExprs(s)
          } else {
            val aliasedVariable = JavaCode.variable(ctx.freshName("aliasedVar"), javaType)
            aliasedExprs += s -> aliasedVariable
            varsPassIn += aliasedVariable
            aliasedVariable
          }
      }
  }

  val params = varsPassIn.filter(!_.isInstanceOf[SimpleExprValue])).map { variable =>
    s"${variable.javaType.getName} ${variable.variableName}"
  }.mkString(", ")

  val funcName = ctx.freshName("nodeName")
  val javaType = CodeGenerator.javaType(dataType)
  val newValue = JavaCode.variable(ctx.freshName("value"), dataType)
  val funcFullName = ctx.addNewFunction(funcName,
    s"""
      |private $javaType $funcName($params) {
      |  $newCode
      |  $setIsNull
      |  return ${eval.value};
      |}
    """.stripMargin))

  eval.value = newValue
  val args = varsPassIn.filter(!_.isInstanceOf[SimpleExprValue])).map { variable =>
    s"${variable.variableName}"
  }

  // Create a code block to assign statements to aliased variables.
  val createVariables = aliasedExprs.foldLeft(EmptyBlock) { (block, (statement, variable)) =>
    block + code"${statement.javaType.getName} $variable = $statement;"
  }
  eval.code = createVariables + code"$javaType $newValue = $funcFullName($args);"
}

How was this patch tested?

Added unite tests.

…for constructing code block.

…ay not aware of necessary block inputs.

viirya · 2018-05-23T03:28:25Z

cc @cloud-fan @hvanhovell @kiszk @mgaido91 @maropu @rednaxelafx @ueshin

SparkQA · 2018-05-23T07:05:02Z

Test build #91014 has finished for PR 21405 at commit 2d77332.

This patch fails due to an unknown error code, -9.
This patch merges cleanly.
This patch adds no public classes.

viirya · 2018-05-23T07:35:45Z

retest this please.

SparkQA · 2018-05-23T10:24:58Z

Test build #91021 has finished for PR 21405 at commit 2d77332.

This patch fails Spark unit tests.
This patch merges cleanly.
This patch adds no public classes.

viirya · 2018-05-23T10:27:57Z

retest this please.

kiszk · 2018-05-23T12:36:01Z

...talyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeBlockSuite.scala

+
+    // We want to replace all occurrences of `expr` with the variable `aliasedParam`.
+    val aliasedCode = code.transformExprValues {
+      case SimpleExprValue("1 + 1", _) => aliasedParam


nit: I know the current code works correctly. How about replacing _ with CodeGenerator.javaClass(IntegerType)?

mgaido91 · 2018-05-23T13:48:11Z

@viirya sorry do you already have some use-cases, goals in mind for this API? May you please explain me them? Thanks.

viirya · 2018-05-23T14:07:07Z

@mgaido91 Yes, as @rednaxelafx said in previous RP, #19813 will be the first use case. I'd like to use this API to transform statement-based expressions in code blocks.

For example, code contains Java code of an expression. We find it is too long and we need to split it out as function. Because statement-based expressions can't be directly passed into. We need to transform them as variables first. I think the usage roughly would looks like:

val collectedStatements = HashMap[ExprValue, ExprValue]
val transformed = code.transformExprValues {
  case s @ SimpleExprValue(_, javaType) =>
    if (collectedStatements.contains(s)) {
      collectedStatements(s)
    } else {
      val aliasedVariable = JavaCode.variable(ctx.freshName("var"), javaType)
      collectedStatements += s -> aliasedVariable
      aliasedVariable
    }
}
val createVariables = collectedStatements.foldLeft(EmptyBlock) { (block, (statement, variable)) =>
  block + code"${statement.javaType.getName} $variable = $statement;"
}
ev.copy(code = createVariables + transformed)

SparkQA · 2018-05-23T14:11:00Z

Test build #91031 has finished for PR 21405 at commit 2d77332.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

cloud-fan · 2018-05-23T15:26:08Z

can we give a concrete use case in the PR description? I think splitting code into methods is a good one, but we need to make it completed, e.g. how the method is generated.

viirya · 2018-05-24T04:18:51Z

@cloud-fan Thanks. I give a use case of splitting code into method in the PR description. I think it can show the basic idea and it is what I hope to make.

mgaido91 · 2018-05-24T08:07:53Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala

+   * Apply a map function to each java expression codes present in this java code, and return a new
+   * java code based on the mapped java expression codes.
+   */
+  def transformExprValues(f: PartialFunction[ExprValue, ExprValue]): this.type = {


am I wrong or we are not updating the exprValues here?

transformExprValues will create a new instance Block because a block is immutable. So once there are change, new block should take new exprValues.

I see, can we add some tests for this?

Ok. Will add them in next commit. Thanks.

Added test for it.

viirya · 2018-05-25T13:34:30Z

@cloud-fan I found that the way to collect method parameters in the use case can be simplified. Updated in the PR description. Please take a look when you have time. Thanks.

viirya · 2018-05-29T10:17:28Z

@cloud-fan @hvanhovell Do you have any comment/suggestion on this change? Thanks.

SparkQA · 2018-05-29T14:00:40Z

Test build #91251 has finished for PR 21405 at commit e30be7a.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

viirya · 2018-06-04T09:45:56Z

ping @cloud-fan @hvanhovell @kiszk

cloud-fan · 2018-07-05T03:58:31Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala

 * Trait representing an opaque fragments of java code.
 */
-trait JavaCode {
+trait JavaCode extends TreeNode[JavaCode] {


shall we only make Block extends TreeNode?

Ok. Currently ExprValue doesn't have to be TreeNode.

Will update in next commit.

cloud-fan · 2018-07-05T06:32:07Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala

  override def hashCode(): Int = value.hashCode() * 31 + javaType.hashCode()
 }

+case class LiteralExpr(override val value: String, override val javaType: Class[_])


do we still need this change?

cloud-fan · 2018-07-05T06:33:20Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala


-  // The expressions to be evaluated inside this block.
+  // All expressions to be evaluated inside this block and underlying blocks.
  def exprValues: Set[ExprValue]


do we still need it?

We can remove it, as discussed in other PR previously.

cloud-fan · 2018-07-05T06:38:36Z

...talyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeBlockSuite.scala

+
+  test ("transform expr in nested blocks") {
+    val expr = JavaCode.expression("1 + 1", IntegerType)
+    val isNull = JavaCode.isNullVariable("expr1_isNull")


not related to this PR, shall we remove isNullVariable and always use expression(..., BooleanType)? Or we can add a bunch of intExpression(...), booleanExpression(...) etc.

I think it is fine and maybe more clear. I was thinking the name of isNullVariable is a bit confusing at the first look actually.

cloud-fan · 2018-07-05T06:48:41Z

...talyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeBlockSuite.scala

+
+    val block = code"${subBlocks(0)}\n${subBlocks(1)}\n${subBlocks(2)}"
+    val transformedBlock = block.transform {
+      case b: Block => b.transformExprValues {


hmmm, I'd image transformExprValues will also transform child blocks.

so the this should be

block.transformExprValues { case SimpleExprValue ... }

Then I may call it transformAllExprValues.

SparkQA · 2018-07-05T07:05:01Z

Test build #92637 has finished for PR 21405 at commit ce8d9ed.

This patch fails due to an unknown error code, -9.
This patch merges cleanly.
This patch adds the following public classes (experimental):
trait JavaCode
trait Block extends TreeNode[Block] with JavaCode

cloud-fan · 2018-07-05T10:13:04Z

LGTM, can you update the demo in your PR description? thanks!

SparkQA · 2018-07-05T11:27:46Z

Test build #92640 has finished for PR 21405 at commit 707dd6f.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2018-07-05T11:40:15Z

Test build #92642 has finished for PR 21405 at commit c9dee44.

This patch passes all tests.
This patch merges cleanly.
This patch adds the following public classes (experimental):
class LiteralValue(val value: String, val javaType: Class[_]) extends ExprValue with Serializable

viirya · 2018-07-05T12:05:14Z

@cloud-fan Updated. Do you think I should address the comment #21405 (comment) to let transformExprValues also transforms child blocks?

cloud-fan · 2018-07-05T12:27:10Z

we can do it later, when we do need transformAllExprValues

cloud-fan · 2018-07-05T12:49:36Z

thanks, merging to master!

Looking forward to seeing how we can solve the method splitting problem with this infra :)

viirya added 21 commits April 30, 2018 13:29

Add API for handling expression code generation.

1df9943

Add new abstraction for expression codegen.

5fe425c

Add basic tests.

00bef6b

Merge remote-tracking branch 'upstream/master' into SPARK-24121

5d9c454

Deal merging conflict.

162deb2

Address comments and add more tests.

d138ee0

Address comments.

ee9a4c0

Remove JavaCode.block. We should always use code string interpolator …

e7cfa28

…for constructing code block.

We should not implicitly convert code block to string. Otherwise we m…

5945c15

…ay not aware of necessary block inputs.

Address comment and trim expected code string.

2b30654

Remove unused import.

aff411b

Address comments.

53b329a

Address some comments.

72faac3

Merge remote-tracking branch 'upstream/master' into SPARK-24121

ffbf4ab

Use code block for newly merged codegen.

d040676

Use Set as method exprValues method returning type.

c378ce2

Address comments.

2ca9741

Merge remote-tracking branch 'upstream/master' into SPARK-24121

d91f111

Merge remote-tracking branch 'upstream/master' into codeblock-api

5c40233

Add API to manipulate blocks and expressions.

8d0b1b9

Merge remote-tracking branch 'upstream/master' into codeblock-api

2d77332

kiszk reviewed May 23, 2018

View reviewed changes

mgaido91 reviewed May 24, 2018

View reviewed changes

Address comments.

e30be7a

viirya mentioned this pull request Jul 5, 2018

[SPARK-24635][SQL] Remove Blocks class from JavaCode class hierarchy #21619

Closed

cloud-fan reviewed Jul 5, 2018

View reviewed changes

viirya added 2 commits July 5, 2018 05:49

Merge remote-tracking branch 'upstream/master' into codeblock-api

21411f5

Only Block extends TreeNode.

ce8d9ed

cloud-fan reviewed Jul 5, 2018

View reviewed changes

viirya added 2 commits July 5, 2018 07:36

Remove exprValues property from Block.

707dd6f

Revert LiteralExpr.

c9dee44

asfgit closed this in 32cfd3e Jul 5, 2018

viirya deleted the codeblock-api branch December 27, 2023 18:21

[SPARK-24361][SQL] Polish code block manipulation API #21405

[SPARK-24361][SQL] Polish code block manipulation API #21405

Uh oh!

Conversation

viirya commented May 23, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

How was this patch tested?

Uh oh!

viirya commented May 23, 2018

Uh oh!

SparkQA commented May 23, 2018

Uh oh!

viirya commented May 23, 2018

Uh oh!

SparkQA commented May 23, 2018

Uh oh!

viirya commented May 23, 2018

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mgaido91 commented May 23, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

viirya commented May 23, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

SparkQA commented May 23, 2018

Uh oh!

cloud-fan commented May 23, 2018

Uh oh!

viirya commented May 24, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya commented May 25, 2018

Uh oh!

viirya commented May 29, 2018

Uh oh!

SparkQA commented May 29, 2018

Uh oh!

viirya commented Jun 4, 2018

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya Jul 5, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya Jul 5, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya commented May 23, 2018 •

edited

Loading

mgaido91 commented May 23, 2018 •

edited

Loading

viirya commented May 23, 2018 •

edited

Loading

viirya commented May 24, 2018 •

edited

Loading

viirya Jul 5, 2018 •

edited

Loading

viirya Jul 5, 2018 •

edited

Loading