[SPARK-59068][SQL][FOLLOWUP] Correct runtime filter validation and test fixtures

common/utils/src/main/resources/error/error-conditions.json

@@ -2292,6 +2292,11 @@
           "The attribute cannot be resolved."
         ]
       },
+      "NOT_IN_FILTER_ATTRIBUTES" : {
+        "message" : [
+          "The attribute must also be returned by `filterAttributes()`."
+        ]
+      },
       "NOT_TOP_LEVEL" : {
         "message" : [
           "The attribute must be top-level, but it is a nested reference."

sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java

@@ -50,7 +50,8 @@ public interface SupportsRuntimeFiltering extends SupportsRuntimeV2Filtering {
    * The provided expressions must be interpreted as a set of filters that are ANDed together.
    * Implementations may use the filters to prune initially planned {@link InputPartition}s.
    * <p>
-   * Spark tracks runtime-filter eligibility by root attribute. If {@link #filterAttributes()}
+   * Spark currently tracks runtime-filter eligibility by root attribute. If
+   * {@link #filterAttributes()}
    * returns a nested reference, this method may receive a filter on another nested field under
    * the same root. Implementations must inspect each filter and use only filters they can apply.
    * Nested paths are encoded in a V1 {@link Filter} as dot-separated names, with each path part

sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java

@@ -63,7 +63,8 @@ public interface SupportsRuntimeV2Filtering extends Scan {
    * The provided expressions must be interpreted as a set of predicates that are ANDed together.
    * Implementations may use the predicates to prune initially planned {@link InputPartition}s.
    * <p>
-   * Spark tracks runtime-filter eligibility by root attribute. If {@link #filterAttributes()}
+   * Spark currently tracks runtime-filter eligibility by root attribute. If
+   * {@link #filterAttributes()}
    * returns a nested reference, this method may receive a predicate on another nested field under
    * the same root. Implementations must inspect each predicate and use only predicates they can
    * apply.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala

@@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.analysis.{NoSuchFunctionException, Unresolv
 import org.apache.spark.sql.catalyst.encoders.EncoderUtils
 import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke}
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan, SampleMethod}
+import org.apache.spark.sql.catalyst.types.DataTypeUtils.fromAttributes
 import org.apache.spark.sql.connector.catalog.{FunctionCatalog, Identifier}
 import org.apache.spark.sql.connector.catalog.functions._
 import org.apache.spark.sql.connector.catalog.functions.ScalarFunction.MAGIC_METHOD_NAME
@@ -70,6 +71,34 @@ object V2ExpressionUtils extends SQLConfHelper with Logging {
     AttributeSet(resolveRefs[NamedExpression](refs.toImmutableArraySeq, plan))
   }
 
+  /**
+   * Resolves data source runtime-filter attributes and wraps resolution failures with connector
+   * context.
+   */
+  private[sql] def resolveDataSourceRuntimeFilterRefs(
+      refs: Array[NamedReference],
+      output: Seq[Attribute],
+      method: String,
+      scanClass: String): AttributeSet = {
+    if (refs.isEmpty) return AttributeSet.empty
+
+    val plan = LocalRelation(output)
+    val resolvedAttrs = refs.map { ref =>
+      try {
+        resolveRef[NamedExpression](ref, plan)
+      } catch {
+        case e: AnalysisException =>
+          throw QueryCompilationErrors.cannotResolveDataSourceRuntimeFilterAttributeError(
+            attribute = ref.fieldNames,
+            method = method,
+            scanClass = scanClass,
+            relationOutput = fromAttributes(output),
+            cause = e)
+      }
+    }
+    AttributeSet(resolvedAttrs)
+  }
+
   /**
    * Converts the array of input V2 [[V2SortOrder]] into their counterparts in catalyst.
    */

sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala

@@ -4680,6 +4680,19 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat
       None)
   }
 
+  def fullyPushedDataSourceRuntimeFilterAttributeNotFilterableError(
+      attribute: Array[String],
+      scanClass: String,
+      relationOutput: StructType): AnalysisException = {
+    invalidDataSourceRuntimeFilterAttributeError(
+      attribute,
+      "fullyPushedFilterAttributes()",
+      scanClass,
+      relationOutput,
+      "NOT_IN_FILTER_ATTRIBUTES",
+      None)
+  }
+
   private def invalidDataSourceRuntimeFilterAttributeError(
       attribute: Array[String],
       method: String,

sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala

@@ -20,10 +20,9 @@ package org.apache.spark.sql.execution.datasources.v2
 import java.util.{Collections, Optional, OptionalLong}
 
 import org.apache.spark.SparkException
-import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.analysis.{MultiInstanceRelation, NamedRelation, TimeTravelSpec}
 import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatistics}
-import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, Expression, NamedExpression, SortOrder, V2ExpressionUtils}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, Expression, SortOrder, V2ExpressionUtils}
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, ExposesMetadataColumns, Histogram, HistogramBin, LeafNode, LogicalPlan, Statistics}
 import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils
@@ -208,23 +207,31 @@ case class DataSourceV2ScanRelation(
   lazy val runtimeFilterAttrs: AttributeSet = {
     checkRuntimeFilteringInterfaces()
     resolvedFullyPushedRuntimeFilterAttrs
-    val filterAttrs = scan match {
-      case s: SupportsRuntimeV2Filtering => s.filterAttributes
-      case s: SupportsRuntimeCatalystFiltering => s.filterAttributes()
-      case _ => Array.empty[NamedReference]
-    }
-    resolveFilterAttrs(filterAttrs, "filterAttributes()")
+    resolvedRuntimeFilterAttrs
+  }
+
+  private[sql] lazy val declaredRuntimeFilterAttrs: Array[NamedReference] = scan match {
+    case s: SupportsRuntimeV2Filtering => s.filterAttributes
+    case s: SupportsRuntimeCatalystFiltering => s.filterAttributes()
+    case _ => Array.empty
   }
 
   private lazy val declaredFullyPushedRuntimeFilterAttrs: Array[NamedReference] = scan match {
     case s: SupportsRuntimeCatalystFiltering => s.fullyPushedFilterAttributes()
     case _ => Array.empty
   }
 
+  private lazy val resolvedRuntimeFilterAttrs: AttributeSet = {
+    resolveFilterAttrs(declaredRuntimeFilterAttrs, "filterAttributes()")
+  }
+
   private lazy val resolvedFullyPushedRuntimeFilterAttrs: AttributeSet = {
-    checkFullyPushedFilterAttrs()
-    resolveFilterAttrs(
+    checkFullyPushedFilterAttrsAreTopLevel()
+    val resolvedAttrs = resolveFilterAttrs(
       declaredFullyPushedRuntimeFilterAttrs, "fullyPushedFilterAttributes()")
+    resolvedRuntimeFilterAttrs
+    checkFullyPushedFilterAttrsAreFilterable()
+    resolvedAttrs
   }
 
   /**
@@ -245,20 +252,8 @@ case class DataSourceV2ScanRelation(
   private def resolveFilterAttrs(
       filterAttrs: Array[NamedReference],
       method: String): AttributeSet = {
-    val resolvedAttrs = filterAttrs.map { ref =>
-      try {
-        V2ExpressionUtils.resolveRef[NamedExpression](ref, this)
-      } catch {
-        case e: AnalysisException =>
-          throw QueryCompilationErrors.cannotResolveDataSourceRuntimeFilterAttributeError(
-            attribute = ref.fieldNames,
-            method = method,
-            scanClass = scan.getClass.getName,
-            relationOutput = fromAttributes(output),
-            cause = e)
-      }
-    }
-    AttributeSet(resolvedAttrs)
+    V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+      filterAttrs, output, method, scan.getClass.getName)
   }
 
   override val nodePatterns: Seq[TreePattern] = Seq(DATA_SOURCE_V2_SCAN_RELATION)
@@ -319,7 +314,7 @@ case class DataSourceV2ScanRelation(
     }
   }
 
-  private def checkFullyPushedFilterAttrs(): Unit = {
+  private def checkFullyPushedFilterAttrsAreTopLevel(): Unit = {
     declaredFullyPushedRuntimeFilterAttrs.find(_.fieldNames.length > 1).foreach { ref =>
       throw QueryCompilationErrors.nestedDataSourceFullyPushedRuntimeFilterAttributeError(
         attribute = ref.fieldNames,
@@ -328,6 +323,20 @@ case class DataSourceV2ScanRelation(
     }
   }
 
+  private def checkFullyPushedFilterAttrsAreFilterable(): Unit = {
+    declaredFullyPushedRuntimeFilterAttrs.find { fullyPushedRef =>
+      !declaredRuntimeFilterAttrs.exists { filterRef =>
+        fullyPushedRef.fieldNames.length == filterRef.fieldNames.length &&
+          fullyPushedRef.fieldNames.lazyZip(filterRef.fieldNames).forall(conf.resolver)
+      }
+    }.foreach { ref =>
+      throw QueryCompilationErrors.fullyPushedDataSourceRuntimeFilterAttributeNotFilterableError(
+        attribute = ref.fieldNames,
+        scanClass = scan.getClass.getName,
+        relationOutput = fromAttributes(output))
+    }
+  }
+
   override def doCanonicalize(): DataSourceV2ScanRelation = {
     this.copy(
       relation = this.relation.copy(

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala

@@ -78,9 +78,10 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
    * Implementations may use the expressions to prune initially planned
    * [[org.apache.spark.sql.connector.read.InputPartition]]s.
    *
-   * Spark tracks runtime-filter eligibility by root attribute. If [[filterAttributes]] returns a
-   * nested reference, an expression may access another nested field under the same root. The scan
-   * must match each access against its own partition layout and use only expressions it can apply.
+   * Spark currently tracks runtime-filter eligibility by root attribute. If [[filterAttributes]]
+   * returns a nested reference, an expression may access another nested field under the same root.
+   * The scan must match each access against its own partition layout and use only expressions it
+   * can apply.
    *
    * Spark may call this method more than once for the same scan instance: a plan can hold several
    * scan nodes sharing one scan (e.g. the two branches of a group-based UPDATE), and each pushes

sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala

@@ -29,7 +29,7 @@ import scala.collection.mutable.{ArrayBuffer, ListBuffer}
 import scala.jdk.CollectionConverters._
 
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BindReferences, Cast, EvalMode, Expression => CatalystExpression, GenericInternalRow, GetStructField, JoinedRow, Literal, MetadataStructFieldWithLogicalName, Predicate => CatalystPredicate}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BoundReference, Cast, EvalMode, Expression => CatalystExpression, GenericInternalRow, GetStructField, JoinedRow, Literal, MetadataStructFieldWithLogicalName, Predicate => CatalystPredicate}
 import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, ArrayData, CaseInsensitiveMap, CharVarcharUtils, DateTimeUtils, GenericArrayData, MapData, ResolveDefaultColumns}
 import org.apache.spark.sql.connector.catalog.constraints.Constraint
 import org.apache.spark.sql.connector.distributions.{Distribution, Distributions}
@@ -231,6 +231,10 @@ abstract class InMemoryBaseTable(
     }
   }
 
+  protected def identityPartitionReferences: Array[NamedReference] = {
+    partitioning.collect { case IdentityTransform(ref) => ref }
+  }
+
   private val UTC = ZoneId.of("UTC")
   private val EPOCH_LOCAL_DATE = Instant.EPOCH.atZone(UTC).toLocalDate
 
@@ -513,9 +517,10 @@ abstract class InMemoryBaseTable(
   }
 
   private def canEvaluate(filter: Filter): Boolean = {
-    if (partitioning.length == 1 && partitioning.head.references.length == 1) {
+    val identityRefs = identityPartitionReferences
+    if (partitioning.length == 1 && identityRefs.length == 1) {
       filter match {
-        case In(attrName, _) if attrName == partitioning.head.references.head.toString => true
+        case In(attrName, _) if attrName == identityRefs.head.toString => true
         case _ => false
       }
     } else {
@@ -633,6 +638,12 @@ abstract class InMemoryBaseTable(
 
     override def toBatch: Batch = this
 
+    protected def identityPartitionAttributes: Array[NamedReference] = {
+      identityPartitionReferences.distinct
+        .filter(ref => readSchema.findNestedField(
+          ref.fieldNames.toImmutableArraySeq, resolver = SQLConf.get.resolver).isDefined)
+    }
+
     override def estimateStatistics(): Statistics = {
       if (data.isEmpty) {
         return InMemoryStats(OptionalLong.of(0L), OptionalLong.of(0L), new util.HashMap())
@@ -717,11 +728,10 @@ abstract class InMemoryBaseTable(
 
   /**
    * Reference implementation of [[SupportsRuntimeCatalystFiltering.filter]] for the in-memory
-   * fixtures: records what was pushed, and for expressions referencing only partition columns
-   * binds them against the partition key and drops partitions that do not match. Binding and
-   * interpreting rather than pattern matching a fixed set of operators is what lets the fixture
-   * honor an arbitrary pushed expression, the same way `PartitionPredicateImpl` does. Mixing
-   * classes supply their own `filterAttributes()`.
+   * fixtures: records what was pushed, and binds expressions referencing only identity partition
+   * columns against the partition key to drop partitions that do not match. Interpreting the
+   * bound expression lets the fixture honor arbitrary pushed expressions. Mixing classes supply
+   * their own `filterAttributes()`.
    */
   trait CatalystRuntimeFilteringScan extends SupportsRuntimeCatalystFiltering {
     self: BatchScanBaseClass =>
@@ -737,18 +747,15 @@ abstract class InMemoryBaseTable(
       filterCalls += 1
       val partAttrs = partitionAttributes
       if (partAttrs.isEmpty) return
-      val partAttrRefs = partAttrs.map(_._2)
 
       expressions.foreach { expr =>
         // Top down, so `s.part` is rewritten before its `s` child is considered.
         val remapped = expr.transformDown {
           case e => partitionAttrFor(e, partAttrs).getOrElse(e)
         }
-        // Only evaluate expressions whose refs are all partition columns, so we can bind
-        // against the partition key InternalRow (same approach as PartitionPredicateImpl).
-        if (remapped.references.forall(r => partAttrRefs.exists(_.exprId == r.exprId))) {
-          val bound = BindReferences.bindReference(remapped, partAttrRefs)
-          val pred = CatalystPredicate.createInterpreted(bound)
+        // Evaluate expressions only when every reference maps to an identity partition-key slot.
+        if (remapped.references.isEmpty) {
+          val pred = CatalystPredicate.createInterpreted(remapped)
           self.data = self.data.filter { p =>
             try {
               pred.eval(p.asInstanceOf[BufferedRows].partitionKey())
@@ -772,37 +779,33 @@ abstract class InMemoryBaseTable(
     def filterCallCount: Int = filterCalls
 
     /**
-     * The `AttributeReference`s standing for the partition key InternalRow fields, in its field
-     * order, each paired with the name-part sequence of its partition column. The parts are kept
-     * unflattened so a quoted top-level column `a.b` (parts `Seq("a.b")`) stays distinct from a
-     * nested column `a`.`b` (parts `Seq("a", "b")`). Example:
-     *   - `PARTITIONED BY (part, s.nested)` -> `(Seq("part"), AttributeReference(part))`, then
-     *     `(Seq("s", "nested"), AttributeReference(s.nested))`
+     * Identity partition columns paired with their bound partition-key slots.
+     *
+     * Only identity transforms expose a source path because their partition-key slot retains the
+     * source value. Name parts stay separate so a quoted top-level column `a.b` remains distinct
+     * from a nested column `a`.`b`.
      */
-    private def partitionAttributes: Seq[(Seq[String], AttributeReference)] = {
-      partitioning.flatMap(_.references()).flatMap { ref =>
-        val path = ref.fieldNames.toImmutableArraySeq
-        val resolver = SQLConf.get.resolver
-        readSchema.findNestedField(path, resolver = resolver)
-          .orElse(tableSchema.findNestedField(path, resolver = resolver)).map {
-          case (_, f) =>
-            path -> AttributeReference(ref.fieldNames.mkString("."), f.dataType, f.nullable)()
-        }
+    private def partitionAttributes: Seq[(Seq[String], BoundReference)] = {
+      partitioning.zipWithIndex.flatMap {
+        case (IdentityTransform(ref), ordinal) =>
+          val path = ref.fieldNames.toImmutableArraySeq
+          val resolver = SQLConf.get.resolver
+          readSchema.findNestedField(path, resolver = resolver)
+            .orElse(tableSchema.findNestedField(path, resolver = resolver)).map {
+            case (_, f) => path -> BoundReference(ordinal, f.dataType, f.nullable)
+          }
+        case _ => None
       }.toSeq
     }
 
     /**
-     * The partition key `AttributeReference` that `e` reads, or None if `e` reads no partition
-     * column. The path `e` reads is compared to each partition column's name parts component-wise
-     * with the resolver, so a quoted top-level column `a.b` cannot collide with a nested column
-     * `a`.`b`. Examples, under `PARTITIONED BY (part, s.nested)` where `nested` is field 0 of `s`:
-     *   - `AttributeReference(part)` -> `AttributeReference(part)`
-     *   - `GetStructField(AttributeReference(s), 0)` -> `AttributeReference(s.nested)`
-     *   - `AttributeReference(s)` -> None if `s` itself is not a partition column, only `s.nested`
+     * The partition-key slot that `e` reads, or None if `e` reads no identity partition column.
+     * The path `e` reads is compared to each partition column's name parts component-wise with the
+     * resolver, so a quoted top-level column `a.b` cannot collide with a nested column `a`.`b`.
      */
     private def partitionAttrFor(
         e: CatalystExpression,
-        partAttrs: Seq[(Seq[String], AttributeReference)]): Option[AttributeReference] = {
+        partAttrs: Seq[(Seq[String], BoundReference)]): Option[BoundReference] = {
       val resolver = SQLConf.get.resolver
       partitionKeyPath(e).flatMap { path =>
         partAttrs.collectFirst {
@@ -841,14 +844,12 @@ abstract class InMemoryBaseTable(
     var pushedFilters: Array[Filter] = Array.empty
 
     override def filterAttributes(): Array[NamedReference] = {
-      partitioning.flatMap(_.references)
-        .filter(ref => readSchema.findNestedField(
-          ref.fieldNames.toImmutableArraySeq, resolver = SQLConf.get.resolver).isDefined)
+      identityPartitionAttributes
     }
 
     override def filter(filters: Array[Filter]): Unit = {
-      if (partitioning.length == 1 && partitioning.head.references().length == 1) {
-        val ref = partitioning.head.references().head
+      if (partitioning.length == 1 && identityPartitionReferences.length == 1) {
+        val ref = identityPartitionReferences.head
         filters.foreach {
           case In(attrName, values) if attrName == ref.toString =>
             val matchingKeys = values.map { value =>

sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala

@@ -107,7 +107,9 @@ class InMemoryCatalystRuntimeFilterTable(
     // Not intersected with `filterAttributes()`, so a table can declare a fully pushed attribute
     // that is not a filter attribute, a combination the interface forbids.
     override def fullyPushedFilterAttributes(): Array[NamedReference] = {
-      partitionAttrs.filter(ref => fullyPushedFilterAttrs.contains(ref.fieldNames.mkString(".")))
+      identityPartitionAttributes.filter { ref =>
+        fullyPushedFilterAttrs.contains(ref.fieldNames.mkString("."))
+      }
     }
   }
 }