Add support for transformation of optional children.

vanzin · Jun 12, 2014 · 2ccb721 · 2ccb721
1 parent 1c04652
commit 2ccb721
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
@@ -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 =>
@@ -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 =>