[SPARK-59012][SQL] Improve SupportsRuntimeCatalystFiltering test coverage and reject nested filter attributes

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

@@ -210,8 +210,7 @@ case class DataSourceV2ScanRelation(
       case s: SupportsRuntimeCatalystFiltering => s.filterAttributes()
       case _ => Array.empty[NamedReference]
     }
-    AttributeSet(V2ExpressionUtils.resolveRefs[Attribute](
-      filterAttrs.toImmutableArraySeq, this))
+    resolveTopLevelFilterAttrs(filterAttrs)
   }
 
   /**
@@ -224,6 +223,21 @@ case class DataSourceV2ScanRelation(
       case s: SupportsRuntimeCatalystFiltering => s.fullyPushedFilterAttributes()
       case _ => Array.empty[NamedReference]
     }
+    resolveTopLevelFilterAttrs(filterAttrs)
+  }
+
+  /**
+   * Resolves the given runtime-filter references against this relation's output. Both runtime
+   * filtering interfaces require each reference to be a top-level attribute of the read schema,
+   * so a nested reference is rejected.
+   */
+  private def resolveTopLevelFilterAttrs(filterAttrs: Array[NamedReference]): AttributeSet = {
+    filterAttrs.find(_.fieldNames.length > 1).foreach { ref =>
+      throw SparkException.internalError(
+        s"Runtime filter attribute '${ref.fieldNames.mkString(".")}' declared by " +
+        s"${scan.getClass.getName} must be a top-level attribute of the scan read schema, " +
+        "but it is a nested reference.")
+    }
     AttributeSet(V2ExpressionUtils.resolveRefs[Attribute](
       filterAttrs.toImmutableArraySeq, this))
   }

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

@@ -42,8 +42,8 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
    *
    * Spark will call [[filter]] if it can derive a runtime filter for any of these attributes.
    * Each reference must be a top-level attribute present in [[Scan.readSchema]]. Nested
-   * references and attributes pruned out of the read schema fail to resolve when Spark builds
-   * the scan relation.
+   * references are rejected, and attributes pruned out of the read schema fail to resolve, when
+   * Spark builds the scan relation.
    */
   def filterAttributes(): Array[NamedReference]
 
@@ -64,8 +64,8 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
    * predicate over it.
    *
    * Each reference must be a top-level attribute present in [[Scan.readSchema]]. Nested
