[SPARK-1822] SchemaRDD.count() should use query optimizer

python/pyspark/sql.py

@@ -268,7 +268,7 @@ def __init__(self, jschema_rdd, sql_ctx):
     def _jrdd(self):
         """
         Lazy evaluation of PythonRDD object. Only done when a user calls methods defined by the
-        L{pyspark.rdd.RDD} super class (map, count, etc.).
+        L{pyspark.rdd.RDD} super class (map, filter, etc.).
         """
         if not hasattr(self, '_lazy_jrdd'):
             self._lazy_jrdd = self._toPython()._jrdd
@@ -321,6 +321,18 @@ def saveAsTable(self, tableName):
         """
         self._jschema_rdd.saveAsTable(tableName)
 
+    def count(self):
+        """
+        Return the number of elements in this RDD.
+
+        >>> srdd = sqlCtx.inferSchema(rdd)
+        >>> srdd.count()
+        3L
+        >>> srdd.count() == srdd.map(lambda x: x).count()
+        True
+        """
+        return self._jschema_rdd.count()
+
     def _toPython(self):
         # We have to import the Row class explicitly, so that the reference Pickler has is
         # pyspark.sql.Row instead of __main__.Row

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

@@ -151,7 +151,7 @@ case class MaxFunction(expr: Expression, base: AggregateExpression) extends Aggr
 case class Count(child: Expression) extends PartialAggregate with trees.UnaryNode[Expression] {
   override def references = child.references
   override def nullable = false
-  override def dataType = IntegerType
+  override def dataType = LongType
   override def toString = s"COUNT($child)"
 
   override def asPartial: SplitEvaluation = {
@@ -295,12 +295,12 @@ case class AverageFunction(expr: Expression, base: AggregateExpression)
 case class CountFunction(expr: Expression, base: AggregateExpression) extends AggregateFunction {
   def this() = this(null, null) // Required for serialization.
 
-  var count: Int = _
+  var count: Long = _
 
   override def update(input: Row): Unit = {
     val evaluatedExpr = expr.map(_.eval(input))
     if (evaluatedExpr.map(_ != null).reduceLeft(_ || _)) {
-      count += 1
+      count += 1L
     }
   }
 

sql/core/src/main/scala/org/apache/spark/sql/SchemaRDD.scala

@@ -274,6 +274,15 @@ class SchemaRDD(
       seed: Long) =
     new SchemaRDD(sqlContext, Sample(fraction, withReplacement, seed, logicalPlan))
 
+  /**
+   * :: Experimental ::
+   * Overriding base RDD implementation to leverage query optimizer
+   */
+  @Experimental
+  override def count(): Long = {
+    groupBy()(Count(Literal(1))).collect().head.getLong(0)
+  }
+
   /**
    * :: Experimental ::
    * Applies the given Generator, or table generating function, to this relation.

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

@@ -108,10 +108,7 @@ class DslQuerySuite extends QueryTest {
   }
 
   test("count") {
-    checkAnswer(
-      testData2.groupBy()(Count(1)),
-      testData2.count()
-    )
+    assert(testData2.count() === testData2.map(_ => 1).count())
   }
 
   test("null count") {
@@ -126,6 +123,10 @@ class DslQuerySuite extends QueryTest {
     )
   }
 
+  test("zero count") {
+    assert(testData4.count() === 0)
+  }
+
   test("inner join where, one match per row") {
     checkAnswer(
       upperCaseData.join(lowerCaseData, Inner).where('n === 'N),

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

@@ -47,6 +47,8 @@ object TestData {
       (1, null) ::
       (2, 2) :: Nil)
 
+  val testData4 = logical.LocalRelation('a.int, 'b.int)
+
   case class UpperCaseData(N: Int, L: String)
   val upperCaseData =
     TestSQLContext.sparkContext.parallelize(