[SPARK-55596][SQL] DSV2 Enhanced Partition Stats Filtering ### What changes were proposed in this pull request? PR for SPARK-55596: SPIP doc: https://docs.google.com/document/d/17vcw411PxSRLWoK-BiLI56UiNdokLWtovF8JZUlDTOo ### Why are the changes needed? Enhance partition filtering for DSV2 data sources that have partition stats. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Add unit test : DataSourceV2EnhancedPartitionFilterSuite ### Was this patch authored or co-authored using generative AI tooling? Tests were generated by Cursor Closes #54459 from szehon-ho/partition_filter. Authored-by: Szehon Ho <szehon.apache@gmail.com> Signed-off-by: Gengliang Wang <gengliang@apache.org>
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/PartitionColumnReference.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/PartitionColumnReference.java new file mode 100644 index 0000000..ef51654 --- /dev/null +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/PartitionColumnReference.java
@@ -0,0 +1,39 @@ +/* + * 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.connector.expressions; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.catalog.Table; + +/** + * A reference to a partition column in {@link Table#partitioning()}. + * <p> + * {@link #fieldNames()} returns the partition column name (or names) as reported by + * the table's partition schema. + * {@link #ordinal()} returns the 0-based position in {@link Table#partitioning()}. + * + * @since 4.2.0 + */ +@Evolving +public interface PartitionColumnReference extends NamedReference { + + /** + * Returns the 0-based ordinal of this partition column in {@link Table#partitioning()}. + */ + int ordinal(); +}
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/filter/PartitionPredicate.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/filter/PartitionPredicate.java new file mode 100644 index 0000000..dbc31aa --- /dev/null +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/filter/PartitionPredicate.java
@@ -0,0 +1,92 @@ +/* + * 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.connector.expressions.filter; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.connector.expressions.NamedReference; +import org.apache.spark.sql.connector.expressions.PartitionColumnReference; + +import static org.apache.spark.sql.connector.expressions.Expression.EMPTY_EXPRESSION; + +/** + * Represents a partition predicate that can be evaluated using {@link Table#partitioning()}. + * <p> + * Connectors are expected to leverage partition predicates for pruning whenever they have + * partition metadata to evaluate them. Use {@link #eval(InternalRow)} to evaluate this + * predicate against a single partition's keys. + * </p> + * + * @since 4.2.0 + */ +@Evolving +public abstract class PartitionPredicate extends Predicate { + + public static final String NAME = "PARTITION_PREDICATE"; + + protected PartitionPredicate() { + super(NAME, EMPTY_EXPRESSION); + } + + /** + * {@inheritDoc} + * <p> + * For PartitionPredicate, returns {@link PartitionColumnReference} instances that identify + * the partition columns (from {@link Table#partitioning()}) referenced by this predicate. + * Each reference's {@link PartitionColumnReference#fieldNames()} gives the partition column + * name; {@link PartitionColumnReference#ordinal()} gives the 0-based position in + * {@link Table#partitioning()}. + * <p> + * <b>Example:</b> Suppose {@code Table.partitioning()} returns three partition + * transforms: {@code [years(ts), months(ts), bucket(32, id)]} with ordinals 0, 1, 2. + * Each {@link PartitionColumnReference} has {@link PartitionColumnReference#fieldNames()} + * (the transform display name, e.g. {@code years(ts)}) and + * {@link PartitionColumnReference#ordinal()}: + * <ul> + * <li>{@code years(ts) = 2026} returns one reference: (fieldNames=[years(ts)], ordinal=0).</li> + * <li>{@code years(ts) = 2026 and months(ts) = 01} returns two references: + * (fieldNames=[years(ts)], ordinal=0), (fieldNames=[months(ts)], ordinal=1).</li> + * <li>{@code bucket(32, id) = 1} returns one reference: + * (fieldNames=[bucket(32, id)], ordinal=2).</li> + * </ul> + * <p> + * Data sources can use these references to decide whether to return a predicate for post-scan + * filtering. For example, sources supporting partition spec evolution should return + * PartitionPredicates that reference later-added partition transforms (incompletely + * partitioned data) to Spark for post-scan filter, while predicates that reference only + * initially-added partition transforms may be fully pushed. + * + * @return array of partition column references + */ + @Override + public abstract NamedReference[] references(); + + /** + * Evaluates this predicate against a single partition's keys. + * <p> + * The caller must pass the <b>full</b> partition key: one value per partition transform in + * {@link Table#partitioning()}, in order. A key for only a subset of referenced columns is not + * supported. + * + * @param partitionKey the full partition key for one partition, ordered according to + * {@link Table#partitioning()}. + * @return true if the partition represented by these keys satisfies this predicate. + */ + public abstract boolean eval(InternalRow partitionKey); +}
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownV2Filters.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownV2Filters.java index 1fec939..5c1a04a 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownV2Filters.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownV2Filters.java
@@ -18,6 +18,8 @@ package org.apache.spark.sql.connector.read; import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.expressions.PartitionColumnReference; +import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate; import org.apache.spark.sql.connector.expressions.filter.Predicate; /** @@ -26,6 +28,12 @@ * Please Note that this interface is preferred over {@link SupportsPushDownFilters}, which uses * V1 {@link org.apache.spark.sql.sources.Filter} and is less efficient due to the * internal -> external data conversion. + * <p> + * <b>Iterative pushdown:</b> When {@link #supportsIterativePushdown()} returns true, + * {@link #pushPredicates(Predicate[])} may be called <i>multiple times</i> on the same + * {@link ScanBuilder} instance with additional predicates (e.g. {@link PartitionPredicate}). + * The implementation must accumulate state across all calls, and + * {@link #pushedPredicates()} must return predicates from all of them. * * @since 3.3.0 */ @@ -34,9 +42,24 @@ /** * Pushes down predicates, and returns predicates that need to be evaluated after scanning. + * Any predicate that the data source cannot fully push down must be returned as-is so that + * Spark can evaluate it after the scan; the data source must not modify or drop such predicates. * <p> * Rows should be returned from the data source if and only if all of the predicates match. * That is, predicates must be interpreted as ANDed together. + * <p> + * This method may be called multiple times with additional predicates (e.g. + * {@link PartitionPredicate} when {@link #supportsIterativePushdown()} returns true). + * The implementation must accumulate state across all calls so that + * {@link #pushedPredicates()} can return predicates from all of them. + * <p> + * For each {@link PartitionPredicate}, the implementation can use + * {@link PartitionPredicate#references()} (each {@link PartitionColumnReference} has + * {@link PartitionColumnReference#ordinal()}) to decide whether to return it for post-scan + * filtering. For example, data sources with + * partition spec evolution may return predicates that reference later-added partition + * transforms (incompletely partitioned data) so Spark evaluates them after the scan, while + * predicates that reference only initially-added partition transforms may be fully pushed. */ Predicate[] pushPredicates(Predicate[] predicates); @@ -55,9 +78,25 @@ * Both case 1 and 2 should be considered as pushed predicates and should be returned * by this method. * <p> + * When iterative pushdown is supported and {@link #pushPredicates(Predicate[])} was called + * multiple times, this method must return predicates from <i>all</i> calls. + * <p> * It's possible that there is no predicates in the query and * {@link #pushPredicates(Predicate[])} is never called, * empty array should be returned for this case. */ Predicate[] pushedPredicates(); + + /** + * Returns true if this data source supports iterative filter pushdown. When true, + * {@link #pushPredicates(Predicate[])} may be called multiple times with additional + * predicates (e.g. {@link PartitionPredicate}). The implementation must accumulate state + * across all calls, and {@link #pushedPredicates()} must return predicates from all of them. + * See the class-level Javadoc for the full contract. + * + * @since 4.2.0 + */ + default boolean supportsIterativePushdown() { + return false; + } }
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java index 50921f3..343221f 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java
@@ -35,6 +35,7 @@ import org.apache.spark.sql.connector.expressions.SortDirection; import org.apache.spark.sql.connector.expressions.SortOrder; import org.apache.spark.sql.connector.expressions.UserDefinedScalarFunc; +import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate; import org.apache.spark.sql.connector.expressions.aggregate.Avg; import org.apache.spark.sql.connector.expressions.aggregate.Max; import org.apache.spark.sql.connector.expressions.aggregate.Min; @@ -78,6 +79,8 @@ return visitLiteral(literal); } else if (expr instanceof NamedReference namedReference) { return visitNamedReference(namedReference); + } else if (expr instanceof PartitionPredicate partitionPredicate) { + return visitPartitionPredicate(partitionPredicate); } else if (expr instanceof Cast cast) { return visitCast(build(cast.expression()), cast.expressionDataType(), cast.dataType()); } else if (expr instanceof Extract extract) { @@ -332,6 +335,10 @@ "_LEGACY_ERROR_TEMP_3207", Map.of("expr", String.valueOf(expr))); } + protected String visitPartitionPredicate(PartitionPredicate partitionPredicate) { + return partitionPredicate.describe(); + } + protected String visitOverlay(String[] inputs) { assert(inputs.length == 3 || inputs.length == 4); if (inputs.length == 3) {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala index fd3d1da..c4c6a60 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala
@@ -34,6 +34,7 @@ import org.apache.spark.sql.connector.expressions.filter.{AlwaysFalse, AlwaysTrue} import org.apache.spark.sql.errors.DataTypeErrors.toSQLId import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.internal.connector.PartitionPredicateImpl import org.apache.spark.sql.types._ import org.apache.spark.util.ArrayImplicits._ @@ -213,6 +214,7 @@ case l: V2Literal[_] => Some(Literal(l.value, l.dataType)) case r: NamedReference => Some(UnresolvedAttribute(r.fieldNames.toImmutableArraySeq)) case c: V2Cast => toCatalyst(c.expression).map(Cast(_, c.dataType, ansiEnabled = true)) + case p: PartitionPredicateImpl => Some(p.expression) case e: GeneralScalarExpression => convertScalarExpr(e) case _ => None }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionColumnReferenceImpl.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionColumnReferenceImpl.scala new file mode 100644 index 0000000..af06920 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionColumnReferenceImpl.scala
@@ -0,0 +1,29 @@ +/* + * 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.internal.connector + +import org.apache.spark.sql.connector.expressions.PartitionColumnReference + +/** + * Implementation of [[PartitionColumnReference]] that carries the position ordinal in + * Table.partitioning() and the partition column name(s) for that position. + */ +private[connector] case class PartitionColumnReferenceImpl( + ordinal: Int, + fieldNames: Array[String]) + extends PartitionColumnReference
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImpl.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImpl.scala new file mode 100644 index 0000000..6c887ab --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImpl.scala
@@ -0,0 +1,112 @@ +/* + * 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.internal.connector + +import org.apache.spark.SparkException +import org.apache.spark.internal.{Logging, LogKeys} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BoundReference, Expression => CatalystExpression, Predicate => CatalystPredicate} +import org.apache.spark.sql.connector.expressions.NamedReference +import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate + +/** + * An implementation for [[PartitionPredicate]] that wraps a Catalyst Expression representing a + * partition filter. + */ +class PartitionPredicateImpl private ( + private val catalystExpr: CatalystExpression, + private val partitionSchema: Seq[AttributeReference]) + extends PartitionPredicate with Logging { + + /** The wrapped partition filter Catalyst Expression. */ + def expression: CatalystExpression = catalystExpr + + /** Bound predicate, computed once and reused for all partition rows. */ + @transient private lazy val boundPredicate: InternalRow => Boolean = { + val boundExpr = catalystExpr.transform { + case a: AttributeReference => + val index = partitionSchema.indexWhere(_.name == a.name) + require(index >= 0, s"Column ${a.name} not found in partition schema") + BoundReference(index, partitionSchema(index).dataType, nullable = a.nullable) + } + val predicate = CatalystPredicate.createInterpreted(boundExpr) + predicate.eval + } + + override def eval(partitionValues: InternalRow): Boolean = { + if (partitionValues.numFields != partitionSchema.length) { + logWarning( + log"Cannot evaluate partition predicate ${MDC(LogKeys.EXPR, catalystExpr.sql)}: " + + log"partition value field count (${MDC(LogKeys.COUNT, partitionValues.numFields)}) " + + log"does not match schema (${MDC(LogKeys.NUM_PARTITIONS, partitionSchema.length)}). " + + log"Including partition in scan result to avoid incorrect filtering.") + return true + } + + try { + boundPredicate(partitionValues) + } catch { + case e: Exception => + logWarning( + log"Failed to evaluate partition predicate ${MDC(LogKeys.EXPR, catalystExpr.sql)}. " + + log"Including partition in scan result to avoid incorrect filtering.", + e) + true + } + } + + @transient override lazy val references: Array[NamedReference] = { + val refNames = catalystExpr.references.map(_.name).toSet + partitionSchema.zipWithIndex + .filter { case (attr, _) => refNames.contains(attr.name) } + .map { case (attr, ordinal) => PartitionColumnReferenceImpl(ordinal, Array(attr.name)) } + .toArray + } + + override def equals(obj: Any): Boolean = obj match { + case other: PartitionPredicateImpl => + catalystExpr.semanticEquals(other.catalystExpr) && partitionSchema == other.partitionSchema + case _ => false + } + + override def hashCode(): Int = { + 31 * catalystExpr.semanticHash() + partitionSchema.hashCode() + } + + override def toString(): String = s"PartitionPredicate(${catalystExpr.sql})" +} + +object PartitionPredicateImpl { + + def apply( + catalystExpr: CatalystExpression, + partitionSchema: Seq[AttributeReference]): PartitionPredicateImpl = { + if (partitionSchema.isEmpty) { + throw SparkException.internalError( + s"Cannot evaluate partition predicate ${catalystExpr.sql}: partition schema is empty") + } + val partitionNames = partitionSchema.map(_.name).toSet + val refNames = catalystExpr.references.map(_.name).toSet + if (!refNames.subsetOf(partitionNames)) { + throw SparkException.internalError( + s"Cannot evaluate partition predicate ${catalystExpr.sql}: expression references " + + s"${refNames.mkString(", ")} not all in partition columns ${partitionNames.mkString(", ")}") + } + new PartitionPredicateImpl(catalystExpr, partitionSchema) + } +}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryEnhancedPartitionFilterTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryEnhancedPartitionFilterTable.scala new file mode 100644 index 0000000..507439b --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryEnhancedPartitionFilterTable.scala
@@ -0,0 +1,168 @@ +/* + * 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.connector.catalog + +import java.util + +import scala.collection.mutable.{ArrayBuffer, Buffer} + +import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.MultipartIdentifierHelper +import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.connector.expressions.filter.{PartitionPredicate, Predicate} +import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters} +import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder} +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.util.ArrayImplicits._ + +/** + * In-memory table whose scan builder implements enhanced partition filtering using + * PartitionPredicates pushed in a second pass. + */ +class InMemoryEnhancedPartitionFilterTable( + name: String, + columns: Array[Column], + partitioning: Array[Transform], + properties: util.Map[String, String]) + extends InMemoryTable(name, columns, partitioning, properties) { + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new InMemoryEnhancedPartitionFilterScanBuilder(schema()) + } + + override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = { + InMemoryBaseTable.maybeSimulateFailedTableWrite(new CaseInsensitiveStringMap(properties)) + InMemoryBaseTable.maybeSimulateFailedTableWrite(info.options) + new InMemoryWriterBuilderWithOverWrite(info) + } + + class InMemoryEnhancedPartitionFilterScanBuilder( + tableSchema: StructType) + extends ScanBuilder + with SupportsPushDownV2Filters + with SupportsPushDownRequiredColumns { + + private var readSchema: StructType = tableSchema + private val partitionPredicates: Buffer[PartitionPredicate] = ArrayBuffer.empty + private val firstPassPushedPredicates: Buffer[Predicate] = ArrayBuffer.empty + + // Default true: accept (push down) partition predicates. + // Set to false to return them for post-scan. + private val acceptPartitionPredicates = + InMemoryEnhancedPartitionFilterTable.this.properties.getOrDefault( + InMemoryEnhancedPartitionFilterTable.AcceptPartitionPredicatesKey, "true") + .toBoolean + + private val acceptDataPredicates = + InMemoryEnhancedPartitionFilterTable.this.properties.getOrDefault( + InMemoryEnhancedPartitionFilterTable.AcceptDataPredicatesKey, "false") + .toBoolean + + override def supportsIterativePushdown(): Boolean = true + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + val partNames = InMemoryEnhancedPartitionFilterTable.this.partCols.flatMap(_.toSeq).toSet + def referencesOnlyPartitionCols(p: Predicate): Boolean = + p.references().forall(ref => partNames.contains(ref.fieldNames().mkString("."))) + def referencesOnlyDataCols(p: Predicate): Boolean = + p.references().forall(ref => !partNames.contains(ref.fieldNames().mkString("."))) + + val returned = ArrayBuffer.empty[Predicate] + + predicates.foreach { + case p: PartitionPredicate => + if (acceptPartitionPredicates) { + partitionPredicates += p + } else { + returned += p + } + case p if referencesOnlyPartitionCols(p) && + InMemoryTableWithV2Filter.supportsPredicates(Array(p)) => + if (acceptPartitionPredicates) { + firstPassPushedPredicates += p + } else { + returned += p + } + case p if acceptDataPredicates && referencesOnlyDataCols(p) => + // Accept: we are mocking a data source that can evaluate this data predicate + firstPassPushedPredicates += p + case p => + returned += p + } + + returned.toArray + } + + override def pushedPredicates(): Array[Predicate] = + (firstPassPushedPredicates ++ partitionPredicates.map(p => p: Predicate)).toArray + + override def pruneColumns(requiredSchema: StructType): Unit = { + readSchema = requiredSchema + } + + override def build(): Scan = { + val allPartitions = data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq + val partNames = + InMemoryEnhancedPartitionFilterTable.this.partCols.map(_.toSeq.quoted) + .toImmutableArraySeq + val partNamesSet = InMemoryEnhancedPartitionFilterTable.this.partCols.flatMap(_.toSeq).toSet + // Only partition predicates can be used for partition key filtering (filtersToKeys). + val firstPassPartitionPredicates = firstPassPushedPredicates.filter { p => + p.references().forall(ref => partNamesSet.contains(ref.fieldNames().mkString("."))) + } + val allKeys = allPartitions.map(_.asInstanceOf[BufferedRows].key) + val matchingKeys = InMemoryTableWithV2Filter.filtersToKeys( + allKeys, partNames, firstPassPartitionPredicates.toArray).toSet + val filteredByFirstPass = allPartitions.filter(p => + matchingKeys.contains(p.asInstanceOf[BufferedRows].key)) + val filteredBySecondPass = filteredByFirstPass.filter { p => + val partRow = p.asInstanceOf[BufferedRows].partitionKey() + partitionPredicates.forall(_.eval(partRow)) + } + InMemoryEnhancedPartitionFilterBatchScan( + filteredBySecondPass, readSchema, tableSchema, partitionPredicates.toSeq) + } + + } + + /** + * Batch scan that stores pushed PartitionPredicates. + */ + case class InMemoryEnhancedPartitionFilterBatchScan( + _data: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + pushedPartitionPredicates: Seq[PartitionPredicate] = Seq.empty) + extends BatchScanBaseClass(_data, readSchema, tableSchema) { + + def getPushedPartitionPredicates: Seq[PartitionPredicate] = pushedPartitionPredicates + } +} + +object InMemoryEnhancedPartitionFilterTable { + /** + * Table property: when "true", accept (do not return) all PartitionPredicates for pushdown. + */ + private[catalog] val AcceptPartitionPredicatesKey = "accept-partition-predicates" + + /** + * Table property: when "true", accept (do not return) data predicates for pushdown (we are + * mocking a data source that can evaluate this particular data predicate). + */ + private[catalog] val AcceptDataPredicatesKey = "accept-data-predicates" +}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableEnhancedPartitionFilterCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableEnhancedPartitionFilterCatalog.scala new file mode 100644 index 0000000..150b406 --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableEnhancedPartitionFilterCatalog.scala
@@ -0,0 +1,49 @@ +/* + * 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.connector.catalog + +import java.util + +import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException +import org.apache.spark.sql.connector.expressions.Transform + +class InMemoryTableEnhancedPartitionFilterCatalog extends InMemoryTableCatalog { + import CatalogV2Implicits._ + + override def createTable( + ident: Identifier, + columns: Array[Column], + partitions: Array[Transform], + properties: util.Map[String, String]): Table = { + if (tables.containsKey(ident)) { + throw new TableAlreadyExistsException(ident.asMultipartIdentifier) + } + + InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties) + + val tableName = s"$name.${ident.quoted}" + val table = new InMemoryEnhancedPartitionFilterTable(tableName, columns, partitions, properties) + tables.put(ident, table) + namespaces.putIfAbsent(ident.namespace.toList, Map()) + table + } + + override def createTable(ident: Identifier, tableInfo: TableInfo): Table = { + createTable(ident, tableInfo.columns(), tableInfo.partitions(), tableInfo.properties()) + } +}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImplSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImplSuite.scala new file mode 100644 index 0000000..432d0df --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImplSuite.scala
@@ -0,0 +1,68 @@ +/* + * 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.internal.connector + +import org.apache.spark.{SparkConf, SparkFunSuite} +import org.apache.spark.serializer.{JavaSerializer, KryoSerializer, SerializerInstance} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, GreaterThan, Literal} +import org.apache.spark.sql.connector.expressions.PartitionColumnReference +import org.apache.spark.sql.types.IntegerType + +class PartitionPredicateImplSuite extends SparkFunSuite { + + test("Kryo serialization: PartitionPredicateImpl works after round-trip") { + val conf = new SparkConf() + val serializer = new KryoSerializer(conf).newInstance() + checkPartitionPredicateImplAfterSerialization(serializer) + } + + test("Java serialization: PartitionPredicateImpl works after round-trip") { + val conf = new SparkConf() + val serializer = new JavaSerializer(conf).newInstance() + checkPartitionPredicateImplAfterSerialization(serializer) + } + + private def checkPartitionPredicateImplAfterSerialization( + serializer: SerializerInstance): Unit = { + val partitionSchema = Seq(AttributeReference("p", IntegerType)()) + val ref = AttributeReference("p", IntegerType)() + val expr = GreaterThan(ref, Literal(5)) + val predicate = PartitionPredicateImpl(expr, partitionSchema) + + val deserialized = serializer.deserialize[PartitionPredicateImpl]( + serializer.serialize(predicate)) + + assert(deserialized.eval(InternalRow(10)) === true) + assert(deserialized.eval(InternalRow(3)) === false) + assert(deserialized.eval(InternalRow(5)) === false) + + val expectedRefsWithOrdinals = Seq(("p", 0)) + assert(refsWithOrdinals(predicate.references.toSeq) === expectedRefsWithOrdinals) + assert(refsWithOrdinals(deserialized.references.toSeq) === expectedRefsWithOrdinals) + + assert(deserialized.equals(predicate)) + } + + private def refsWithOrdinals(refs: Seq[AnyRef]): Seq[(String, Int)] = refs.map { + case r: PartitionColumnReference => + (r.fieldNames().mkString("."), r.ordinal()) + case other => + fail(s"Expected PartitionColumnReference, got ${other.getClass.getName}: $other") + } +}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala index 8cf2bd8..8843fe1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala
@@ -28,6 +28,7 @@ import org.apache.spark.sql.connector.write.RowLevelOperation.Command.MERGE import org.apache.spark.sql.execution.datasources.DataSourceStrategy import org.apache.spark.sql.sources.Filter +import org.apache.spark.sql.types.StructType /** * A rule that builds scans for group-based row-level operations. @@ -47,9 +48,10 @@ val table = relation.table.asRowLevelOperationTable val scanBuilder = table.newScanBuilder(relation.options) + val partitionSchema = PushDownUtils.getPartitionPredicateSchema(relation) val (pushedFilters, evaluatedFilters, postScanFilters) = - pushFilters(cond, relation.output, scanBuilder) + pushFilters(cond, relation.output, scanBuilder, partitionSchema) val pushedFiltersStr = if (pushedFilters.isLeft) { pushedFilters.swap @@ -97,13 +99,14 @@ private def pushFilters( cond: Expression, tableAttrs: Seq[AttributeReference], - scanBuilder: ScanBuilder) + scanBuilder: ScanBuilder, + partitionSchema: Option[StructType]) : (Either[Seq[Filter], Seq[V2Filter]], Seq[Expression], Seq[Expression]) = { val (filtersWithSubquery, filtersWithoutSubquery) = findTableFilters(cond, tableAttrs) val (pushedFilters, postScanFiltersWithoutSubquery) = - PushDownUtils.pushFilters(scanBuilder, filtersWithoutSubquery) + PushDownUtils.pushFilters(scanBuilder, filtersWithoutSubquery, partitionSchema) val postScanFilterSetWithoutSubquery = ExpressionSet(postScanFiltersWithoutSubquery) val evaluatedFilters = filtersWithoutSubquery.filterNot { filter =>
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala index 7c7b6d5..4a87a50 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala
@@ -19,27 +19,36 @@ import scala.collection.mutable -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Expression, NamedExpression, SchemaPruning} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Expression, ExpressionSet, NamedExpression, PythonUDF, SchemaPruning, SubqueryExpression} import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.CharVarcharUtils -import org.apache.spark.sql.connector.expressions.SortOrder +import org.apache.spark.sql.connector.expressions.{IdentityTransform, SortOrder, Transform} import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownFilters, SupportsPushDownLimit, SupportsPushDownOffset, SupportsPushDownRequiredColumns, SupportsPushDownTableSample, SupportsPushDownTopN, SupportsPushDownV2Filters} -import org.apache.spark.sql.execution.datasources.DataSourceStrategy +import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, DataSourceUtils} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters +import org.apache.spark.sql.internal.connector.{PartitionPredicateImpl, SupportsPushDownCatalystFilters} import org.apache.spark.sql.sources -import org.apache.spark.sql.types.StructType -import org.apache.spark.util.ArrayImplicits._ +import org.apache.spark.sql.types.{StructField, StructType} +import org.apache.spark.util.ArrayImplicits.SparkArrayOps import org.apache.spark.util.collection.Utils object PushDownUtils { + /** - * Pushes down filters to the data source reader + * Pushes down filters to the data source reader. * + * @param scanBuilder The scan builder to push filters to. + * @param filters Catalyst filter expressions to push down. + * @param partitionSchema The schema of [[Table#partitioning() partitioning]]. + * When set and the scan supports V2 filters, + * [[PartitionPredicate]] can be pushed for a second pass. * @return pushed filter and post-scan filters. */ - def pushFilters(scanBuilder: ScanBuilder, filters: Seq[Expression]) + def pushFilters( + scanBuilder: ScanBuilder, + filters: Seq[Expression], + partitionSchema: Option[StructType]) : (Either[Seq[sources.Filter], Seq[Predicate]], Seq[Expression]) = { scanBuilder match { case r: SupportsPushDownFilters => @@ -76,13 +85,12 @@ (postScanFilters ++ untranslatableExprs).toImmutableArraySeq) case r: SupportsPushDownV2Filters => - // A map from translated data source leaf node filters to original catalyst filter - // expressions. For a `And`/`Or` predicate, it is possible that the predicate is partially - // pushed down. This map can be used to construct a catalyst filter expression from the - // input filter, or a superset(partial push down filter) of the input filter. + // Divide the filters into those translatable and untranslatable to data source filters. + // For the translated filters, we will try to push them down to the data source, + // and the data source will return the filters that it cannot guarantee to be true + // for all returned rows. val translatedFilterToExpr = mutable.HashMap.empty[Predicate, Expression] val translatedFilters = mutable.ArrayBuffer.empty[Predicate] - // Catalyst filter expression that can't be translated to data source filters. val untranslatableExprs = mutable.ArrayBuffer.empty[Expression] for (filterExpr <- filters) { @@ -96,18 +104,20 @@ } } - // Data source filters that need to be evaluated again after scanning. which means - // the data source cannot guarantee the rows returned can pass these filters. - // As a result we must return it so Spark can plan an extra filter operator. - val postScanFilters = r.pushPredicates(translatedFilters.toArray).map { predicate => + val rejectedFilters = r.pushPredicates(translatedFilters.toArray).map { predicate => DataSourceV2Strategy.rebuildExpressionFromFilter(predicate, translatedFilterToExpr) } - // Normally translated filters (postScanFilters) are simple filters that can be evaluated - // faster, while the untranslated filters are complicated filters that take more time to - // evaluate, so we want to evaluate the postScanFilters filters first. - (Right(r.pushedPredicates.toImmutableArraySeq), - (postScanFilters ++ untranslatableExprs).toImmutableArraySeq) + val remainingFilters = (rejectedFilters ++ untranslatableExprs).toSeq + val postScanFilters = if (partitionSchema.isEmpty || !r.supportsIterativePushdown) { + remainingFilters + } else { + pushPartitionPredicates(r, partitionSchema.get, remainingFilters) + } + + val orderedPostScanFilters = prioritizeFilters(postScanFilters, + ExpressionSet(untranslatableExprs)) + (Right(r.pushedPredicates.toImmutableArraySeq), orderedPostScanFilters) case r: SupportsPushDownCatalystFilters => val postScanFilters = r.pushFilters(filters) (Right(r.pushedFilters.toImmutableArraySeq), postScanFilters) @@ -116,6 +126,79 @@ } /** + * Normally translated filters (postScanFilters) are simple filters that can be + * evaluated faster, while the untranslated filters are complicated filters + * that take more time to evaluate, so we want to evaluate the translatable + * filters first. + */ + private def prioritizeFilters( + filters: Seq[Expression], + untranslatableFilterSet: ExpressionSet): Seq[Expression] = { + val (translatable, untranslatable) = filters.partition(!untranslatableFilterSet.contains(_)) + translatable ++ untranslatable + } + + /** + * If the scan supports iterative filtering, convert partition filters to + * PartitionPredicates (see SPARK-55596) and push them down in another pass. + */ + private def pushPartitionPredicates( + scanBuilder: SupportsPushDownV2Filters, + partitionSchema: StructType, + remainingFilters: Seq[Expression]): Seq[Expression] = { + val (partitionFilters, nonPartitionFilters) = + DataSourceUtils.getPartitionFiltersAndDataFilters(partitionSchema, remainingFilters) + val (pushable, nonPushable) = partitionFilters.partition(isPushablePartitionFilter) + val partitionAttrs = toAttributes(partitionSchema) + val partitionPredicates = pushable.map(expr => PartitionPredicateImpl(expr, partitionAttrs)) + val rejectedPartitionFilters = scanBuilder.pushPredicates(partitionPredicates.toArray).map { + predicate => predicate.asInstanceOf[PartitionPredicateImpl].expression + } + nonPartitionFilters ++ nonPushable ++ rejectedPartitionFilters + } + + /** + * Returns a table's partitioning expression schema as a StructType, if creation of a + * PartitionPredicate is supported for the schema. + * Currently only supported if all partitioning expressions are identity transforms on simple + * (single-name, non-nested) field references. + * + * @return Some(StructType) representing partition transform expression types, if schema + * is supported for PartitionPredicate. None if not supported. + */ + def getPartitionPredicateSchema(relation: DataSourceV2Relation): Option[StructType] = { + val transforms = relation.table.partitioning + val fields = transforms.flatMap(toSupportedPartitionField(_, relation)) + Option.when(transforms.nonEmpty && fields.length == transforms.length)(StructType(fields)) + } + + /** + * Returns a StructField for the given partition transform if it is + * supported for iterative partition predicate push down. + */ + private def toSupportedPartitionField( + transform: Transform, + relation: DataSourceV2Relation): Option[StructField] = { + transform match { + case t: IdentityTransform if t.ref.fieldNames.length == 1 => + val colName = t.ref.fieldNames.head + relation.output + .find(_.name == colName) + .map(attr => StructField(colName, attr.dataType, attr.nullable)) + case _ => + None + } + } + + /** + * Returns true if the given filter expression is safe to push as a partition predicate + * when using iterative pushdown: it must be deterministic, contain + * no subquery, and no PythonUDF. + */ + private def isPushablePartitionFilter(f: Expression): Boolean = + f.deterministic && !SubqueryExpression.hasSubquery(f) && !f.exists(_.isInstanceOf[PythonUDF]) + + /** * Pushes down TableSample to the data source Scan */ def pushTableSample(scanBuilder: ScanBuilder, sample: TableSampleInfo): Boolean = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala index adfd5ce..9a25752 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
@@ -80,8 +80,9 @@ // `pushedFilters` will be pushed down and evaluated in the underlying data sources. // `postScanFilters` need to be evaluated after the scan. // `postScanFilters` and `pushedFilters` can overlap, e.g. the parquet row group filter. + val partitionSchema = PushDownUtils.getPartitionPredicateSchema(sHolder.relation) val (pushedFilters, postScanFiltersWithoutSubquery) = PushDownUtils.pushFilters( - sHolder.builder, normalizedFiltersWithoutSubquery) + sHolder.builder, normalizedFiltersWithoutSubquery, partitionSchema) val pushedFiltersStr = if (pushedFilters.isLeft) { pushedFilters.swap .getOrElse(throw new NoSuchElementException("The left node doesn't have pushedFilters"))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedPartitionFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedPartitionFilterSuite.scala new file mode 100644 index 0000000..be6e78f --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedPartitionFilterSuite.scala
@@ -0,0 +1,445 @@ +/* + * 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.connector + +import java.util.Locale + +import org.scalatest.BeforeAndAfter + +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.{Expression, In, PredicateHelper, ScalaUDF} +import org.apache.spark.sql.connector.catalog.BufferedRows +import org.apache.spark.sql.connector.catalog.InMemoryEnhancedPartitionFilterTable +import org.apache.spark.sql.connector.catalog.InMemoryTableEnhancedPartitionFilterCatalog +import org.apache.spark.sql.connector.expressions.PartitionColumnReference +import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate +import org.apache.spark.sql.execution.ExplainUtils.stripAQEPlan +import org.apache.spark.sql.execution.FilterExec +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Tests for enhanced partition filter pushdown with tables whose scan builder handles + * PartitionPredicates in a second pass of partition filter pushdown, for those + * Catalyst Expression filters that are not translatable to DSV2, or are returned by DSV2 + * in the first pushdown. + * + * Pushdown cases (Translated/Untranslatable, Partition Filter/Data Filter, 1st/2nd Pass): + * 1. Translated, Data Filter, 1st Pass Returned -> Post-Scan Filters + * 2. Translated, Data Filter, 1st Pass Accepted -> Pushed Down + * 3. Translated, Partition Filter, 1st Pass Returned, 2nd Pass Returned -> Post-Scan Filters + * 4. Translated, Partition Filter, 1st Pass Returned, 2nd Pass Accepted -> Pushed Down + * 5. Translated, Partition Filter, 1st Pass Accepted -> Pushed Down + * 6. Untranslatable, Data Filter -> Post-Scan Filters + * 7. Untranslatable, Partition Filter, 2nd Pass Returned -> Post-Scan Filters + * 8. Untranslatable, Partition Filter, 2nd Pass Accepted -> Pushed Down + */ +class DataSourceV2EnhancedPartitionFilterSuite + extends QueryTest with SharedSparkSession with BeforeAndAfter with PredicateHelper { + + protected val v2Source = classOf[FakeV2ProviderWithCustomSchema].getName + protected val partFilterTableName = "testpartfilter.t" + + protected def registerCatalog(name: String, clazz: Class[_]): Unit = { + spark.conf.set(s"spark.sql.catalog.$name", clazz.getName) + } + + before { + registerCatalog("testpartfilter", classOf[InMemoryTableEnhancedPartitionFilterCatalog]) + } + + after { + spark.sessionState.catalogManager.reset() + } + + test("case 1: translated data filter returned in first pass is in post-scan") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'keep'), ('a', 'drop'), ('b', 'other')") + + // Translated, Data Filter; 1st Pass Returned -> Post-Scan Filters. + // Returned because it is a data filter (on data column), not a partition filter. + // No filter pushdown + val df = sql(s"SELECT * FROM $partFilterTableName WHERE data = 'keep'") + checkAnswer(df, Seq(Row("a", "keep"))) + assertPushedPartitionPredicates(df, 0) + assertScanReturnsPartitionKeys(df, Set("a", "b")) + } + } + + test("case 2: translated data filter accepted in first pass is pushed down") { + // Translated, Data Filter; 1st Pass Accepted -> Pushed Down (not in post-scan). + withTable(partFilterTableName) { + // We mock a data source that can evaluate and reject + // this data predicate (accept-data-predicates); + // the test uses data IS NOT NULL, which always evaluates to true here. + // No partition filter pushdown + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col) TBLPROPERTIES('accept-data-predicates' = 'true')") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'x'), ('b', 'y'), ('c', 'z')") + + val df = sql(s"SELECT * FROM $partFilterTableName WHERE data IS NOT NULL") + checkAnswer(df, Seq(Row("a", "x"), Row("b", "y"), Row("c", "z"))) + assertPushedPartitionPredicates(df, 0) + assertScanReturnsPartitionKeys(df, Set("a", "b", "c")) + assert(!df.queryExecution.executedPlan.exists(_.isInstanceOf[FilterExec]), + "Data filter accepted in first pass should not appear as post-scan Filter") + } + } + + test("case 3: filter returned in both first and second pass") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col) " + + "TBLPROPERTIES('accept-partition-predicates' = 'false')") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'x'), ('b', 'y'), ('c', 'z')") + + // Translated, Partition Filter; 1st Pass Returned, 2nd Pass Returned + // No partition filter pushdown + val df = sql(s"SELECT * FROM $partFilterTableName WHERE part_col IN ('a')") + checkAnswer(df, Seq(Row("a", "x"))) + assertPushedPartitionPredicates(df, 0) + assertScanReturnsPartitionKeys(df, Set("a", "b", "c")) + } + } + + test("case 4: first pass partition predicate returned by source applied in second pass") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'x'), ('b', 'y'), ('c', 'z')") + + // Translated, Partition Filter; 1st Pass Returned, 2nd Pass Accepted + // Partition Filter pushdown + val df = sql(s"SELECT * FROM $partFilterTableName WHERE part_col IN ('a', 'b')") + checkAnswer(df, Seq(Row("a", "x"), Row("b", "y"))) + assertPushedPartitionPredicates(df, 1) + assertScanReturnsPartitionKeys(df, Set("a", "b")) + assertReferencedPartitionColumnOrdinals(df, Array(0), Array("part_col")) + } + } + + test("case 5: first pass partition filter still works") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'x'), ('b', 'y'), ('c', 'z')") + + // Translated, Partition Filter; 1st Pass Accepted + // Partition Filter pushdown + val df = sql(s"SELECT * FROM $partFilterTableName WHERE part_col = 'b'") + checkAnswer(df, Seq(Row("b", "y"))) + assertPushedPartitionPredicates(df, 0) + assertScanReturnsPartitionKeys(df, Set("b")) + } + } + + test("case 6: untranslatable data filters are applied after scan") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $partFilterTableName VALUES " + + "('a', 'keep'), ('a', 'drop'), ('b', 'keep'), ('b', 'other')") + + // Untranslatable, Data Filter + // No filter pushdown + spark.udf.register("is_keep", (s: String) => s != null && s == "keep") + + val df = sql(s"SELECT * FROM $partFilterTableName WHERE is_keep(data)") + checkAnswer(df, Seq(Row("a", "keep"), Row("b", "keep"))) + assertPushedPartitionPredicates(df, 0) + assertScanReturnsPartitionKeys(df, Set("a", "b")) + } + } + + test("case 7: untranslatable partition filter returned in second pass is in post-scan") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col) " + + "TBLPROPERTIES('accept-partition-predicates' = 'false')") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'x'), ('b', 'y'), ('bc', 'z')") + + // Untranslatable, Partition Filter; 2nd Pass Returned + // No partition filter pushdown + val df = sql(s"SELECT * FROM $partFilterTableName WHERE part_col LIKE 'b%'") + checkAnswer(df, Seq(Row("b", "y"), Row("bc", "z"))) + assertPushedPartitionPredicates(df, 0) + assertScanReturnsPartitionKeys(df, Set("a", "b", "bc")) + } + } + + test("case 8: untranslatable partition-only expression handled by second pass") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'x'), ('b', 'y'), ('bc', 'z')") + + // Untranslatable, Partition Filter; 2nd Pass Accepted + // Partition Filter push down + val df = sql(s"SELECT * FROM $partFilterTableName WHERE part_col LIKE 'b%'") + checkAnswer(df, Seq(Row("b", "y"), Row("bc", "z"))) + assertPushedPartitionPredicates(df, 1) + assertScanReturnsPartitionKeys(df, Set("b", "bc")) + assertReferencedPartitionColumnOrdinals(df, Array(0), Array("part_col")) + } + } + + test("case 8: Second-pass PartitionPredicate filter works for UDF filter on partition column") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'x'), ('A', 'y'), ('b', 'z')") + + spark.udf.register("my_upper", (s: String) => + if (s == null) null else s.toUpperCase(Locale.ROOT)) + + // Untranslatable, Partition Filter; 2nd Pass Accepted + // Partition Filter push down + val df = sql(s"SELECT * FROM $partFilterTableName WHERE my_upper(part_col) = 'A'") + checkAnswer(df, Seq(Row("a", "x"), Row("A", "y"))) + assertPushedPartitionPredicates(df, 1) + assertScanReturnsPartitionKeys(df, Set("a", "A")) + assertReferencedPartitionColumnOrdinals(df, Array(0), Array("part_col")) + } + } + + test("referenced partition column ordinals: partition predicate same column twice " + + "has de-duped ordinals") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'x'), ('b', 'y'), ('bc', 'z')") + + val df = sql(s"SELECT * FROM $partFilterTableName WHERE part_col LIKE 'b%' " + + "OR part_col = 'a'") + checkAnswer(df, Seq(Row("a", "x"), Row("b", "y"), Row("bc", "z"))) + assertPushedPartitionPredicates(df, 1) + assertScanReturnsPartitionKeys(df, Set("a", "b", "bc")) + assertReferencedPartitionColumnOrdinals(df, Array(0), Array("part_col")) + } + } + + test("referenced partition column ordinals: one non-first partition column in second-pass") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (p0 string, p1 string, p2 string, data string) " + + s"USING $v2Source PARTITIONED BY (p0, p1, p2)") + sql(s"INSERT INTO $partFilterTableName VALUES " + + "('a', 'x', '1', 'd1'), ('a', 'y', '1', 'd2'), " + + "('a', 'x', '2', 'd3'), ('b', 'x', '1', 'd4')") + + // Untranslatable, Partition Filter; 2nd Pass Accepted + // Partition filter push down + val df = sql(s"SELECT * FROM $partFilterTableName WHERE p1 LIKE 'x%'") + checkAnswer(df, Seq( + Row("a", "x", "1", "d1"), Row("a", "x", "2", "d3"), Row("b", "x", "1", "d4"))) + assertPushedPartitionPredicates(df, 1) + assertReferencedPartitionColumnOrdinals(df, Array(1), Array("p0", "p1", "p2")) + } + } + + test("referenced partition column ordinals: two non-first partition columns in second-pass") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (p0 string, p1 string, p2 string, data string) " + + s"USING $v2Source PARTITIONED BY (p0, p1, p2)") + sql(s"INSERT INTO $partFilterTableName VALUES " + + "('a', 'x', '1', 'd1'), ('a', 'y', '1', 'd2'), " + + "('a', 'x', '2', 'd3'), ('b', 'x', '1', 'd4')") + + spark.udf.register("concat2", (a: String, b: String) => + if (a == null || b == null) null else a + b) + + // Untranslatable, Partition Filter; 2nd Pass Accepted + // Partition filter pushdown + val df = sql(s"SELECT * FROM $partFilterTableName WHERE concat2(p1, p2) = 'x1'") + checkAnswer(df, Seq(Row("a", "x", "1", "d1"), Row("b", "x", "1", "d4"))) + assertPushedPartitionPredicates(df, 1) + assertReferencedPartitionColumnOrdinals(df, Array(1, 2), Array("p0", "p1", "p2")) + } + } + + test("non-deterministic partition filter not pushed as PartitionPredicate") { + // Same checks as FileSourceStrategy/PruneFileSourcePartitions: non-deterministic + // partition filters must not be pushed as PartitionPredicate; they are applied after scan. + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'x'), ('b', 'y'), ('c', 'z')") + + spark.udf.register("nondet_identity", udf((s: String) => s).asNondeterministic()) + + val df = sql(s"SELECT * FROM $partFilterTableName WHERE nondet_identity(part_col) = 'a'") + checkAnswer(df, Seq(Row("a", "x"))) + assertPushedPartitionPredicates(df, 0) + } + } + + test("partition filter with subquery is not pushed as PartitionPredicate") { + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $partFilterTableName VALUES ('a', 'x'), ('b', 'y'), ('c', 'z')") + + withView("subq") { + sql("CREATE TEMP VIEW subq AS SELECT 'a' AS c") + val df = sql(s"SELECT * FROM $partFilterTableName WHERE part_col IN (SELECT c FROM subq)") + checkAnswer(df, Seq(Row("a", "x"))) + assertPushedPartitionPredicates(df, 0) + } + } + } + + test("all eight pushdown cases: translatable filters before untranslatable " + + "in post-scan Filter") { + // Cases 1, 3, 6, 7 end in post-scan (accept-partition-predicates=false). + // Post-scan filters are ordered with translatable first. + withTable(partFilterTableName) { + sql(s"CREATE TABLE $partFilterTableName (part_col string, data string) USING $v2Source " + + "PARTITIONED BY (part_col) TBLPROPERTIES('accept-partition-predicates' = 'false')") + sql(s"INSERT INTO $partFilterTableName VALUES " + + "('a', 'keep'), ('a', 'drop'), ('A', 'keep'), ('b', 'keep'), ('b', 'other')") + + spark.udf.register("is_keep", (s: String) => s != null && s == "keep") + spark.udf.register("my_upper", (s: String) => + if (s == null) null else s.toUpperCase(Locale.ROOT)) + + val df = sql(s"SELECT * FROM $partFilterTableName WHERE " + + "part_col IN ('a', 'A') AND data = 'keep' AND is_keep(data) AND my_upper(part_col) = 'A'") + checkAnswer(df, Seq(Row("a", "keep"), Row("A", "keep"))) + assertScanReturnsPartitionKeys(df, Set("a", "A", "b")) + assertTranslatableBeforeUntranslatableInPostScan(df) + + // Reversed filter order in the query; post-scan order still translatable first. + val dfReversed = sql(s"SELECT * FROM $partFilterTableName WHERE " + + "my_upper(part_col) = 'A' AND is_keep(data) AND data = 'keep' AND part_col IN ('a', 'A')") + checkAnswer(dfReversed, Seq(Row("a", "keep"), Row("A", "keep"))) + assertScanReturnsPartitionKeys(dfReversed, Set("a", "A", "b")) + assertTranslatableBeforeUntranslatableInPostScan(dfReversed) + } + } + + private def assertTranslatableBeforeUntranslatableInPostScan(df: DataFrame): Unit = { + val postScanFilterExec = df.queryExecution.executedPlan.collect { + case f @ FilterExec(_, _) if f.exists(_.isInstanceOf[BatchScanExec]) => f + }.headOption.getOrElse(fail("Expected a post-scan FilterExec above BatchScanExec")) + + val predicates = splitConjunctivePredicates(postScanFilterExec.condition) + // Untranslatable: UDFs and predicates that may be rejected by the scan (e.g. IN) + def isUntranslatable(pred: Expression): Boolean = + pred.exists(_.isInstanceOf[ScalaUDF]) || pred.exists(_.isInstanceOf[In]) + val untranslatableIndices = predicates.indices.filter(i => isUntranslatable(predicates(i))) + val translatableIndices = predicates.indices.filter(i => !isUntranslatable(predicates(i))) + + assert( + untranslatableIndices.isEmpty || translatableIndices.isEmpty || + translatableIndices.max < untranslatableIndices.min, + s"Translatable filters must appear before untranslatable filters in post-scan " + + s"condition; predicates: ${predicates.mkString(", ")}; " + + s"untranslatable indices: $untranslatableIndices, translatable indices: " + + s"$translatableIndices") + } + + /** + * Collects pushed partition predicates from the plan when the scan is our + * test in-memory scan. + */ + private def getPushedPartitionPredicates( + df: DataFrame): Seq[PartitionPredicate] = { + val batchScan = stripAQEPlan(df.queryExecution.executedPlan).collectFirst { + case b: BatchScanExec => b + }.getOrElse(fail("Expected BatchScanExec in plan")) + batchScan.batch match { + case s: InMemoryEnhancedPartitionFilterTable#InMemoryEnhancedPartitionFilterBatchScan => + s.getPushedPartitionPredicates + case _ => Seq.empty + } + } + + /** + * Asserts that the number of pushed partition predicates (second pass) in the plan + * matches the expected count. Use for tests that run a query against the in-memory + * enhanced partition filter table. + */ + private def assertPushedPartitionPredicates( + df: DataFrame, + expectedCount: Int): Unit = { + val predicates = getPushedPartitionPredicates(df) + assert(predicates.size === expectedCount, + s"Expected $expectedCount pushed partition predicate(s), got ${predicates.size}: $predicates") + } + + /** + * Asserts that each pushed partition predicate's references() (PartitionColumnReference, + * each with ordinal()) match the expected ordinals and partition column names. + * + * @param df the query result + * @param expectedOrdinals expected 0-based ordinals from Table.partitioning() + * @param expectedPartitionColumnNames partition column names by ordinal + * (names(ordinal) is the name for that partition column) + */ + private def assertReferencedPartitionColumnOrdinals( + df: DataFrame, + expectedOrdinals: Array[Int], + expectedPartitionColumnNames: Array[String]): Unit = { + val predicates = getPushedPartitionPredicates(df) + val names = expectedPartitionColumnNames + predicates.foreach { p => + val refs = p.references() + val ordinals = refs.map(_.asInstanceOf[PartitionColumnReference].ordinal()).sorted + assert(ordinals.sameElements(expectedOrdinals.sorted), + s"Expected references().map(_.ordinal()) " + + s"${expectedOrdinals.sorted.mkString("[", ", ", "]")}, " + + s"got ${ordinals.mkString("[", ", ", "]")}") + + refs.foreach { ref => + assert(ref.isInstanceOf[PartitionColumnReference], + s"Expected PartitionColumnReference, got ${ref.getClass.getName}") + val partRef = ref.asInstanceOf[PartitionColumnReference] + assert(partRef.fieldNames().nonEmpty, + s"PartitionColumnReference.ordinal=${partRef.ordinal()} has empty fieldNames") + assert(partRef.ordinal() < names.length, + s"PartitionColumnReference.ordinal=${partRef.ordinal()} " + + s"out of range for names length ${names.length}") + val expectedName = names(partRef.ordinal()) + val actualName = partRef.fieldNames().mkString(".") + assert(actualName === expectedName, + s"PartitionColumnReference.ordinal=${partRef.ordinal()}: " + + s"expected fieldNames '${expectedName}', got '${actualName}'") + } + } + } + + /** + * Asserts that the scan reads exactly the given set of partition keys (single-partition + * column tables use keyString() which is the partition value). + */ + private def assertScanReturnsPartitionKeys( + df: DataFrame, + expectedPartitionKeys: Set[String]): Unit = { + val batchScan = df.queryExecution.executedPlan.collectFirst { + case b: BatchScanExec => b + }.getOrElse(fail("Expected BatchScanExec in plan")) + val partitions = batchScan.batch.planInputPartitions() + assert(partitions.length === expectedPartitionKeys.size, + s"Expected ${expectedPartitionKeys.size} partition(s), got ${partitions.length}") + val partKeys = partitions.map(_.asInstanceOf[BufferedRows].keyString()).toSet + assert(partKeys === expectedPartitionKeys, + s"Partition keys should be $expectedPartitionKeys, got $partKeys") + } +}