diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileIndex.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileIndex.scala index 0291a5fd28a7..b4ed1c4a0a54 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileIndex.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileIndex.scala @@ -133,6 +133,20 @@ trait FileIndex { def listFiles( partitionFilters: Seq[Expression], dataFilters: Seq[Expression]): Seq[PartitionDirectory] + /** + * Given the `partitionFilters` Spark will use to prune files for this index, returns the + * subset that this index guarantees are fully applied for all returned rows. Only these are + * dropped from the row-level `FilterExec` above the scan; any filter not returned is still + * evaluated after the scan for correctness. + * + * The default returns all `partitionFilters`, matching standard file-source behavior where + * partition filters are guaranteed by file listing. Override to omit filters that are only + * partially applied - for example with partition evolution, where not all data files are + * partitioned by a given partition column and may contain both matching and non-matching rows. + */ + def canFullPushDownPartitionFilter(partitionFilters: Seq[Expression]): Seq[Expression] = + partitionFilters + /** * Returns the list of files that will be read when scanning this relation. This call may be * very expensive for large tables. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala index e2427222d8eb..1aa00a1a1482 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala @@ -212,8 +212,10 @@ object FileSourceStrategy extends Strategy with PredicateHelper with Logging { .flatMap(DataSourceStrategy.translateFilter(_, supportNestedPredicatePushdown)) logInfo(log"Pushed Filters: ${MDC(PUSHED_FILTERS, pushedFilters.mkString(","))}") - // Predicates with both partition keys and attributes need to be evaluated after the scan. - val afterScanFilters = filterSet -- partitionKeyFilters.filter(_.references.nonEmpty) + // Drop partition-key filters that the FileIndex guarantees are fully applied from the + // FilterExec above the scan. By default all partition-key filters are dropped. + val afterScanFilters = filterSet -- fsRelation.location.canFullPushDownPartitionFilter( + partitionKeyFilters.filter(_.references.nonEmpty).toSeq) val maxToStringFields = fsRelation.sparkSession.conf.get(SQLConf.MAX_TO_STRING_FIELDS) logInfo(log"Post-Scan Filters: ${MDC(POST_SCAN_FILTERS, afterScanFilters.simpleString(maxToStringFields))}") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategySuite.scala index f38996adbd99..03d43b4a78ed 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategySuite.scala @@ -252,6 +252,27 @@ class FileSourceStrategySuite extends SharedSparkSession { assert(getPhysicalFilters(df2) contains resolve(df2, "(p1 + c2) = 2")) } + test("FileIndex canFullPushDownPartitionFilter controls final FilterExec removal") { + val (table, testIndex) = + createTableWithCanFullPushDownPartitionFilterIndex( + files = Seq( + "p1=1/file1" -> 10, + "p1=2/file2" -> 10)) + + // By default the index returns the partition filters, matching legacy behavior, so the + // partition-key filter is removed from the FilterExec above the scan. + val defaultPartitionFilter = table.where("p1 = 1") + assert(!getPhysicalFilters(defaultPartitionFilter).contains( + resolve(defaultPartitionFilter, "p1 = 1"))) + + // Returning nothing means no partition filter is fully pushed, so the partition-key + // filter stays in the FilterExec above the scan. + testIndex.setCanFullPushDownPartitionFilter(_ => Nil) + val explicitNoFullyPushed = table.where("p1 = 1") + assert(getPhysicalFilters(explicitNoFullyPushed).contains( + resolve(explicitNoFullyPushed, "p1 = 1"))) + } + test("bucketed table") { val table = createTable( @@ -707,6 +728,57 @@ class FileSourceStrategySuite extends SharedSparkSession { } } + private class CanFullPushDownPartitionFilterTestFileIndex( + sparkSession: SparkSession, + rootPathsSpecified: Seq[Path], + parameters: Map[String, String], + userSpecifiedSchema: Option[StructType]) + extends InMemoryFileIndex( + sparkSession, + rootPathsSpecified, + parameters, + userSpecifiedSchema) { + + private var canFullPushDown: Seq[Expression] => Seq[Expression] = identity + + override def canFullPushDownPartitionFilter( + partitionFilters: Seq[Expression]): Seq[Expression] = + canFullPushDown(partitionFilters) + + def setCanFullPushDownPartitionFilter(f: Seq[Expression] => Seq[Expression]): Unit = { + canFullPushDown = f + } + } + + private def createTableWithCanFullPushDownPartitionFilterIndex( + files: Seq[(String, Int)]): (DataFrame, CanFullPushDownPartitionFilterTestFileIndex) = { + val tempDir = Utils.createTempDir() + files.foreach { + case (name, size) => + val file = new File(tempDir, name) + assert(file.getParentFile.exists() || Utils.createDirectory(file.getParentFile)) + util.stringToFile(file, "*".repeat(size)) + } + + val dataSchema = StructType(Nil) + .add("c1", IntegerType) + .add("c2", IntegerType) + val fileIndex = new CanFullPushDownPartitionFilterTestFileIndex( + spark, + Seq(new Path(tempDir.getCanonicalPath)), + Map.empty[String, String], + Some(dataSchema)) + val relation = HadoopFsRelation( + location = fileIndex, + partitionSchema = fileIndex.partitionSchema, + dataSchema = dataSchema, + bucketSpec = None, + fileFormat = TestFileFormat(), + options = Map.empty)(spark) + + Dataset.ofRows(spark, LogicalRelation(relation)) -> fileIndex + } + def getFileScanRDD(df: DataFrame): FileScanRDD = { df.queryExecution.executedPlan.collect { case scan: DataSourceScanExec if scan.inputRDDs().head.isInstanceOf[FileScanRDD] =>