Add support of user-provided standard deviation to ApproxCountDistinct.

Author: larvaboy

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala

@@ -319,7 +319,12 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
     COUNT ~> "(" ~ "*" <~ ")" ^^ { case _ => Count(Literal(1)) } |
     COUNT ~> "(" ~ expression <~ ")" ^^ { case dist ~ exp => Count(exp) } |
     COUNT ~> "(" ~> DISTINCT ~> expression <~ ")" ^^ { case exp => CountDistinct(exp :: Nil) } |
-    APPROXIMATE ~> COUNT ~> "(" ~> DISTINCT ~> expression <~ ")" ^^ { case exp => ApproxCountDistinct(exp) } |
+    APPROXIMATE ~> COUNT ~> "(" ~> DISTINCT ~> expression <~ ")" ^^ {
+      case exp => ApproxCountDistinct(exp)
+    } |
+    APPROXIMATE ~> "(" ~> floatLit ~ ")" ~ COUNT ~ "(" ~ DISTINCT ~ expression <~ ")" ^^ {
+      case s ~ _ ~ _ ~ _ ~ _ ~ e => ApproxCountDistinct(e, s.toDouble)
+    } |
     FIRST ~> "(" ~> expression <~ ")" ^^ { case exp => First(exp) } |
     AVG ~> "(" ~> expression <~ ")" ^^ { case exp => Average(exp) } |
     MIN ~> "(" ~> expression <~ ")" ^^ { case exp => Min(exp) } |

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala

@@ -171,38 +171,38 @@ case class CountDistinct(expressions: Seq[Expression]) extends AggregateExpressi
   override def newInstance() = new CountDistinctFunction(expressions, this)
 }
 
-case class ApproxCountDistinctPartition(child: Expression)
+case class ApproxCountDistinctPartition(child: Expression, relativeSD: Double)
   extends AggregateExpression with trees.UnaryNode[Expression] {
   override def references = child.references
   override def nullable = false
   override def dataType = child.dataType
   override def toString = s"APPROXIMATE COUNT(DISTINCT $child)"
-  override def newInstance() = new ApproxCountDistinctPartitionFunction(child, this)
+  override def newInstance() = new ApproxCountDistinctPartitionFunction(child, this, relativeSD)
 }
 
-case class ApproxCountDistinctMerge(child: Expression)
+case class ApproxCountDistinctMerge(child: Expression, relativeSD: Double)
   extends AggregateExpression with trees.UnaryNode[Expression] {
   override def references = child.references
   override def nullable = false
   override def dataType = IntegerType
   override def toString = s"APPROXIMATE COUNT(DISTINCT $child)"
-  override def newInstance() = new ApproxCountDistinctMergeFunction(child, this)
+  override def newInstance() = new ApproxCountDistinctMergeFunction(child, this, relativeSD)
 }
 
-object ApproxCountDistinct {
-  val RelativeSD = 0.05
-}
-
-case class ApproxCountDistinct(child: Expression)
+case class ApproxCountDistinct(child: Expression, relativeSD: Double = 0.05)
   extends PartialAggregate with trees.UnaryNode[Expression] {
   override def references = child.references
   override def nullable = false
   override def dataType = IntegerType
   override def toString = s"APPROXIMATE COUNT(DISTINCT $child)"
 
   override def asPartial: SplitEvaluation = {
-    val partialCount = Alias(ApproxCountDistinctPartition(child), "PartialApproxCountDistinct")()
-    SplitEvaluation(ApproxCountDistinctMerge(partialCount.toAttribute), partialCount :: Nil)
+    val partialCount =
+      Alias(ApproxCountDistinctPartition(child, relativeSD), "PartialApproxCountDistinct")()
+
+    SplitEvaluation(
+      ApproxCountDistinctMerge(partialCount.toAttribute, relativeSD),
+      partialCount :: Nil)
   }
 
   override def newInstance() = new CountDistinctFunction(child :: Nil, this)
@@ -307,11 +307,14 @@ case class CountFunction(expr: Expression, base: AggregateExpression) extends Ag
   override def eval(input: Row): Any = count
 }
 
-case class ApproxCountDistinctPartitionFunction(expr: Expression, base: AggregateExpression)
+case class ApproxCountDistinctPartitionFunction(
+    expr: Expression,
+    base: AggregateExpression,
+    relativeSD: Double)
   extends AggregateFunction {
-  def this() = this(null, null) // Required for serialization.
+  def this() = this(null, null, 0) // Required for serialization.
 
-  private val hyperLogLog = new HyperLogLog(ApproxCountDistinct.RelativeSD)
+  private val hyperLogLog = new HyperLogLog(relativeSD)
 
   override def update(input: Row): Unit = {
     val evaluatedExpr = expr.eval(input)
@@ -323,11 +326,14 @@ case class ApproxCountDistinctPartitionFunction(expr: Expression, base: Aggregat
   override def eval(input: Row): Any = hyperLogLog
 }
 
-case class ApproxCountDistinctMergeFunction(expr: Expression, base: AggregateExpression)
+case class ApproxCountDistinctMergeFunction(
+    expr: Expression,
+    base: AggregateExpression,
+    relativeSD: Double)
   extends AggregateFunction {
-  def this() = this(null, null) // Required for serialization.
+  def this() = this(null, null, 0) // Required for serialization.
 
-  private val hyperLogLog = new HyperLogLog(ApproxCountDistinct.RelativeSD)
+  private val hyperLogLog = new HyperLogLog(relativeSD)
 
   override def update(input: Row): Unit = {
     val evaluatedExpr = expr.eval(input)

sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

@@ -111,6 +111,12 @@ class SQLQuerySuite extends QueryTest {
       3)
   }
 
+  test("approximate count distinct with user provided standard deviation") {
+    checkAnswer(
+      sql("SELECT APPROXIMATE(0.04) COUNT(DISTINCT a) FROM testData2"),
+      3)
+  }
+
   // No support for primitive nulls yet.
   ignore("null count") {
     checkAnswer(