Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[SPARK-9165][SQL] codegen for CreateArray, CreateStruct and CreateNamedStruct #7537

Closed
wants to merge 3 commits into from
Closed

[SPARK-9165][SQL] codegen for CreateArray, CreateStruct and CreateNamedStruct #7537

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
39cefb8
codegen for createArray createStruct & createNamedStruct
yjshen Jul 20, 2015
5e90f0a
resolve comments: classOf
yjshen Jul 21, 2015
3a6dce6
use infix notion in createArray test
yjshen Jul 22, 2015
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 60 additions & 5 deletions ...atalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,18 @@

package org.apache.spark.sql.catalyst.expressions

import scala.collection.mutable

import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
import org.apache.spark.sql.catalyst.expressions.codegen.{GeneratedExpressionCode, CodeGenContext}
import org.apache.spark.sql.catalyst.util.TypeUtils
import org.apache.spark.sql.types._

/**
* Returns an Array containing the evaluation of all children expressions.
*/
case class CreateArray(children: Seq[Expression]) extends Expression with CodegenFallback {
case class CreateArray(children: Seq[Expression]) extends Expression {

override def foldable: Boolean = children.forall(_.foldable)

Expand All @@ -45,14 +47,31 @@ case class CreateArray(children: Seq[Expression]) extends Expression with Codege
children.map(_.eval(input))
}

override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
val arraySeqClass = classOf[mutable.ArraySeq[Any]].getName
s"""
boolean ${ev.isNull} = false;
$arraySeqClass<Object> ${ev.primitive} = new $arraySeqClass<Object>(${children.size});
""" +
children.zipWithIndex.map { case (e, i) =>
val eval = e.gen(ctx)
eval.code + s"""
if (${eval.isNull}) {
${ev.primitive}.update($i, null);
} else {
${ev.primitive}.update($i, ${eval.primitive});
}
"""
}.mkString("\n")
}

override def prettyName: String = "array"
}

/**
* Returns a Row containing the evaluation of all children expressions.
* TODO: [[CreateStruct]] does not support codegen.
*/
case class CreateStruct(children: Seq[Expression]) extends Expression with CodegenFallback {
case class CreateStruct(children: Seq[Expression]) extends Expression {

override def foldable: Boolean = children.forall(_.foldable)

Expand All @@ -76,6 +95,24 @@ case class CreateStruct(children: Seq[Expression]) extends Expression with Codeg
InternalRow(children.map(_.eval(input)): _*)
}

override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
val rowClass = classOf[GenericMutableRow].getName
s"""
boolean ${ev.isNull} = false;
final $rowClass ${ev.primitive} = new $rowClass(${children.size});
""" +
children.zipWithIndex.map { case (e, i) =>
val eval = e.gen(ctx)
eval.code + s"""
if (${eval.isNull}) {
${ev.primitive}.update($i, null);
} else {
${ev.primitive}.update($i, ${eval.primitive});
}
"""
}.mkString("\n")
}

override def prettyName: String = "struct"
}

Expand All @@ -84,7 +121,7 @@ case class CreateStruct(children: Seq[Expression]) extends Expression with Codeg
*
* @param children Seq(name1, val1, name2, val2, ...)
*/
case class CreateNamedStruct(children: Seq[Expression]) extends Expression with CodegenFallback {
case class CreateNamedStruct(children: Seq[Expression]) extends Expression {

private lazy val (nameExprs, valExprs) =
children.grouped(2).map { case Seq(name, value) => (name, value) }.toList.unzip
Expand Down Expand Up @@ -122,5 +159,23 @@ case class CreateNamedStruct(children: Seq[Expression]) extends Expression with
InternalRow(valExprs.map(_.eval(input)): _*)
}

override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
val rowClass = classOf[GenericMutableRow].getName
s"""
boolean ${ev.isNull} = false;
final $rowClass ${ev.primitive} = new $rowClass(${valExprs.size});
""" +
valExprs.zipWithIndex.map { case (e, i) =>
val eval = e.gen(ctx)
eval.code + s"""
if (${eval.isNull}) {
${ev.primitive}.update($i, null);
} else {
${ev.primitive}.update($i, ${eval.primitive});
}
"""
}.mkString("\n")
}

override def prettyName: String = "named_struct"
}
16 changes: 16 additions & 0 deletions sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,22 @@ class ComplexTypeSuite extends SparkFunSuite with ExpressionEvalHelper {
checkEvaluation(getArrayStructFields(nullArrayStruct, "a"), null)
}

test("CreateArray") {
val intSeq = Seq(5, 10, 15, 20, 25)
val longSeq = intSeq.map(_.toLong)
val strSeq = intSeq.map(_.toString)
checkEvaluation(CreateArray(intSeq.map(Literal(_))), intSeq, EmptyRow)
checkEvaluation(CreateArray(longSeq.map(Literal(_))), longSeq, EmptyRow)
checkEvaluation(CreateArray(strSeq.map(Literal(_))), strSeq, EmptyRow)
Copy link
Contributor

Choose a reason for hiding this comment

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

Maybe add a test that includes null values?

Copy link
Member Author

Choose a reason for hiding this comment

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

added below.


val intWithNull = intSeq.map(Literal(_)) :+ Literal.create(null, IntegerType)
val longWithNull = longSeq.map(Literal(_)) :+ Literal.create(null, LongType)
val strWithNull = strSeq.map(Literal(_)) :+ Literal.create(null, StringType)
checkEvaluation(CreateArray(intWithNull), intSeq :+ null, EmptyRow)
checkEvaluation(CreateArray(longWithNull), longSeq :+ null, EmptyRow)
checkEvaluation(CreateArray(strWithNull), strSeq :+ null, EmptyRow)
}

test("CreateStruct") {
val row = create_row(1, 2, 3)
val c1 = 'a.int.at(0)
Expand Down