[SPARK-16355] [SPARK-16354] [SQL] Fix Bugs When LIMIT/TABLESAMPLE is Non-foldable, Zero or Negative

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala

@@ -46,6 +46,15 @@ trait CheckAnalysis extends PredicateHelper {
     }).length > 1
   }
 
+  private def checkLimitClause(limitExpr: Expression): Unit = {
+    val numRows = limitExpr.eval().asInstanceOf[Int]
+    if (numRows < 0) {
+      failAnalysis(
+        s"number_rows in limit clause must be equal to or greater than 0. " +
+          s"number_rows:$numRows")
+    }
+  }
+
   def checkAnalysis(plan: LogicalPlan): Unit = {
     // We transform up and order the rules so as to catch the first possible failure instead
     // of the result of cascading resolution failures.
@@ -251,6 +260,10 @@ trait CheckAnalysis extends PredicateHelper {
                 s"but one table has '${firstError.output.length}' columns and another table has " +
                 s"'${s.children.head.output.length}' columns")
 
+          case GlobalLimit(limitExpr, _) => checkLimitClause(limitExpr)
+
+          case LocalLimit(limitExpr, _) => checkLimitClause(limitExpr)
+
           case p if p.expressions.exists(ScalarSubquery.hasCorrelatedScalarSubquery) =>
             p match {
               case _: Filter | _: Aggregate | _: Project => // Ok

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

@@ -660,7 +660,12 @@ case class GlobalLimit(limitExpr: Expression, child: LogicalPlan) extends UnaryN
   }
   override lazy val statistics: Statistics = {
     val limit = limitExpr.eval().asInstanceOf[Int]
-    val sizeInBytes = (limit: Long) * output.map(a => a.dataType.defaultSize).sum
+    var sizeInBytes = (limit: Long) * output.map(a => a.dataType.defaultSize).sum
+    if (sizeInBytes == 0) {
+      // sizeInBytes can't be zero, or sizeInBytes of BinaryNode will also be zero
+      // (product of children).
+      sizeInBytes = 1
+    }
     child.statistics.copy(sizeInBytes = sizeInBytes)
   }
 }
@@ -675,7 +680,12 @@ case class LocalLimit(limitExpr: Expression, child: LogicalPlan) extends UnaryNo
   }
   override lazy val statistics: Statistics = {
     val limit = limitExpr.eval().asInstanceOf[Int]
-    val sizeInBytes = (limit: Long) * output.map(a => a.dataType.defaultSize).sum
+    var sizeInBytes = (limit: Long) * output.map(a => a.dataType.defaultSize).sum
+    if (sizeInBytes == 0) {
+      // sizeInBytes can't be zero, or sizeInBytes of BinaryNode will also be zero
+      // (product of children).
+      sizeInBytes = 1
+    }
     child.statistics.copy(sizeInBytes = sizeInBytes)
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala

@@ -352,6 +352,12 @@ class AnalysisErrorSuite extends AnalysisTest {
     "Generators are not supported outside the SELECT clause, but got: Sort" :: Nil
   )
 
+  errorTest(
+    "num_rows in limit clause must be equal to or greater than 0",
+    listRelation.limit(-1),
+    "number_rows in limit clause must be equal to or greater than 0. number_rows:-1" :: Nil
+  )
+
   errorTest(
     "more than one generators in SELECT",
     listRelation.select(Explode('list), Explode('list)),

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

@@ -592,6 +592,19 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
     }
   }
 
+  test("negative in LIMIT or TABLESAMPLE") {
+    val expected = "number_rows in limit clause must be equal to or greater than 0. number_rows:-1"
+    var e = intercept[AnalysisException] {
+      sql("SELECT * FROM testData TABLESAMPLE (-1 rows)")
+    }.getMessage
+    assert(e.contains(expected))
+
+    e = intercept[AnalysisException] {
+      sql("SELECT * FROM testData LIMIT -1")
+    }.getMessage
+    assert(e.contains(expected))
+  }
+
   test("aggregates with nulls") {
     checkAnswer(
       sql("SELECT SKEWNESS(a), KURTOSIS(a), MIN(a), MAX(a)," +

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

@@ -17,10 +17,12 @@
 
 package org.apache.spark.sql
 
+import org.apache.spark.sql.catalyst.plans.logical.{GlobalLimit, Join, LocalLimit}
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types._
 
 class StatisticsSuite extends QueryTest with SharedSQLContext {
+  import testImplicits._
 
   test("SPARK-15392: DataFrame created from RDD should not be broadcasted") {
     val rdd = sparkContext.range(1, 100).map(i => Row(i, i))
@@ -31,4 +33,46 @@ class StatisticsSuite extends QueryTest with SharedSQLContext {
       spark.sessionState.conf.autoBroadcastJoinThreshold)
   }
 
+  test("estimates the size of limit") {
+    withTempTable("test") {
+      Seq(("one", 1), ("two", 2), ("three", 3), ("four", 4)).toDF("k", "v")
+        .createOrReplaceTempView("test")
+      Seq((0, 1), (1, 24), (2, 48)).foreach { case (limit, expected) =>
+        val df = sql(s"""SELECT * FROM test limit $limit""")
+
+        val sizesGlobalLimit = df.queryExecution.analyzed.collect { case g: GlobalLimit =>
+          g.statistics.sizeInBytes
+        }
+        assert(sizesGlobalLimit.size === 1, s"Size wrong for:\n ${df.queryExecution}")
+        assert(sizesGlobalLimit(0).equals(BigInt(expected)),
+          s"expected exact size 24 for table 'test', got: ${sizesGlobalLimit(0)}")
+
+        val sizesLocalLimit = df.queryExecution.analyzed.collect { case l: LocalLimit =>
+          l.statistics.sizeInBytes
+        }
+        assert(sizesLocalLimit.size === 1, s"Size wrong for:\n ${df.queryExecution}")
+        assert(sizesLocalLimit(0).equals(BigInt(expected)),
+          s"expected exact size 24 for table 'test', got: ${sizesLocalLimit(0)}")
+      }
+    }
+  }
+
+  test("estimates the size of a limit 0 on outer join") {
+    withTempTable("test") {
+      Seq(("one", 1), ("two", 2), ("three", 3), ("four", 4)).toDF("k", "v")
+        .createOrReplaceTempView("test")
+      val df1 = spark.table("test")
+      val df2 = spark.table("test").limit(0)
+      val df = df1.join(df2, Seq("k"), "left")
+
+      val sizes = df.queryExecution.analyzed.collect { case g: Join =>
+        g.statistics.sizeInBytes
+      }
+
+      assert(sizes.size === 1, s"Size wrong for:\n ${df.queryExecution}")
+      assert(sizes(0).equals(BigInt(96)),
+        s"expected exact size 96 for table 'test', got: ${sizes(0)}")
+    }
+  }
+
 }