LuciferYang commented on code in PR #55579: URL: https://github.com/apache/spark/pull/55579#discussion_r4489850200
########## 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 + + /** + * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps + * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over + * them is evaluated again above the scan. + */ + override def filterAttributes(): Array[NamedReference] = { + val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet + readPartitionSchema.fieldNames + .filter(name => readFields.contains(normalizePartitionName(name))) + .map(FieldReference.column) + } + + /** + * Keeps the listed partition directories whose values satisfy `expressions`, then splits and + * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`. + * + * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec` + * above the scan applies the filter again. + */ + override def planInputPartitionsWithRuntimeFilters( + expressions: Array[Expression]): Array[InputPartition] = { + if (expressions.isEmpty) return allPartitions.toArray + val readFields = readPartitionSchema.fieldNames.map(normalizePartitionName).toSet + val notApplicable = expressions.filterNot( + _.references.forall(a => readFields.contains(normalizePartitionName(a.name)))) + if (notApplicable.nonEmpty) { Review Comment: Done in a76b80e435f. The check is in the `fieldIndex.getOrElse`, against `fileIndex.partitionSchema`, with a comment that `PushDownUtils` passes only filters over `filterAttributes()`. A direct call with a filter over a partition column the scan does not read now evaluates it on the directory values instead of failing. ########## 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 = Review Comment: Done in a76b80e435f. Both memos are `TransientBestEffortLazyVal`, and `once`, the holders and the `AtomicReference` import are gone. ########## connector/avro/src/main/scala/org/apache/spark/sql/v2/avro/AvroScan.scala: ########## @@ -41,7 +41,7 @@ case class AvroScan( options: CaseInsensitiveStringMap, pushedFilters: Array[Filter], partitionFilters: Seq[Expression] = Seq.empty, - dataFilters: Seq[Expression] = Seq.empty) extends FileScan { + dataFilters: Seq[Expression] = Seq.empty) extends RuntimePartitionPruningFileScan { Review Comment: Done in a76b80e435f, in all four scans. -- 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]
