Implement unhandled filters for Parquet

Author: HyukjinKwon

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

@@ -19,6 +19,8 @@ package org.apache.spark.sql.execution.datasources.parquet
 
 import java.io.Serializable
 
+import scala.collection.mutable.ArrayBuffer
+
 import org.apache.parquet.filter2.predicate.FilterApi._
 import org.apache.parquet.filter2.predicate._
 import org.apache.parquet.io.api.Binary
@@ -207,12 +209,31 @@ private[sql] object ParquetFilters {
      */
   }
 
+  /**
+   *  Return referenced columns in [[sources.Filter]].
+   */
+  def referencedColumns(schema: StructType, predicate: sources.Filter): Array[String] = {
+    val dataTypeOf = schema.map(f => f.name -> f.dataType).toMap
+    val referencedColumns = ArrayBuffer.empty[String]
+    def getDataTypeOf(name: String): DataType = {
+      referencedColumns += name
+      dataTypeOf(name)
+    }
+    createParquetFilter(getDataTypeOf, predicate)
+    referencedColumns.distinct.toArray
+  }
+
   /**
    * Converts data sources filters to Parquet filter predicates.
    */
   def createFilter(schema: StructType, predicate: sources.Filter): Option[FilterPredicate] = {
     val dataTypeOf = schema.map(f => f.name -> f.dataType).toMap
+    createParquetFilter(dataTypeOf, predicate)
+  }
 
+  private def createParquetFilter(
+      dataTypeOf: String => DataType,
+      predicate: sources.Filter): Option[FilterPredicate] = {
     relaxParquetValidTypeMap
 
     // NOTE:
@@ -265,18 +286,18 @@ private[sql] object ParquetFilters {
         // Pushing one side of AND down is only safe to do at the top level.
         // You can see ParquetRelation's initializeLocalJobFunc method as an example.
         for {
-          lhsFilter <- createFilter(schema, lhs)
-          rhsFilter <- createFilter(schema, rhs)
+          lhsFilter <- createParquetFilter(dataTypeOf, lhs)
+          rhsFilter <- createParquetFilter(dataTypeOf, rhs)
         } yield FilterApi.and(lhsFilter, rhsFilter)
 
       case sources.Or(lhs, rhs) =>
         for {
-          lhsFilter <- createFilter(schema, lhs)
-          rhsFilter <- createFilter(schema, rhs)
+          lhsFilter <- createParquetFilter(dataTypeOf, lhs)
+          rhsFilter <- createParquetFilter(dataTypeOf, rhs)
         } yield FilterApi.or(lhsFilter, rhsFilter)
 
       case sources.Not(pred) =>
-        createFilter(schema, pred).map(FilterApi.not)
+        createParquetFilter(dataTypeOf, pred).map(FilterApi.not)
 
       case _ => None
     }

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

@@ -133,6 +133,11 @@ private[sql] class ParquetRelation(
       .map(_.toBoolean)
       .getOrElse(sqlContext.conf.getConf(SQLConf.PARQUET_SCHEMA_MERGING_ENABLED))
 
+  // When merging schemas is enabled and the column of the given filter does not exist,
+  // Parquet emits an exception which is an issue of Parquet (PARQUET-389).
+  private val safeParquetFilterPushDown =
+    sqlContext.conf.parquetFilterPushDown && !shouldMergeSchemas
+
   private val mergeRespectSummaries =
     sqlContext.conf.getConf(SQLConf.PARQUET_SCHEMA_RESPECT_SUMMARIES)
 
@@ -288,20 +293,23 @@ private[sql] class ParquetRelation(
     }
   }
 
+  override def unhandledFilters(filters: Array[Filter]): Array[Filter] = {
+    if (safeParquetFilterPushDown) {
+      filters.filter(ParquetFilters.createFilter(dataSchema, _).isEmpty)
+    } else {
+      filters
+    }
+  }
+
   override def buildInternalScan(
       requiredColumns: Array[String],
       filters: Array[Filter],
       inputFiles: Array[FileStatus],
       broadcastedConf: Broadcast[SerializableConfiguration]): RDD[InternalRow] = {
     val useMetadataCache = sqlContext.getConf(SQLConf.PARQUET_CACHE_METADATA)
-    val parquetFilterPushDown = sqlContext.conf.parquetFilterPushDown
+    val parquetFilterPushDown = safeParquetFilterPushDown
     val assumeBinaryIsString = sqlContext.conf.isParquetBinaryAsString
     val assumeInt96IsTimestamp = sqlContext.conf.isParquetINT96AsTimestamp
-
-    // When merging schemas is enabled and the column of the given filter does not exist,
-    // Parquet emits an exception which is an issue of Parquet (PARQUET-389).
-    val safeParquetFilterPushDown = !shouldMergeSchemas && parquetFilterPushDown
-
     // Parquet row group size. We will use this value as the value for
     // mapreduce.input.fileinputformat.split.minsize and mapred.min.split.size if the value
     // of these flags are smaller than the parquet row group size.
@@ -315,7 +323,7 @@ private[sql] class ParquetRelation(
         dataSchema,
         parquetBlockSize,
         useMetadataCache,
-        safeParquetFilterPushDown,
+        parquetFilterPushDown,
         assumeBinaryIsString,
         assumeInt96IsTimestamp) _
 
@@ -568,6 +576,15 @@ private[sql] object ParquetRelation extends Logging {
     conf.set(ParquetInputFormat.READ_SUPPORT_CLASS, classOf[CatalystReadSupport].getName)
 
     // Try to push down filters when filter push-down is enabled.
+    val safeRequiredColumns = if (parquetFilterPushDown) {
+      val referencedColumns = filters
+        // Collects all columns referenced in Parquet filter predicates.
+        .flatMap(filter => ParquetFilters.referencedColumns(dataSchema, filter))
+      (requiredColumns ++ referencedColumns).distinct
+    } else {
+      requiredColumns
+    }
+
     if (parquetFilterPushDown) {
       filters
         // Collects all converted Parquet filter predicates. Notice that not all predicates can be
@@ -579,7 +596,7 @@ private[sql] object ParquetRelation extends Logging {
     }
 
     conf.set(CatalystReadSupport.SPARK_ROW_REQUESTED_SCHEMA, {
-      val requestedSchema = StructType(requiredColumns.map(dataSchema(_)))
+      val requestedSchema = StructType(safeRequiredColumns.map(dataSchema(_)))
       CatalystSchemaConverter.checkFieldNames(requestedSchema).json
     })