-   * references and attributes pruned out of the read schema fail to resolve when Spark builds
-   * the scan relation.
+   * references are rejected, and attributes pruned out of the read schema fail to resolve, when
+   * Spark builds the scan relation.
    */
   def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty
 

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, JoinedRow, Literal, MetadataStructFieldWithLogicalName, Predicate => CatalystPredicate}
+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.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}
@@ -718,22 +718,24 @@ abstract class InMemoryBaseTable(
     protected def tableSchema: StructType
 
     private val catalystPredicates = ArrayBuffer.empty[CatalystExpression]
+    private var filterCalls = 0
 
     override def filter(expressions: Array[CatalystExpression]): Unit = {
       catalystPredicates ++= expressions
+      filterCalls += 1
       val partAttrs = partitionAttributes
       if (partAttrs.isEmpty) return
+      val partAttrRefs = partAttrs.map(_._2)
 
-      val resolver = SQLConf.get.resolver
       expressions.foreach { expr =>
-        val remapped = expr.transform {
-          case a: AttributeReference =>
-            partAttrs.find(p => resolver(p.name, a.name)).getOrElse(a)
+        // 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 => partAttrs.exists(_.exprId == r.exprId))) {
-          val bound = BindReferences.bindReference(remapped, partAttrs)
+        if (remapped.references.forall(r => partAttrRefs.exists(_.exprId == r.exprId))) {
+          val bound = BindReferences.bindReference(remapped, partAttrRefs)
           val pred = CatalystPredicate.createInterpreted(bound)
           self.data = self.data.filter { p =>
             try {
@@ -755,15 +757,60 @@ abstract class InMemoryBaseTable(
     /** Predicates recorded by [[filter]], for test assertions only. */
     def pushedCatalystPredicates: Seq[CatalystExpression] = catalystPredicates.toSeq
 
-    /** AttributeReferences matching the partition-key InternalRow field order. */
-    private def partitionAttributes: Seq[AttributeReference] = {
+    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))`
+     */
+    private def partitionAttributes: Seq[(Seq[String], AttributeReference)] = {
       partitioning.flatMap(_.references()).flatMap { ref =>
-        val name = ref.fieldNames.mkString(".")
-        readSchema.find(_.name == name).orElse(tableSchema.find(_.name == name)).map { f =>
-          AttributeReference(f.name, f.dataType, f.nullable)()
+        val path = ref.fieldNames.toImmutableArraySeq
+        readSchema.findNestedField(path).orElse(tableSchema.findNestedField(path)).map {
+          case (_, f) =>
+            path -> AttributeReference(ref.fieldNames.mkString("."), f.dataType, f.nullable)()
         }
       }.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`
+     */
+    private def partitionAttrFor(
+        e: CatalystExpression,
+        partAttrs: Seq[(Seq[String], AttributeReference)]): Option[AttributeReference] = {
+      val resolver = SQLConf.get.resolver
+      partitionKeyPath(e).flatMap { path =>
+        partAttrs.collectFirst {
+          case (parts, attr) if parts.length == path.length &&
+            parts.lazyZip(path).forall((part, name) => resolver(part, name)) => attr
+        }
+      }
+    }
+
+    /**
+     * The name parts `e` reads, or None if it reads neither a column nor a struct field. Each
+     * `GetStructField` ordinal is the field's position in its parent struct. Examples:
+     *   - `AttributeReference(a)` -> `Seq("a")`, the top level column a
+     *   - `GetStructField(AttributeReference(a), 0)` -> `Seq("a", "b")`, the nested column a.b
+     *   - `GetStructField(GetStructField(AttributeReference(a), 0), 0)` -> `Seq("a", "b", "c")`
+     */
+    private def partitionKeyPath(e: CatalystExpression): Option[Seq[String]] = e match {
+      case a: AttributeReference => Some(Seq(a.name))
+      case g: GetStructField =>
+        partitionKeyPath(g.child).map(parent => parent :+ g.childSchema(g.ordinal).name)
+      case _ => None
+    }
   }
 
   case class InMemoryBatchScan(

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

@@ -21,7 +21,9 @@ import java.util
 
 import InMemoryCatalystRuntimeFilterTable._
 
-import org.apache.spark.sql.connector.expressions.{NamedReference, Transform}
+import org.apache.spark.sql.connector.catalog.constraints.Constraint
+import org.apache.spark.sql.connector.distributions.{Distribution, Distributions}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference, SortOrder, Transform}
 import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
@@ -41,8 +43,17 @@ class InMemoryCatalystRuntimeFilterTable(
     name: String,
     columns: Array[Column],
     partitioning: Array[Transform],
-    properties: util.Map[String, String])
-  extends InMemoryTableWithV2Filter(name, columns, partitioning, properties) {
+    properties: util.Map[String, String],
+    constraints: Array[Constraint] = Array.empty,
+    distribution: Distribution = Distributions.unspecified(),
+    ordering: Array[SortOrder] = Array.empty,
+    numPartitions: Option[Int] = None,
+    advisoryPartitionSize: Option[Long] = None,
+    isDistributionStrictlyRequired: Boolean = true,
+    numRowsPerSplit: Int = Int.MaxValue)
+  extends InMemoryTableWithV2Filter(name, columns, partitioning, properties, constraints,
+    distribution, ordering, numPartitions, advisoryPartitionSize, isDistributionStrictlyRequired,
+    numRowsPerSplit) {
 
   override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
     new InMemoryCatalystRuntimeFilterScanBuilder(schema, options)
@@ -74,25 +85,35 @@ class InMemoryCatalystRuntimeFilterTable(
       Option(InMemoryCatalystRuntimeFilterTable.this.properties.get(FilterAttributesKey))
         .map(_.split(",").map(_.trim).toSet)
 
-    override def filterAttributes(): Array[NamedReference] = {
+    private val fullyPushedFilterAttrs: Set[String] = Option(
+      InMemoryCatalystRuntimeFilterTable.this.properties.get(FullyPushedFilterAttributesKey))
+      .map(_.split(",").map(_.trim).toSet)
+      .getOrElse(Set.empty)
+
+    /**
+     * The partition columns, each named by the top level read schema column it lives under, the
+     * form both interface methods require. Columns pruned out of the read schema are dropped,
+     * since neither method may name one. Examples:
+     *   - `PARTITIONED BY (part)` -> `"part"`
+     *   - `PARTITIONED BY (s.nested)` -> `"s"`, the struct column holding the partition field
+     */
+    private def partitionAttrNames: Array[String] = {
       val scanFields = readSchema.fields.map(_.name).toSet
-      partitioning.flatMap(_.references()).filter { ref =>
-        val name = ref.fieldNames.mkString(".")
-        scanFields.contains(name) &&
-          restrictedFilterAttrs.forall(_.contains(name))
-      }
+      partitioning.flatMap(_.references()).map(_.fieldNames.head).distinct
+        .filter(scanFields.contains)
     }
 
-    override def fullyPushedFilterAttributes(): Array[NamedReference] = {
-      val fullyPushedFilterAttrs = Option(
-        InMemoryCatalystRuntimeFilterTable.this.properties.get(FullyPushedFilterAttributesKey))
-        .map(_.split(",").map(_.trim).toSet)
-        .getOrElse(Set.empty)
-      filterAttributes().filter { ref =>
-        fullyPushedFilterAttrs.contains(ref.fieldNames.mkString("."))
-      }
+    override def filterAttributes(): Array[NamedReference] = {
+      partitionAttrNames
+        .filter(name => restrictedFilterAttrs.forall(_.contains(name)))
+        .map(FieldReference.column)
     }
 
+    // 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] = {
+      partitionAttrNames.filter(fullyPushedFilterAttrs.contains).map(FieldReference.column)
+    }
   }
 }
 

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

@@ -20,7 +20,9 @@ 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
+import org.apache.spark.sql.connector.catalog.constraints.Constraint
+import org.apache.spark.sql.connector.distributions.Distribution
+import org.apache.spark.sql.connector.expressions.{SortOrder, Transform}
 
 class InMemoryTableCatalystRuntimeFilterCatalog extends InMemoryTableCatalog {
   import CatalogV2Implicits._
@@ -48,3 +50,41 @@ class InMemoryTableCatalystRuntimeFilterCatalog extends InMemoryTableCatalog {
     createTable(ident, tableInfo.columns(), tableInfo.partitions(), tableInfo.properties)
   }
 }
+
+/**
+ * The [[InMemoryCatalog]] counterpart of [[InMemoryTableCatalystRuntimeFilterCatalog]]: it hands
+ * out tables whose scans take runtime filters as Catalyst expressions, and honors
+ * `numRowsPerSplit` so that a partition key can have several splits.
+ */
+class InMemoryCatalystRuntimeFilterCatalog extends InMemoryCatalog {
+  import CatalogV2Implicits._
+
+  // scalastyle:off argcount
+  override def createTable(
+      ident: Identifier,
+      columns: Array[Column],
+      partitions: Array[Transform],
+      properties: util.Map[String, String],
+      distribution: Distribution,
+      ordering: Array[SortOrder],
+      requiredNumPartitions: Option[Int],
+      advisoryPartitionSize: Option[Long],
+      constraints: Array[Constraint],
+      distributionStrictlyRequired: Boolean,
+      numRowsPerSplit: Int): Table = {
+    // scalastyle:on argcount
+    if (tables.containsKey(ident)) {
+      throw new TableAlreadyExistsException(ident.asMultipartIdentifier)
+    }
+
+    InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties)
+
+    val tableName = s"$name.${ident.quoted}"
+    val table = new InMemoryCatalystRuntimeFilterTable(
+      tableName, columns, partitions, properties, constraints, distribution, ordering,
+      requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, numRowsPerSplit)
+    tables.put(ident, table)
+    namespaces.putIfAbsent(ident.namespace.toList, Map())
+    table
+  }
+}

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

@@ -21,7 +21,9 @@ import java.util
 
 import org.scalatest.Assertions.assert
 
-import org.apache.spark.sql.connector.expressions.{FieldReference, LiteralValue, NamedReference, Transform}
+import org.apache.spark.sql.connector.catalog.constraints.Constraint
+import org.apache.spark.sql.connector.distributions.{Distribution, Distributions}
+import org.apache.spark.sql.connector.expressions.{FieldReference, LiteralValue, NamedReference, SortOrder, Transform}
 import org.apache.spark.sql.connector.expressions.filter.{And, Predicate}
 import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder, SupportsRuntimeV2Filtering}
 import org.apache.spark.sql.connector.write.{LogicalWriteInfo, SupportsOverwriteV2, WriteBuilder, WriterCommitMessage}
@@ -33,8 +35,17 @@ class InMemoryTableWithV2Filter(
     name: String,
     columns: Array[Column],
     partitioning: Array[Transform],
-    properties: util.Map[String, String])
-  extends InMemoryBaseTable(name, columns, partitioning, properties) with SupportsDeleteV2 {
+    properties: util.Map[String, String],
+    constraints: Array[Constraint] = Array.empty,
+    distribution: Distribution = Distributions.unspecified(),
+    ordering: Array[SortOrder] = Array.empty,
+    numPartitions: Option[Int] = None,
+    advisoryPartitionSize: Option[Long] = None,
+    isDistributionStrictlyRequired: Boolean = true,
+    numRowsPerSplit: Int = Int.MaxValue)
+  extends InMemoryBaseTable(name, columns, partitioning, properties, constraints, distribution,
+    ordering, numPartitions, advisoryPartitionSize, isDistributionStrictlyRequired,
+    numRowsPerSplit) with SupportsDeleteV2 {
 
   override def canDeleteWhere(predicates: Array[Predicate]): Boolean = {
     InMemoryTableWithV2Filter.supportsPredicates(predicates)

sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala

@@ -33,7 +33,7 @@ import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._
 import org.apache.spark.sql.catalyst.optimizer.BuildRight
 import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, Inner, LeftAnti}
 import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
-import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog}
+import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableCatalystRuntimeFilterCatalog, InMemoryTableWithV2FilterCatalog}
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive._
 import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
@@ -2690,6 +2690,32 @@ class DynamicPartitionPruningV2FilterSuiteAEOn
     extends DynamicPartitionPruningV2FilterSuite
   with EnableAdaptiveExecutionSuite
 
+/**
+ * Runs the DSv2 dynamic partition pruning tests against scans that receive runtime filters as
+ * Catalyst expressions, via
+ * [[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]], rather than as
+ * connector predicates. This is the counterpart of [[DynamicPartitionPruningV2FilterSuite]],
+ * which covers the same tests for
+ * [[org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering]].
+ */
+abstract class DynamicPartitionPruningV2CatalystFilterSuite
+    extends DynamicPartitionPruningV2Suite {
+
+  override protected def initState(): Unit = {
+    super.initState()
+    spark.conf.set("spark.sql.catalog.testcat",
+      classOf[InMemoryTableCatalystRuntimeFilterCatalog].getName)
+  }
+}
+
+class DynamicPartitionPruningV2CatalystFilterSuiteAEOff
+    extends DynamicPartitionPruningV2CatalystFilterSuite
+  with DisableAdaptiveExecutionSuite
+
+class DynamicPartitionPruningV2CatalystFilterSuiteAEOn
+    extends DynamicPartitionPruningV2CatalystFilterSuite
+  with EnableAdaptiveExecutionSuite
+
 private object DppMaterializedInputTestState {
   private val counters = TrieMap.empty[String, AtomicInteger]