Push down In filter to Parquet.

Author: viirya

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilters.scala

@@ -201,6 +201,9 @@ private[sql] object ParquetFilters {
   }
 
   private val makeInSet: PartialFunction[DataType, (String, Set[Any]) => FilterPredicate] = {
+    case BooleanType =>
+      (n: String, v: Set[Any]) =>
+        FilterApi.userDefined(booleanColumn(n), SetInFilter(v.asInstanceOf[Set[java.lang.Boolean]]))
     case IntegerType =>
       (n: String, v: Set[Any]) =>
         FilterApi.userDefined(intColumn(n), SetInFilter(v.asInstanceOf[Set[java.lang.Integer]]))
@@ -252,6 +255,9 @@ private[sql] object ParquetFilters {
       case sources.IsNotNull(name) =>
         makeNotEq.lift(dataTypeOf(name)).map(_(name, null))
 
+      case sources.In(name, values) =>
+        makeInSet.lift(dataTypeOf(name)).map(_(name, values.toSet))
+
       case sources.EqualTo(name, value) =>
         makeEq.lift(dataTypeOf(name)).map(_(name, value))
       case sources.Not(sources.EqualTo(name, value)) =>

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala

@@ -112,6 +112,11 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
       checkFilterPredicate('_1 === true, classOf[Eq[_]], true)
       checkFilterPredicate('_1 <=> true, classOf[Eq[_]], true)
       checkFilterPredicate('_1 !== true, classOf[NotEq[_]], false)
+
+      checkFilterPredicate(('_1 in(true)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], true)
+      checkFilterPredicate(('_1 in(false)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], false)
     }
   }
 
@@ -138,6 +143,11 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
 
       checkFilterPredicate(!('_1 < 4), classOf[GtEq[_]], 4)
       checkFilterPredicate('_1 < 2 || '_1 > 3, classOf[Operators.Or], Seq(Row(1), Row(4)))
+
+      checkFilterPredicate(('_1 in(1, 2)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row(1), Row(2)))
+      checkFilterPredicate(('_1 in(3, 4)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row(3), Row(4)))
     }
   }
 
@@ -164,6 +174,11 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
 
       checkFilterPredicate(!('_1 < 4), classOf[GtEq[_]], 4)
       checkFilterPredicate('_1 < 2 || '_1 > 3, classOf[Operators.Or], Seq(Row(1), Row(4)))
+
+      checkFilterPredicate(('_1 in(1L, 2L)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row(1L), Row(2L)))
+      checkFilterPredicate(('_1 in(3L, 4L)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row(3L), Row(4L)))
     }
   }
 
@@ -190,6 +205,11 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
 
       checkFilterPredicate(!('_1 < 4), classOf[GtEq[_]], 4)
       checkFilterPredicate('_1 < 2 || '_1 > 3, classOf[Operators.Or], Seq(Row(1), Row(4)))
+
+      checkFilterPredicate(('_1 in(1.0f, 2.0f)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row(1.0f), Row(2.0f)))
+      checkFilterPredicate(('_1 in(3.0f, 4.0f)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row(3.0f), Row(4.0f)))
     }
   }
 
@@ -216,6 +236,11 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
 
       checkFilterPredicate(!('_1 < 4), classOf[GtEq[_]], 4)
       checkFilterPredicate('_1 < 2 || '_1 > 3, classOf[Operators.Or], Seq(Row(1), Row(4)))
+
+      checkFilterPredicate(('_1 in(1.0, 2.0)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row(1.0), Row(2.0)))
+      checkFilterPredicate(('_1 in(3.0, 4.0)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row(3.0), Row(4.0)))
     }
   }
 
@@ -244,6 +269,11 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
 
       checkFilterPredicate(!('_1 < "4"), classOf[GtEq[_]], "4")
       checkFilterPredicate('_1 < "2" || '_1 > "3", classOf[Operators.Or], Seq(Row("1"), Row("4")))
+
+      checkFilterPredicate(('_1 in("1", "2")).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row("1"), Row("2")))
+      checkFilterPredicate(('_1 in("3", "4")).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row("3"), Row("4")))
     }
 
     withParquetDataFrame((1 to 4).map(i => Tuple1(i.toString * 5 + "test"))) { implicit df =>
@@ -298,6 +328,11 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
       checkBinaryFilterPredicate(!('_1 < 4.b), classOf[GtEq[_]], 4.b)
       checkBinaryFilterPredicate(
         '_1 < 2.b || '_1 > 3.b, classOf[Operators.Or], Seq(Row(1.b), Row(4.b)))
+
+      checkFilterPredicate(('_1 in(1.b, 2.b)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row(1.b), Row(2.b)))
+      checkFilterPredicate(('_1 in(3.b, 4.b)).asInstanceOf[Predicate],
+        classOf[UserDefinedByInstance[_, _]], Seq(Row(3.b), Row(4.b)))
     }
   }