create a distinctSet for uniqueness constraint

Author: gatorsmile

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -1213,31 +1213,10 @@ object EliminateDistinct extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     // Eliminate the useless distinct.
     // Distinct has been replaced by Aggregate in the rule ReplaceDistinctWithAggregate
-    case a @ Aggregate(grouping, aggs, child) if isDistinct(a) && isDistinct(child) => child
-  }
-
-  // propagate the distinct property from the child
-  @tailrec
-  private def isDistinct(plan: LogicalPlan): Boolean = plan match {
-    // Distinct(left) or Aggregate(left.output, left.output, _) always returns distinct results
-    case _: Distinct => true
-    case Aggregate(grouping, aggs, _) if grouping.nonEmpty && grouping == aggs => true
-    // BinaryNode:
-    case p @ Join(_, _, LeftSemi, _) => isDistinct(p.left)
-    case p: Intersect => isDistinct(p.left)
-    case p: Except => isDistinct(p.left)
-    // UnaryNode:
-    case p: Project if p.child.outputSet.subsetOf(p.outputSet) => isDistinct(p.child)
-    case p: Aggregate if p.child.outputSet.subsetOf(p.outputSet) => isDistinct(p.child)
-    case p: Filter => isDistinct(p.child)
-    case p: GlobalLimit => isDistinct(p.child)
-    case p: LocalLimit => isDistinct(p.child)
-    case p: Sort => isDistinct(p.child)
-    case p: BroadcastHint => isDistinct(p.child)
-    case p: Sample => isDistinct(p.child)
-    case p: SubqueryAlias => isDistinct(p.child)
-    // Others:
-    case o => false
+    case a @ Aggregate(grouping, aggs, child)
+        if child.distinctSet.nonEmpty && child.distinctSet.subsetOf(AttributeSet(aggs)) &&
+          a.isForDistinct =>
+      child
   }
 }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala

@@ -101,6 +101,11 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]] extends TreeNode[PlanT
    */
   protected def validConstraints: Set[Expression] = Set.empty
 
+  /**
+   * The set of attributes whose combination can uniquely identify a row.
+   */
+  def distinctSet: AttributeSet = AttributeSet.empty
+
   /**
    * Returns the set of attributes that are output by this node.
    */

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala

@@ -51,6 +51,14 @@ case class Project(projectList: Seq[NamedExpression], child: LogicalPlan) extend
     !expressions.exists(!_.resolved) && childrenResolved && !hasSpecialExpressions
   }
 
+  override def distinctSet: AttributeSet = {
+    if (child.outputSet.nonEmpty && child.outputSet.subsetOf(outputSet)) {
+      child.distinctSet
+    } else {
+      AttributeSet.empty
+    }
+  }
+
   override def validConstraints: Set[Expression] =
     child.constraints.union(getAliasedConstraints(projectList))
 }
@@ -107,6 +115,8 @@ case class Filter(condition: Expression, child: LogicalPlan)
 
   override def maxRows: Option[Long] = child.maxRows
 
+  override def distinctSet: AttributeSet = child.distinctSet
+
   override protected def validConstraints: Set[Expression] =
     child.constraints.union(splitConjunctivePredicates(condition).toSet)
 }
@@ -137,6 +147,8 @@ case class Intersect(left: LogicalPlan, right: LogicalPlan) extends SetOperation
       leftAttr.withNullability(leftAttr.nullable && rightAttr.nullable)
     }
 
+  override def distinctSet: AttributeSet = left.outputSet
+
   override protected def validConstraints: Set[Expression] =
     leftConstraints.union(rightConstraints)
 
@@ -168,6 +180,8 @@ case class Except(left: LogicalPlan, right: LogicalPlan) extends SetOperation(le
   /** We don't use right.output because those rows get excluded from the set. */
   override def output: Seq[Attribute] = left.output
 
+  override def distinctSet: AttributeSet = left.outputSet
+
   override protected def validConstraints: Set[Expression] = leftConstraints
 
   override lazy val resolved: Boolean =
@@ -265,6 +279,9 @@ case class Join(
     }
   }
 
