From e0b5e7c80f7b59944309aaa46f86ed8d4e62cd4a Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Fri, 18 Dec 2015 12:11:36 -0800 Subject: [PATCH] [SPARK-12409][SQL]add filter (IN, AND, OR) --- .../sql/execution/datasources/jdbc/JDBCRDD.scala | 13 +++++++++++++ .../scala/org/apache/spark/sql/jdbc/JDBCSuite.scala | 12 ++++++++++++ 2 files changed, 25 insertions(+) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala index 2d38562e0901..1da5e0f0e8c6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala @@ -269,6 +269,13 @@ private[sql] class JDBCRDD( case stringValue: String => s"'${escapeSql(stringValue)}'" case timestampValue: Timestamp => "'" + timestampValue + "'" case dateValue: Date => "'" + dateValue + "'" + case objectValue: Array[Object] => { + val str = objectValue.map { + case string: String => s"'${escapeSql(string)}'" + case other => s"${escapeSql(other.toString)}" + } + str.mkString(",") + } case _ => value } @@ -288,6 +295,12 @@ private[sql] class JDBCRDD( case GreaterThanOrEqual(attr, value) => s"$attr >= ${compileValue(value)}" case IsNull(attr) => s"$attr IS NULL" case IsNotNull(attr) => s"$attr IS NOT NULL" + case In(attr, value) => s"$attr IN (${compileValue(value)})" + case Not(In(attr, value)) => s"$attr NOT IN (${compileValue(value)})" + case Or(filter1, filter2) => + "(" + compileFilter (filter1) + ") OR (" + compileFilter (filter2) + ")" + case And(filter1, filter2) => + "(" + compileFilter (filter1) + ") AND (" + compileFilter (filter2) + ")" case _ => null } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala index 2b91f62c2fa2..9582fa93f6a0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala @@ -184,6 +184,18 @@ class JDBCSuite extends SparkFunSuite with BeforeAndAfter with SharedSQLContext assert(stripSparkFilter(sql("SELECT * FROM foobar WHERE NAME != 'fred'")).collect().size == 2) assert(stripSparkFilter(sql("SELECT * FROM nulltypes WHERE A IS NULL")).collect().size == 1) assert(stripSparkFilter(sql("SELECT * FROM nulltypes WHERE A IS NOT NULL")).collect().size == 0) + assert(stripSparkFilter(sql( + "SELECT * FROM foobar WHERE NAME IN ('mary', 'fred')")).collect().size === 2) + assert(stripSparkFilter(sql( + "SELECT * FROM foobar WHERE NAME NOT IN ('mary', 'fred')")).collect().size === 1) + assert(stripSparkFilter(sql( + "SELECT * FROM foobar WHERE THEID IN (1,3)")).collect().size === 2) + assert(stripSparkFilter(sql( + "SELECT * FROM foobar WHERE THEID NOT IN (1,3)")).collect().size === 1) + assert(stripSparkFilter(sql( + "SELECT * FROM foobar WHERE THEID=1 OR NAME = 'mary'")).collect().size === 2) + assert(stripSparkFilter(sql( + "SELECT * FROM foobar WHERE THEID = 1 OR NAME = 'mary' AND THEID = 2")).collect().size === 2) } test("SELECT * WHERE (quoted strings)") {