Skip to content

Commit

Permalink
[SQL] Support transforming TreeNodes with Option children.
Browse files Browse the repository at this point in the history
Thanks goes to @marmbrus for his implementation.

Author: Michael Armbrust <[email protected]>
Author: Zongheng Yang <[email protected]>

Closes #1074 from concretevitamin/option-treenode and squashes the following commits:

ef27b85 [Zongheng Yang] Merge pull request #1 from marmbrus/pr/1074
73133c2 [Michael Armbrust] TreeNodes can't be inner classes.
ab78420 [Zongheng Yang] Add a test.
2ccb721 [Michael Armbrust] Add support for transformation of optional children.
  • Loading branch information
marmbrus committed Jun 15, 2014
1 parent 7dd9fc6 commit 269fc62
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,14 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] {
} else {
arg
}
case Some(arg: TreeNode[_]) if children contains arg =>
val newChild = arg.asInstanceOf[BaseType].transformDown(rule)
if (!(newChild fastEquals arg)) {
changed = true
Some(newChild)
} else {
Some(arg)
}
case m: Map[_,_] => m
case args: Traversable[_] => args.map {
case arg: TreeNode[_] if children contains arg =>
Expand Down Expand Up @@ -231,6 +239,14 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] {
} else {
arg
}
case Some(arg: TreeNode[_]) if children contains arg =>
val newChild = arg.asInstanceOf[BaseType].transformUp(rule)
if (!(newChild fastEquals arg)) {
changed = true
Some(newChild)
} else {
Some(arg)
}
case m: Map[_,_] => m
case args: Traversable[_] => args.map {
case arg: TreeNode[_] if children contains arg =>
Expand Down Expand Up @@ -273,7 +289,8 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] {
} catch {
case e: java.lang.IllegalArgumentException =>
throw new TreeNodeException(
this, s"Failed to copy node. Is otherCopyArgs specified correctly for $nodeName?")
this, s"Failed to copy node. Is otherCopyArgs specified correctly for $nodeName? "
+ s"Exception message: ${e.getMessage}.")
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,17 @@ import scala.collection.mutable.ArrayBuffer
import org.scalatest.FunSuite

import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.types.{StringType, NullType}

case class Dummy(optKey: Option[Expression]) extends Expression {
def children = optKey.toSeq
def references = Set.empty[Attribute]
def nullable = true
def dataType = NullType
override lazy val resolved = true
type EvaluatedType = Any
def eval(input: Row) = null.asInstanceOf[Any]
}

class TreeNodeSuite extends FunSuite {
test("top node changed") {
Expand Down Expand Up @@ -75,4 +86,20 @@ class TreeNodeSuite extends FunSuite {

assert(expected === actual)
}

test("transform works on nodes with Option children") {
val dummy1 = Dummy(Some(Literal("1", StringType)))
val dummy2 = Dummy(None)
val toZero: PartialFunction[Expression, Expression] = { case Literal(_, _) => Literal(0) }

var actual = dummy1 transformDown toZero
assert(actual === Dummy(Some(Literal(0))))

actual = dummy1 transformUp toZero
assert(actual === Dummy(Some(Literal(0))))

actual = dummy2 transform toZero
assert(actual === Dummy(None))
}

}

0 comments on commit 269fc62

Please sign in to comment.