[SPARK-57499][SQL] Fix column pruning and invalid plans in variant extraction pushdown on DSv2 scans

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala

@@ -34,12 +34,12 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
 import org.apache.spark.sql.connector.expressions.{SortOrder => V2SortOrder}
 import org.apache.spark.sql.connector.expressions.aggregate.{Aggregation, Avg, Count, CountStar, Max, Min, Sum}
 import org.apache.spark.sql.connector.expressions.filter.Predicate
-import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownAggregates, SupportsPushDownFilters, SupportsPushDownJoin, SupportsPushDownVariantExtractions, V1Scan, VariantExtraction}
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownAggregates, SupportsPushDownFilters, SupportsPushDownJoin, SupportsPushDownRequiredColumns, SupportsPushDownVariantExtractions, V1Scan, VariantExtraction}
 import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, VariantInRelation, VariantMetadata}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.internal.connector.VariantExtractionImpl
 import org.apache.spark.sql.sources
-import org.apache.spark.sql.types.{DataType, DecimalType, IntegerType, StringType, StructField, StructType}
+import org.apache.spark.sql.types.{DataType, DecimalType, IntegerType, StringType, StructField, StructType, VariantType}
 import org.apache.spark.sql.util.SchemaUtils._
 import org.apache.spark.util.ArrayImplicits._
 import org.apache.spark.util.Utils
@@ -358,11 +358,30 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
     case agg: Aggregate => rewriteAggregate(agg)
   }
 
