You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ge...@apache.org on 2026/03/18 04:07:12 UTC
(spark) branch master updated: [SPARK-55596][SQL] DSV2 Enhanced Partition Stats Filtering
This is an automated email from the ASF dual-hosted git repository.
gengliang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new a4dcd6d14d75 [SPARK-55596][SQL] DSV2 Enhanced Partition Stats Filtering
a4dcd6d14d75 is described below
commit a4dcd6d14d7553226eeefacdd27ceb8ed3449c78
Author: Szehon Ho <sz...@gmail.com>
AuthorDate: Tue Mar 17 21:06:59 2026 -0700
[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 <sz...@gmail.com>
Signed-off-by: Gengliang Wang <ge...@apache.org>
---
.../expressions/PartitionColumnReference.java | 39 ++
.../expressions/filter/PartitionPredicate.java | 92 +++++
.../connector/read/SupportsPushDownV2Filters.java | 39 ++
.../sql/connector/util/V2ExpressionSQLBuilder.java | 7 +
.../catalyst/expressions/V2ExpressionUtils.scala | 2 +
.../connector/PartitionColumnReferenceImpl.scala | 29 ++
.../connector/PartitionPredicateImpl.scala | 112 ++++++
.../InMemoryEnhancedPartitionFilterTable.scala | 168 ++++++++
...MemoryTableEnhancedPartitionFilterCatalog.scala | 49 +++
.../connector/PartitionPredicateImplSuite.scala | 68 ++++
.../GroupBasedRowLevelOperationScanPlanning.scala | 9 +-
.../execution/datasources/v2/PushDownUtils.scala | 127 +++++-
.../datasources/v2/V2ScanRelationPushDown.scala | 3 +-
.../DataSourceV2EnhancedPartitionFilterSuite.scala | 445 +++++++++++++++++++++
14 files changed, 1163 insertions(+), 26 deletions(-)
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 000000000000..ef516545812b
--- /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 000000000000..dbc31aa1a458
--- /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 1fec939aeb47..5c1a04a3a3dc 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 @@ import org.apache.spark.sql.connector.expressions.filter.Predicate;
* 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 @@ public interface SupportsPushDownV2Filters extends ScanBuilder {
/**
* 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 @@ public interface SupportsPushDownV2Filters extends ScanBuilder {
* 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 50921f3de0b4..343221f91a18 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.NullOrdering;
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 @@ public class V2ExpressionSQLBuilder {
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 @@ public class V2ExpressionSQLBuilder {
"_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 fd3d1daae674..c4c6a60ce931 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.{BucketTransform, Cast => V2Ca
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 @@ object V2ExpressionUtils extends SQLConfHelper with Logging {
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 000000000000..af069201483e
--- /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 000000000000..6c887ab7d807
--- /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 000000000000..507439b9d06c
--- /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 000000000000..150b4062ef32
--- /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 000000000000..432d0df7d8d4
--- /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 8cf2bd844bf4..8843fe105ec3 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.read.ScanBuilder
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 @@ object GroupBasedRowLevelOperationScanPlanning extends Rule[LogicalPlan] with Pr
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 @@ object GroupBasedRowLevelOperationScanPlanning extends Rule[LogicalPlan] with Pr
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 7c7b6d5488d8..4a87a50c6576 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 @@ package org.apache.spark.sql.execution.datasources.v2
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 @@ object PushDownUtils {
(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 @@ object PushDownUtils {
}
}
- // 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)
@@ -115,6 +125,79 @@ object PushDownUtils {
}
}
+ /**
+ * 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
*/
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 adfd5ceacd67..9a25752ccada 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 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
// `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 000000000000..be6e78fcebcc
--- /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")
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org