gengliangwang commented on a change in pull request #32060:
URL: https://github.com/apache/spark/pull/32060#discussion_r608359329
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/predicates.scala
##########
@@ -26,6 +26,7 @@ import
org.apache.spark.sql.catalyst.expressions.BindReferences.bindReference
import org.apache.spark.sql.catalyst.expressions.codegen._
import org.apache.spark.sql.catalyst.expressions.codegen.Block._
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LeafNode,
LogicalPlan, Project}
+import org.apache.spark.sql.catalyst.trees.TreePattern._
Review comment:
Nit:
`import org.apache.spark.sql.catalyst.trees.TreePattern.IN`
As per https://github.com/databricks/scala-style-guide#imports, we should
avoid using wildcard imports
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
##########
@@ -90,6 +139,66 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
extends Product {
*/
private val tags: mutable.Map[TreeNodeTag[_], Any] = mutable.Map.empty
+ /**
+ * A bit set of tree patterns for this TreeNode and its subtree. If this
TreeNode node and its
+ * subtree contains a pattern `P`, the corresponding bit for `P.id` is set
in this BitSet.
+ */
+ override lazy val treePatternBits: BitSet = {
+ val bits: BitSet = new BitSet(TreePattern.maxId)
+ // Propagate node pattern bits
+ val nodePatternIterator = nodePatterns.iterator
+ while (nodePatternIterator.hasNext) {
+ bits.set(nodePatternIterator.next().id)
+ }
+ // Propagate children's pattern bits
+ val childIterator = children.iterator
Review comment:
ditto
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
##########
@@ -48,6 +52,27 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]]
@transient
lazy val outputSet: AttributeSet = AttributeSet(output)
+ // Override `treePatternBits` to propagate bits for its expressions.
+ override lazy val treePatternBits: BitSet = {
+ val bits: BitSet = new BitSet(TreePattern.maxId)
+ // Propagate node pattern bits
+ val nodeTypeIterator = nodePatterns.iterator
Review comment:
nit: This is the same as
```
nodePatterns.foreach {...}
```
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
##########
@@ -77,8 +80,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 {
+ 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
Review comment:
ditto
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
##########
@@ -48,6 +52,27 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]]
@transient
lazy val outputSet: AttributeSet = AttributeSet(output)
+ // Override `treePatternBits` to propagate bits for its expressions.
+ override lazy val treePatternBits: BitSet = {
+ val bits: BitSet = new BitSet(TreePattern.maxId)
+ // Propagate node pattern bits
+ val nodeTypeIterator = nodePatterns.iterator
+ while (nodeTypeIterator.hasNext) {
+ bits.set(nodeTypeIterator.next().id)
+ }
+ // Propagate children's pattern bits
+ val childIterator = children.iterator
Review comment:
ditto
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
##########
@@ -90,6 +139,66 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
extends Product {
*/
private val tags: mutable.Map[TreeNodeTag[_], Any] = mutable.Map.empty
+ /**
+ * A bit set of tree patterns for this TreeNode and its subtree. If this
TreeNode node and its
+ * subtree contains a pattern `P`, the corresponding bit for `P.id` is set
in this BitSet.
+ */
+ override lazy val treePatternBits: BitSet = {
+ val bits: BitSet = new BitSet(TreePattern.maxId)
+ // Propagate node pattern bits
+ val nodePatternIterator = nodePatterns.iterator
Review comment:
ditto
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
##########
@@ -48,6 +52,27 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]]
@transient
lazy val outputSet: AttributeSet = AttributeSet(output)
+ // Override `treePatternBits` to propagate bits for its expressions.
+ override lazy val treePatternBits: BitSet = {
+ val bits: BitSet = new BitSet(TreePattern.maxId)
+ // Propagate node pattern bits
+ val nodeTypeIterator = nodePatterns.iterator
+ while (nodeTypeIterator.hasNext) {
+ bits.set(nodeTypeIterator.next().id)
+ }
+ // Propagate children's pattern bits
+ val childIterator = children.iterator
+ while (childIterator.hasNext) {
+ bits.union(childIterator.next().treePatternBits)
+ }
+ // Propagate expressions' pattern bits
+ val exprIterator = expressions.iterator
Review comment:
ditto
--
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]