-  def pushDownVariants(plan: LogicalPlan): LogicalPlan = plan.transformDown {
-    case p@PhysicalOperation(projectList, filters, sHolder @ ScanBuilderHolder(_, _,
-        builder: SupportsPushDownVariantExtractions))
-        if conf.getConf(org.apache.spark.sql.internal.SQLConf.PUSH_VARIANT_INTO_SCAN) =>
-      pushVariantExtractions(p, projectList, filters, sHolder, builder)
+  // Two-visit protocol explanation:
+  // `transformDown` is pre-order and always recurses into children, because
+  // `pushVariantExtractions` returns `originalPlan` unchanged (same object reference).
+  // This means a tree like Project(pl, ScanBuilderHolder) gets visited twice:
+  //   (1) Outer visit: PhysicalOperation collapses Project->ScanBuilderHolder and
+  //       yields the real (projectList, filters, sHolder). This is the authoritative
+  //       visit where column pruning decisions are made.
+  //   (2) Inner leaf visit: PhysicalOperation on a bare ScanBuilderHolder LeafNode
+  //       yields (sHolder.output, Nil, sHolder) -- projectList = full schema output.
+  //       Without a guard, this would add fullVariant for *every* variant column.
+  //
+  // The `pushedVariants.isEmpty` guard prevents the inner visit from re-running once
+  // the outer visit has completed (successfully or not). An empty-mapping sentinel
+  // `Some(new VariantInRelation())` is written by the outer visit when there is
+  // nothing to push; `buildScanWithPushedVariants` requires `mapping.nonEmpty` to
+  // fire, so the sentinel correctly bypasses it while still suppressing the inner visit.
+  def pushDownVariants(plan: LogicalPlan): LogicalPlan = {
+    plan.transformDown {
+      case p@PhysicalOperation(projectList, filters, sHolder @ ScanBuilderHolder(_, _,
+          builder: SupportsPushDownVariantExtractions))
+          if conf.getConf(org.apache.spark.sql.internal.SQLConf.PUSH_VARIANT_INTO_SCAN) &&
+             sHolder.pushedVariants.isEmpty =>
+        pushVariantExtractions(p, projectList, filters, sHolder, builder)
+    }
   }
 
   /**
@@ -418,14 +437,48 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
     for ((a, defaultValue) <- schemaAttributes.zip(defaultValues)) {
       variants.addVariantFields(a.exprId, a.dataType, defaultValue, Nil)
     }
-    if (variants.mapping.isEmpty) return originalPlan
+    if (variants.mapping.isEmpty) {
+      // No variant columns in the schema. Mark as attempted so the guard in
+      // pushDownVariants prevents a spurious second visit on the leaf node.
+      sHolder.pushedVariants = Some(new VariantInRelation())
+      return originalPlan
+    }
 
     // Collect requested fields from project list and filters
     projectList.foreach(variants.collectRequestedFields)
     filters.foreach(variants.collectRequestedFields)
 
-    // If no variant columns remain after collection, return original plan
-    if (variants.mapping.forall(_._2.isEmpty)) return originalPlan
+    // Drop a variant whose sole requested field targets VariantType (a whole-variant read: bare
+    // `SELECT v`, 2-arg `variant_get(v, '$.a')`, or a column lifted to feed a variant_get above a
+    // barrier). Shredding it to a lone variant-typed slot saves no I/O and is mishandled: the
+    // reader collapses such a slot to a boolean placeholder while catalyst keeps it VariantType, so
+    // holder.output and readSchema disagree. Leaving it raw avoids that. The predicate matches the
+    // reader's collapse condition by target type, not the exact `fullVariant` value -- the latter
+    // would miss 2-arg variant_get and non-UTC sessions (different path/tz, same broken slot).
+    // When the field coexists with real extractions the struct has >= 2 slots and is not collapsed,
+    // so this is scoped to the sole-field case. Runs before the IsNull/IsNotNull injection so
+    // presence-only references, which legitimately shred to a placeholder, are unaffected.
+    variants.mapping.values.foreach { pathToFields =>
+      pathToFields.filterInPlace { case (_, fields) =>
+        !(fields.size == 1 && fields.head._1.targetType.isInstanceOf[VariantType])
+      }
+    }
+    variants.mapping.filterInPlace { case (_, pathToFields) => pathToFields.nonEmpty }
+
+    // If a variant column is referenced only via IsNull/IsNotNull (e.g. WHERE isnotnull(v)),
+    // collectRequestedFields adds nothing to its field map (the IsNull/IsNotNull branch is a
+    // no-op). Inject a fullVariant entry for each such column so rewriteType generates the
+    // placeholder struct and the extraction is pushed on this (outer) visit. Without this,
+    // the inner leaf visit would fire with sHolder.output as its projectList, causing all
+    // variant columns -- including unreferenced siblings -- to receive fullVariant treatment.
+    val referencedAttrs = AttributeSet((projectList ++ filters).flatMap(_.references))
+    for (a <- sHolder.relation.output) {
+      if (variants.mapping.contains(a.exprId) &&
+          variants.mapping(a.exprId).values.forall(_.isEmpty) &&
+          referencedAttrs.contains(a)) {
+        variants.collectRequestedFields(a)
+      }
+    }
 
     // Build individual VariantExtraction for each field access
     // Track which extraction corresponds to which (attr, field, ordinal)
@@ -475,20 +528,31 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
       }
     }
 
-    // Call the API to push down variant extractions
-    if (extractionInfo.isEmpty) return originalPlan
+    // No extraction was requested for any variant column. Set the sentinel to suppress
+    // the inner leaf visit and avoid a spurious empty push.
+    if (extractionInfo.isEmpty) {
+      sHolder.pushedVariants = Some(new VariantInRelation())
+      return originalPlan
+    }
 
     // Companion extractions can only be honored by readers that support cast-error deferral. If
     // none were generated, the pushdown carries only non-strict accesses (`try_variant_get`, plain
     // variant reads, casts to variant/string) that are safe regardless of deferral support.
-    if (hasCompanionExtraction && !builder.supportsDeferCastError()) return originalPlan
+    if (hasCompanionExtraction && !builder.supportsDeferCastError()) {
+      // Set the sentinel like the other early returns, else the leaf is re-visited (double-visit).
+      sHolder.pushedVariants = Some(new VariantInRelation())
+      return originalPlan
+    }
 
     val extractions: Array[VariantExtraction] = extractionInfo.map(_._1).toArray
     val pushedResults = builder.pushVariantExtractions(extractions)
 
     // Filter to only the accepted extractions
     val acceptedExtractions = extractionInfo.zip(pushedResults).filter(_._2).map(_._1)
-    if (acceptedExtractions.isEmpty) return originalPlan
+    if (acceptedExtractions.isEmpty) {
+      sHolder.pushedVariants = Some(new VariantInRelation())
+      return originalPlan
+    }
 
     // Group accepted extractions by attribute to rebuild the struct schemas
     val extractionsByAttr = acceptedExtractions.groupBy(_._2)
@@ -513,6 +577,31 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
     sHolder.pushedVariantAttributeMap = attributeMap
     sHolder.output = newOutput
 
+    // Commit the required top-level columns to the builder now, while projectList and
+    // filters are in scope. This mirrors how rewriteAggregate sets holder.output to only
+    // the pushed-down aggregate columns before build() -- both patterns ensure that
+    // buildScanWithPushed* can call build() and zip holder.output with readSchema() directly.
+    //
+    // projectList/filters reference original ExprIds (this function returns originalPlan
+    // unchanged). holder.relation.output carries those original ExprIds, so AttributeSet
+    // matching works directly against it.
+    //
+    // attributeMap maps old ExprId -> new AttributeReference (with a fresh ExprId for rewritten
+    // variant columns). Invert it to map new ExprId -> old ExprId so we can filter holder.output
+    // (which carries new ExprIds) against requiredColumns (which carries old ExprIds).
+    sHolder.builder match {
+      case r: SupportsPushDownRequiredColumns =>
+        val requiredColumns = AttributeSet((projectList ++ filters).flatMap(_.references))
+        val neededRelOutput = sHolder.relation.output.filter(requiredColumns.contains)
+        val newToOldExprId = attributeMap.map { case (oldId, newAttr) => newAttr.exprId -> oldId }
+        val oldExprIdToRelAttr = sHolder.relation.output.map(a => a.exprId -> a).toMap
+        sHolder.output = sHolder.output.filter { a =>
+          oldExprIdToRelAttr.get(newToOldExprId(a.exprId)).exists(requiredColumns.contains)
+        }
+        r.pruneColumns(neededRelOutput.toStructType)
+      case _ => // builder does not support column pruning; holder.output stays full-schema
+    }
+
     // Return the original plan unchanged - transformation happens in buildScanWithPushedVariants
     originalPlan
   }
@@ -794,13 +883,17 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
 
   def buildScanWithPushedVariants(plan: LogicalPlan): LogicalPlan = plan.transform {
     case p@PhysicalOperation(projectList, filters, holder: ScanBuilderHolder)
-        if holder.pushedVariants.isDefined =>
+        if holder.pushedVariants.exists(_.mapping.nonEmpty) =>
       val variants = holder.pushedVariants.get
       val attributeMap = holder.pushedVariantAttributeMap
 
       // Build the scan
       val scan = holder.builder.build()
       val realOutput = toAttributes(scan.readSchema())
+      assert(realOutput.length == holder.output.length,
+        s"The data source returns ${realOutput.length} columns but the plan expected " +
+        s"${holder.output.length}: scan=[${realOutput.map(_.name).mkString(",")}], " +
+        s"plan=[${holder.output.map(_.name).mkString(",")}]")
       val wrappedScan = getWrappedScan(scan, holder)
       // Note: holder.pushedFilterExpressions is not propagated here because the output schema
       // changes with variant extraction. When validConstraints is wired up, this needs revisiting.