[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/InMemoryTableCatalog.scala

@@ -193,14 +193,39 @@ class BasicInMemoryTableCatalog extends TableCatalog {
     InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties)
 
     val tableName = s"$name.${ident.quoted}"
-    val table = new InMemoryTable(tableName, columns, partitions, properties, constraints,
-      distribution, ordering, requiredNumPartitions, advisoryPartitionSize,
-      distributionStrictlyRequired, numRowsPerSplit)
+    val table = newInMemoryTable(
+      tableName, columns, partitions, properties, constraints, distribution, ordering,
+      requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, numRowsPerSplit,
+      util.UUID.randomUUID().toString)
     tables.put(ident, table)
     namespaces.putIfAbsent(ident.namespace.toList, Map())
     table
   }
 
+  /**
+   * Builds the in-memory table this catalog serves. Subclasses that expose a specialized table
+   * type must override this so both CREATE and ALTER reconstruct the same class.
+   */
+  // scalastyle:off argcount
+  protected def newInMemoryTable(
+      name: String,
+      columns: Array[Column],
+      partitioning: Array[Transform],
+      properties: util.Map[String, String],
+      constraints: Array[Constraint],
+      distribution: Distribution,
+      ordering: Array[SortOrder],
+      requiredNumPartitions: Option[Int],
+      advisoryPartitionSize: Option[Long],
+      distributionStrictlyRequired: Boolean,
+      numRowsPerSplit: Int,
+      id: String): InMemoryBaseTable = {
+    // scalastyle:on argcount
+    new InMemoryTable(name, columns, partitioning, properties, constraints, distribution,
+      ordering, requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired,
+      numRowsPerSplit, id)
+  }
+
   override def alterTable(ident: Identifier, changes: TableChange*): Table = {
     val table = loadTable(ident).asInstanceOf[InMemoryBaseTable]
     val properties = CatalogV2Util.applyPropertiesChanges(table.properties, changes)
@@ -238,22 +263,13 @@ class BasicInMemoryTableCatalog extends TableCatalog {
     val currentVersion = table.version()
     val columnsWithIds = InMemoryBaseTable.assignMissingIds(
       CatalogV2Util.structTypeToV2Columns(schema))
+    val reconstructedId = Option(table.id()).getOrElse(util.UUID.randomUUID().toString)
     val newTable = table match {
-      case _: InMemoryTable =>
-        new InMemoryTable(
-          name = table.name,
-          columns = columnsWithIds,
-          partitioning = finalPartitioning,
-          properties = properties,
-          constraints = constraints,
-          id = table.id)
-          .alterTableWithData(table.data, schemaAfterDrops)
-      case _: InMemoryTableWithV2Filter =>
-        new InMemoryTableWithV2Filter(
-          name = table.name,
-          columns = columnsWithIds,
-          partitioning = finalPartitioning,
-          properties = properties)
+      case _: InMemoryTable | _: InMemoryTableWithV2Filter =>
+        newInMemoryTable(
+          table.name, columnsWithIds, finalPartitioning, properties, constraints,
+          table.distribution, table.ordering, table.numPartitions, table.advisoryPartitionSize,
+          table.isDistributionStrictlyRequired, table.numRowsPerSplit, reconstructedId)
           .alterTableWithData(table.data, schemaAfterDrops)
       case other =>
         throw new UnsupportedOperationException(

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

@@ -19,32 +19,43 @@ 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._
-
-  override def createTable(
-      ident: Identifier,
+/**
+ * Mix-in that constructs [[InMemoryCatalystRuntimeFilterTable]] from the shared in-memory
+ * catalog factory used by both CREATE TABLE and ALTER TABLE.
+ */
+trait InMemoryCatalystRuntimeFilterTableFactory { self: BasicInMemoryTableCatalog =>
+  // scalastyle:off argcount
+  override protected def newInMemoryTable(
+      name: String,
       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 InMemoryCatalystRuntimeFilterTable(
-      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)
+      partitioning: Array[Transform],
+      properties: util.Map[String, String],
+      constraints: Array[Constraint],
+      distribution: Distribution,
+      ordering: Array[SortOrder],
+      requiredNumPartitions: Option[Int],
+      advisoryPartitionSize: Option[Long],
+      distributionStrictlyRequired: Boolean,
+      numRowsPerSplit: Int,
+      id: String): InMemoryBaseTable = {
+    // scalastyle:on argcount
+    new InMemoryCatalystRuntimeFilterTable(
+      name, columns, partitioning, properties, constraints, distribution, ordering,
+      requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, numRowsPerSplit)
   }
 }
+
+class InMemoryTableCatalystRuntimeFilterCatalog extends InMemoryTableCatalog
+    with InMemoryCatalystRuntimeFilterTableFactory
+
+/**
+ * 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
+    with InMemoryCatalystRuntimeFilterTableFactory