LuciferYang commented on code in PR #55579: URL: https://github.com/apache/spark/pull/55579#discussion_r4489844535
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala: ########## @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2 + +import java.util.Locale +import java.util.concurrent.atomic.AtomicReference + +import scala.util.control.NonFatal + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate} +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} +import org.apache.spark.sql.connector.read.InputPartition +import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory} +import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering + +/** + * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over + * its partition columns (SPARK-30628). The built-in file formats mix this in. + * + * Both read paths start from one listing of the file index and split and pack the files they + * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that + * reports a partitioning or an ordering, should not mix this in. + */ +trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering { + + @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]() + @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]() + + /** + * The file index listing and its packing, each taken once per scan so planning and execution + * read one snapshot. A `lazy val` would compute them while holding the scan's monitor. + */ + private def listing: Seq[PartitionDirectory] = + once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters)) + + private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing)) + + private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T = + Option(ref.get()).getOrElse { + ref.compareAndSet(null, compute) + ref.get() + } + + override def planInputPartitions(): Array[InputPartition] = allPartitions.toArray Review Comment: Done in a76b80e435f. The trait now declares `final override protected def partitions = allPartitions()`, and `planInputPartitions()` is final as well, since overriding it skips the runtime-filtered path in the same way. A subclass of a built-in scan can still override `filePartitions`, and both read paths go through it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
