gengliangwang commented on a change in pull request #32060:
URL: https://github.com/apache/spark/pull/32060#discussion_r611140191



##########
File path: 
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
##########
@@ -77,8 +82,54 @@ object CurrentOrigin {
 // A tag of a `TreeNode`, which defines name and type
 case class TreeNodeTag[T](name: String)
 
+// A wrapper of BitSet for pattern enums.
+trait TreePatternBits {
+  protected val treePatternBits: BitSet
+
+  /**
+   * @param t, the tree pattern enum to be tested.
+   * @return true if the bit for `t` is set; false otherwise.
+   */
+  @inline final def containsPattern(t: TreePattern): Boolean = {
+    treePatternBits.get(t.id)
+  }
+
+  /**
+   * @param patterns, a sequence of tree pattern enums to be tested.
+   * @return true if every bit for `patterns` is set; false otherwise.
+   */
+  final def containsAllPatterns(patterns: TreePattern*): Boolean = {
+    val iterator = patterns.iterator
+    while (iterator.hasNext) {
+      if (!containsPattern(iterator.next)) {
+        return false
+      }
+    }
+    true
+  }
+
+  /**
+   * @param patterns, a sequence of tree pattern enums to be tested.
+   * @return true if at least one bit for `patterns` is set; false otherwise.
+   */
+  final def containsAnyPattern(patterns: TreePattern*): Boolean = {
+    val iterator = patterns.iterator
+    while (iterator.hasNext) {
+      if (containsPattern(iterator.next)) {
+        return true
+      }
+    }
+    false
+  }
+}
+
+// A functor that always returns true.
+object AlwaysProcess{
+  val fn: TreePatternBits => Boolean = { _ => true}
+}
+
 // scalastyle:off
-abstract class TreeNode[BaseType <: TreeNode[BaseType]] extends Product {
+abstract class TreeNode[BaseType <: TreeNode[BaseType]] extends Product with 
TreePatternBits{

Review comment:
       nit: one space after `TreePatternBits`

##########
File path: 
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
##########
@@ -77,8 +82,54 @@ object CurrentOrigin {
 // A tag of a `TreeNode`, which defines name and type
 case class TreeNodeTag[T](name: String)
 
+// A wrapper of BitSet for pattern enums.
+trait TreePatternBits {
+  protected val treePatternBits: BitSet
+
+  /**
+   * @param t, the tree pattern enum to be tested.
+   * @return true if the bit for `t` is set; false otherwise.
+   */
+  @inline final def containsPattern(t: TreePattern): Boolean = {
+    treePatternBits.get(t.id)
+  }
+
+  /**
+   * @param patterns, a sequence of tree pattern enums to be tested.
+   * @return true if every bit for `patterns` is set; false otherwise.
+   */
+  final def containsAllPatterns(patterns: TreePattern*): Boolean = {
+    val iterator = patterns.iterator
+    while (iterator.hasNext) {
+      if (!containsPattern(iterator.next)) {
+        return false
+      }
+    }
+    true
+  }
+
+  /**
+   * @param patterns, a sequence of tree pattern enums to be tested.
+   * @return true if at least one bit for `patterns` is set; false otherwise.
+   */
+  final def containsAnyPattern(patterns: TreePattern*): Boolean = {
+    val iterator = patterns.iterator
+    while (iterator.hasNext) {
+      if (containsPattern(iterator.next)) {
+        return true
+      }
+    }
+    false
+  }
+}
+
+// A functor that always returns true.
+object AlwaysProcess{

Review comment:
       nit: one space after AlwaysProcess
   
   

##########
File path: 
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
##########
@@ -77,8 +82,54 @@ object CurrentOrigin {
 // A tag of a `TreeNode`, which defines name and type
 case class TreeNodeTag[T](name: String)
 
+// A wrapper of BitSet for pattern enums.
+trait TreePatternBits {

Review comment:
       Nit: how about putting TreePatternBits in a new file. There are already 
1k lines code in `TreeNode.scala`




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]



---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to