Github user mgaido91 commented on a diff in the pull request: https://github.com/apache/spark/pull/21193#discussion_r186254876 --- Diff: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/javaCode.scala --- @@ -112,6 +112,112 @@ object JavaCode { def isNullExpression(code: String): SimpleExprValue = { expression(code, BooleanType) } + + def block(code: String): Block = { + CodeBlock(codeParts = Seq(code), blockInputs = Seq.empty) + } +} + +/** + * A trait representing a block of java code. + */ +trait Block extends JavaCode { + + // The expressions to be evaluated inside this block. + def exprValues: Seq[ExprValue] + + // This will be called during string interpolation. + override def toString: String = _marginChar match { + case Some(c) => code.stripMargin(c) + case _ => code + } + + var _marginChar: Option[Char] = None + + def stripMargin(c: Char): this.type = { + _marginChar = Some(c) + this + } + + def stripMargin: this.type = { + _marginChar = Some('|') + this + } + + def + (other: Block): Block +} + +object Block { + implicit def blockToString(block: Block): String = block.toString + + implicit def blocksToBlock(blocks: Seq[Block]): Block = Blocks(blocks) + + implicit class BlockHelper(val sc: StringContext) extends AnyVal { + def code(args: Any*): Block = { + sc.checkLengths(args) + if (sc.parts.length == 0) { + EmptyBlock + } else { + args.foreach { + case _: ExprValue => + case _: Int | _: Long | _: Float | _: Double | _: String => + case _: Block => + case other => throw new IllegalArgumentException( + s"Can not interpolate ${other.getClass.getName} into code block.") + } + CodeBlock(sc.parts, args) + } + } + } +} + +/** + * A block of java code. Including a sequence of code parts and some inputs to this block. + * The actual java code is generated by embedding the inputs into the code parts. + */ +case class CodeBlock(codeParts: Seq[String], blockInputs: Seq[Any]) extends Block { + override def exprValues: Seq[ExprValue] = { + blockInputs.flatMap { + case b: Block => b.exprValues + case e: ExprValue => Seq(e) + case _ => Seq.empty + } + } + + override def code: String = { + val strings = codeParts.iterator + val inputs = blockInputs.iterator + var buf = new StringBuffer(strings.next) + while (strings.hasNext) { + buf append inputs.next + buf append strings.next + } + buf.toString + } + + override def + (other: Block): Block = other match { + case c: CodeBlock => Blocks(Seq(this, c)) + case b: Blocks => Blocks(Seq(this) ++ b.blocks) + case EmptyBlock => this + } +} + +case class Blocks(blocks: Seq[Block]) extends Block { + override def exprValues: Seq[ExprValue] = blocks.flatMap(_.exprValues) + override def code: String = blocks.map(_.toString).mkString --- End diff -- +1 for @viirya 's answer
--- --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org