[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>

Author: szehon-ho

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();
+}

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);
+}

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 @@ 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;
+  }
 }

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 @@ public String build(Expression expr) {
       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 @@ protected String visitUnexpectedExpr(Expression expr) throws IllegalArgumentExce
       "_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) {

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
   }

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

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)
+  }
+}