+  override def distinctSet: AttributeSet =
+    if (joinType == LeftSemi) left.distinctSet else AttributeSet.empty
+
   override protected def validConstraints: Set[Expression] = {
     joinType match {
       case Inner if condition.isDefined =>
@@ -312,6 +329,7 @@ case class Join(
  */
 case class BroadcastHint(child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = child.output
+  override def distinctSet: AttributeSet = child.distinctSet
 
   // We manually set statistics of BroadcastHint to smallest value to make sure
   // the plan wrapped by BroadcastHint will be considered to broadcast later.
@@ -367,6 +385,7 @@ case class Sort(
     child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = child.output
   override def maxRows: Option[Long] = child.maxRows
+  override def distinctSet: AttributeSet = child.distinctSet
 }
 
 /** Factory for constructing new `Range` nodes. */
@@ -422,6 +441,19 @@ case class Aggregate(
   override def output: Seq[Attribute] = aggregateExpressions.map(_.toAttribute)
   override def maxRows: Option[Long] = child.maxRows
 
+  override def distinctSet: AttributeSet = {
+    if (isForDistinct) {
+      AttributeSet(aggregateExpressions)
+    } else if (child.outputSet.nonEmpty && child.outputSet.subsetOf(outputSet)) {
+      child.distinctSet
+    } else {
+      AttributeSet.empty
+    }
+  }
+
+  def isForDistinct: Boolean =
+    groupingExpressions.nonEmpty && groupingExpressions == aggregateExpressions
+
   override def validConstraints: Set[Expression] =
     child.constraints.union(getAliasedConstraints(aggregateExpressions))
 
@@ -443,6 +475,8 @@ case class Window(
   override def output: Seq[Attribute] =
     child.output ++ windowExpressions.map(_.toAttribute)
 
+  override def distinctSet: AttributeSet = child.distinctSet
+
   def windowOutputSet: AttributeSet = AttributeSet(windowExpressions.map(_.toAttribute))
 }
 
@@ -585,6 +619,7 @@ object Limit {
 
 case class GlobalLimit(limitExpr: Expression, child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = child.output
+  override def distinctSet: AttributeSet = child.distinctSet
   override def maxRows: Option[Long] = {
     limitExpr match {
       case IntegerLiteral(limit) => Some(limit)
@@ -600,6 +635,7 @@ case class GlobalLimit(limitExpr: Expression, child: LogicalPlan) extends UnaryN
 
 case class LocalLimit(limitExpr: Expression, child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = child.output
+  override def distinctSet: AttributeSet = child.distinctSet
   override def maxRows: Option[Long] = {
     limitExpr match {
       case IntegerLiteral(limit) => Some(limit)
@@ -615,6 +651,7 @@ case class LocalLimit(limitExpr: Expression, child: LogicalPlan) extends UnaryNo
 
 case class SubqueryAlias(alias: String, child: LogicalPlan) extends UnaryNode {
 
+  override def distinctSet: AttributeSet = child.distinctSet
   override def output: Seq[Attribute] = child.output.map(_.withQualifier(Some(alias)))
 }
 
@@ -638,6 +675,7 @@ case class Sample(
     val isTableSample: java.lang.Boolean = false) extends UnaryNode {
 
   override def output: Seq[Attribute] = child.output
+  override def distinctSet: AttributeSet = child.distinctSet
 
   override def statistics: Statistics = {
     val ratio = upperBound - lowerBound
@@ -658,6 +696,7 @@ case class Sample(
 case class Distinct(child: LogicalPlan) extends UnaryNode {
   override def maxRows: Option[Long] = child.maxRows
   override def output: Seq[Attribute] = child.output
+  override def distinctSet: AttributeSet = child.outputSet
 }
 
 /**