add test cases for CTE feature

jacek-lewandowski · Mar 3, 2015 · 614182f · 614182f
1 parent 32e415b
commit 614182f
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 8 deletions.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -17,16 +17,11 @@
 
 package org.apache.spark.sql
 
-import org.apache.spark.sql.test.TestSQLContext
-import org.scalatest.BeforeAndAfterAll
-
-import org.apache.spark.sql.functions._
-import org.apache.spark.sql.catalyst.errors.TreeNodeException
-import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
-import org.apache.spark.sql.types._
-
 import org.apache.spark.sql.TestData._
+import org.apache.spark.sql.test.TestSQLContext
 import org.apache.spark.sql.test.TestSQLContext.{udf => _, _}
+import org.apache.spark.sql.types._
+import org.scalatest.BeforeAndAfterAll
 
 
 class SQLQuerySuite extends QueryTest with BeforeAndAfterAll {
@@ -287,6 +282,20 @@ class SQLQuerySuite extends QueryTest with BeforeAndAfterAll {
       mapData.collect().take(1).map(Row.fromTuple).toSeq)
   }
 
+  test("CTE feature") {
+    checkAnswer(
+      sql("with q1 as (select * from testData limit 10) select * from q1"),
+      testData.take(10).toSeq)
+
+    checkAnswer(
+      sql("""
+            |with q1 as (select * from testData where key= '5'),
+            |q2 as (select * from testData where key = '4')
+            |select * from q1 union all select * from q2""".stripMargin),
+      Row(5, "5") :: Row(4, "4") :: Nil)
+
+  }
+
   test("date row") {
     checkAnswer(sql(
       """select cast("2015-01-28" as date) from testData limit 1"""),

diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
@@ -207,6 +207,26 @@ class SQLQuerySuite extends QueryTest {
     }
   }
 
+  test("CTE feature") {
+    checkAnswer(
+      sql("with q1 as (select key from src) select * from q1 where key = 5"),
+      sql("select * from (select key from src) q1 where key = 5").collect().toSeq
+    )
+
+    checkAnswer(
+      sql("""
+            |with q1 as (select * from src where key= '5'),
+            |q2 as (select * from src s2 where key = '4')
+            |select * from q1 union all select * from q2"""),
+      sql(
+        """
+          |select * from (select * from src where key = '5') q2
+          |union all
+          |select * from (select * from src where key = '4') q1
+        """.stripMargin).collect().toSeq
+    )
+  }
+
   test("command substitution") {
     sql("set tbl=src")
     checkAnswer(