[SPARK-59020][SQL] Support DSv2 Parquet shredded Variant predicate pushdown

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

@@ -110,3 +110,7 @@ abstract class FileScanBuilder(
   val partitionNameSet: Set[String] =
     partitionSchema.fields.map(PartitioningUtils.getColName(_, isCaseSensitive)).toSet
 }
+
+private[v2] trait SupportsPushDownVariantPredicateFilters {
+  def pushVariantPredicateFilters(filters: Array[Filter]): Array[Filter]
+}

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

@@ -929,6 +929,8 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
       val variants = holder.pushedVariants.get
       val attributeMap = holder.pushedVariantAttributeMap
 
+      pushDownVariantPredicateFilters(filters, variants, attributeMap, holder.builder)
+
       // Build the scan
       val scan = holder.builder.build()
       val realOutput = toAttributes(scan.readSchema())
@@ -973,6 +975,26 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
       Project(rewrittenProjectList, withFilter)
   }
 
+  private def pushDownVariantPredicateFilters(
+      filters: Seq[Expression],
+      variants: VariantInRelation,
+      attributeMap: Map[ExprId, AttributeReference],
+      builder: ScanBuilder): Unit = {
+    builder match {
+      case s: SupportsPushDownVariantPredicateFilters =>
+        val pushedFilters = filters.flatMap { filter =>
+          val rewritten = variants.rewriteExpr(filter, attributeMap)
+          if (rewritten.semanticEquals(filter)) {
+            None
+          } else {
+            DataSourceStrategy.translateFilter(rewritten, supportNestedPredicatePushdown = true)
+          }
+        }
+        s.pushVariantPredicateFilters(pushedFilters.toArray)
+      case _ =>
+    }
+  }
+
   def pruneColumns(plan: LogicalPlan): LogicalPlan = plan.transform {
     case ScanOperation(project, filtersStayUp, filtersPushDown, sHolder: ScanBuilderHolder) =>
       // column pruning

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

@@ -37,7 +37,7 @@ import org.apache.spark.sql.catalyst.util.RebaseDateTime.RebaseSpec
 import org.apache.spark.sql.connector.expressions.aggregate.Aggregation
 import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader}
 import org.apache.spark.sql.execution.WholeStageCodegenExec
-import org.apache.spark.sql.execution.datasources.{AggregatePushDownUtils, DataSourceUtils, PartitionedFile, RecordReaderIterator}
+import org.apache.spark.sql.execution.datasources.{AggregatePushDownUtils, DataSourceUtils, PartitionedFile, RecordReaderIterator, VariantMetadata}
 import org.apache.spark.sql.execution.datasources.parquet._
 import org.apache.spark.sql.execution.datasources.v2._
 import org.apache.spark.sql.internal.SQLConf
@@ -55,6 +55,7 @@ import org.apache.spark.util.{SerializableConfiguration, Utils}
  * @param readDataSchema Required schema of Parquet files.
  * @param partitionSchema Schema of partitions.
  * @param filters Filters to be pushed down in the batch scan.
+ * @param variantPredicateFilters Filters on pushed-down Variant extraction fields.
  * @param aggregation Aggregation to be pushed down in the batch scan.
  * @param options The options of Parquet datasource that are set for the read.
  */
@@ -65,6 +66,7 @@ case class ParquetPartitionReaderFactory(
     readDataSchema: StructType,
     partitionSchema: StructType,
     filters: Array[Filter],
+    variantPredicateFilters: Array[Filter],
     aggregation: Option[Aggregation],
     options: ParquetOptions) extends FilePartitionReaderFactory with Logging {
   private val isCaseSensitive = sqlConf.caseSensitiveAnalysis
@@ -85,6 +87,15 @@ case class ParquetPartitionReaderFactory(
   private val pushDownInFilterThreshold = sqlConf.parquetFilterPushDownInFilterThreshold
   private val datetimeRebaseModeInRead = options.datetimeRebaseModeInRead
   private val int96RebaseModeInRead = options.int96RebaseModeInRead
+  private val variantExtractionSchema =
+    if (sqlConf.getConf(SQLConf.VARIANT_SHREDDED_PREDICATE_PUSHDOWN_ENABLED) &&
+        readDataSchema.existsRecursively(VariantMetadata.isVariantStruct)) {
+      Some(readDataSchema)
+    } else {
+      None
+    }
+  private val filtersForParquet =
+    filters ++ variantExtractionSchema.map(_ => variantPredicateFilters).getOrElse(Array.empty)
 
   private val parquetReaderCallback = new ParquetReaderCallback()
 
@@ -245,8 +256,9 @@ case class ParquetPartitionReaderFactory(
           pushDownStringPredicate,
           pushDownInFilterThreshold,
           isCaseSensitive,
-          datetimeRebaseSpec)
-        filters
+          datetimeRebaseSpec,
+          variantExtractionSchema = variantExtractionSchema)
+        filtersForParquet
           // Collects all converted Parquet filter predicates. Notice that not all predicates can be
           // converted (`ParquetFilters.createFilter` returns an `Option`). That's why a `flatMap`
           // is used here.

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

@@ -48,7 +48,8 @@ case class ParquetScan(
     pushedAggregate: Option[Aggregation] = None,
     partitionFilters: Seq[Expression] = Seq.empty,
     dataFilters: Seq[Expression] = Seq.empty,
-    pushedVariantExtractions: Array[VariantExtraction] = Array.empty) extends FileScan {
+    pushedVariantExtractions: Array[VariantExtraction] = Array.empty,
+    pushedVariantPredicateFilters: Array[Filter] = Array.empty) extends FileScan {
   override def isSplitable(path: Path): Boolean = {
     // If aggregate is pushed down, only the file footer will be read once,
     // so file should not be split across multiple tasks.
@@ -176,6 +177,7 @@ case class ParquetScan(
       effectiveSchema,
       readPartitionSchema,
       pushedFilters,
+      pushedVariantPredicateFilters,
       pushedAggregate,
       new ParquetOptions(options.asCaseSensitiveMap.asScala.toMap, conf))
   }
@@ -190,9 +192,11 @@ case class ParquetScan(
       val pushedVariantEqual =
         java.util.Arrays.equals(pushedVariantExtractions.asInstanceOf[Array[Object]],
           p.pushedVariantExtractions.asInstanceOf[Array[Object]])
+      val pushedVariantPredicateFiltersEqual =
+        equivalentFilters(pushedVariantPredicateFilters, p.pushedVariantPredicateFilters)
       super.equals(p) && dataSchema == p.dataSchema && options == p.options &&
         equivalentFilters(pushedFilters, p.pushedFilters) && pushedDownAggEqual &&
-        pushedVariantEqual
+        pushedVariantEqual && pushedVariantPredicateFiltersEqual
     case _ => false
   }
 
@@ -214,9 +218,17 @@ case class ParquetScan(
     } else {
       "[]"
     }
+    val variantPredicateFilterMetadata =
+      if (pushedVariantPredicateFilters.nonEmpty) {
+        Map("PushedVariantPredicateFilters" ->
+          seqToString(pushedVariantPredicateFilters.toImmutableArraySeq))
+      } else {
+        Map.empty[String, String]
+      }
     super.getMetaData() ++ Map("PushedFilters" -> seqToString(pushedFilters.toImmutableArraySeq)) ++
       Map("PushedAggregation" -> pushedAggregationsStr) ++
       Map("PushedGroupBy" -> pushedGroupByStr) ++
-      Map("PushedVariantExtractions" -> variantExtractionStr)
+      Map("PushedVariantExtractions" -> variantExtractionStr) ++
+      variantPredicateFilterMetadata
   }
 }

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

@@ -25,8 +25,10 @@ import org.apache.spark.sql.connector.expressions.aggregate.Aggregation
 import org.apache.spark.sql.connector.read.{SupportsPushDownAggregates, SupportsPushDownVariantExtractions, VariantExtraction}
 import org.apache.spark.sql.execution.datasources.{AggregatePushDownUtils, PartitioningAwareFileIndex}
 import org.apache.spark.sql.execution.datasources.parquet.{ParquetFilters, SparkToParquetSchemaConverter}
-import org.apache.spark.sql.execution.datasources.v2.FileScanBuilder
-import org.apache.spark.sql.internal.LegacyBehaviorPolicy
+import org.apache.spark.sql.execution.datasources.v2.{
+  FileScanBuilder,
+  SupportsPushDownVariantPredicateFilters}
+import org.apache.spark.sql.internal.{LegacyBehaviorPolicy, SQLConf}
 import org.apache.spark.sql.sources.Filter
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
@@ -40,7 +42,8 @@ case class ParquetScanBuilder(
     options: CaseInsensitiveStringMap)
   extends FileScanBuilder(sparkSession, fileIndex, dataSchema)
     with SupportsPushDownAggregates
-    with SupportsPushDownVariantExtractions {
+    with SupportsPushDownVariantExtractions
+    with SupportsPushDownVariantPredicateFilters {
   lazy val hadoopConf = {
     val caseSensitiveMap = options.asCaseSensitiveMap.asScala.toMap
     // Hadoop Configurations are case sensitive.
@@ -53,6 +56,8 @@ case class ParquetScanBuilder(
 
   private var pushedVariantExtractions = Array.empty[VariantExtraction]
 
+  private var pushedVariantPredicateFilters = Array.empty[Filter]
+
   override protected val supportsNestedSchemaPruning: Boolean = true
 
   override def pushDataFilters(dataFilters: Array[Filter]): Array[Filter] = {
@@ -66,15 +71,6 @@ case class ParquetScanBuilder(
       val isCaseSensitive = sqlConf.caseSensitiveAnalysis
       val parquetSchema =
         new SparkToParquetSchemaConverter(sparkSession.sessionState.conf).convert(readDataSchema())
-      // Shredded-variant predicate pushdown (SPARK-55817) is not wired here: it applies to the
-      // DSv1 path only. DSv2 does rewrite variant extractions into `v.`0`` struct accesses, but
-      // only in `V2ScanRelationPushDown.buildScanWithPushedVariants`, which runs *after*
-      // `pushDownFilters`. So the filters reaching this method are still `variant_get(v, ...)`
-      // predicates, which do not translate to a source `Filter` at all -- there is no
-      // shredded-variant logical name for ParquetFilters to resolve here, and nothing would be
-      // reported convertible even with a variantExtractionSchema. DSv2 reads remain correct (the
-      // variant filter is applied post-scan); they just do not get row-group skipping on shredded
-      // columns.
       val parquetFilters = new ParquetFilters(
         parquetSchema,
         pushDownDate,
@@ -121,6 +117,18 @@ case class ParquetScanBuilder(
     Array.fill(extractions.length)(true)
   }
 
+  override def pushVariantPredicateFilters(filters: Array[Filter]): Array[Filter] = {
+    val sqlConf = sparkSession.sessionState.conf
+    pushedVariantPredicateFilters =
+      if (sqlConf.parquetFilterPushDown &&
+          sqlConf.getConf(SQLConf.VARIANT_SHREDDED_PREDICATE_PUSHDOWN_ENABLED)) {
+        filters
+      } else {
+        Array.empty[Filter]
+      }
+    pushedVariantPredicateFilters
+  }
+
   override def build(): ParquetScan = {
     // the `finalSchema` is either pruned in pushAggregation (if aggregates are
     // pushed down), or pruned in readDataSchema() (in regular column pruning). These
@@ -130,6 +138,6 @@ case class ParquetScanBuilder(
     }
     ParquetScan(sparkSession, hadoopConf, fileIndex, dataSchema, finalSchema,
       readPartitionSchema(), pushedDataFilters, options, pushedAggregations,
-      partitionFilters, dataFilters, pushedVariantExtractions)
+      partitionFilters, dataFilters, pushedVariantExtractions, pushedVariantPredicateFilters)
   }
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantShreddingFilterPushdownSuite.scala

@@ -36,12 +36,11 @@ import org.apache.spark.util.AccumulatorContext
  * the physical leaf and guards it so a row group is skipped only when the leaf cannot match AND
  * every value for the path is provably in the leaf (see `makeShreddedFilter`).
  *
- * Scope: the optimization fires on the DSv1 read path only. DSv2 does rewrite variant extractions
- * into `v.`0`` struct accesses, but only after filter pushdown has run, so the filters offered to
- * the Parquet scan builder are still `variant_get(v, ...)` predicates and cannot be pushed for
- * row-group skipping (see the comment in ParquetScanBuilder). DSv2 reads remain correct -- the
- * variant filter is applied post-scan -- they just do not skip row groups. These tests therefore
- * assert skipping only on DSv1, and assert correctness on both DSv1 and DSv2.
+ * Scope: the optimization fires on both DSv1 and DSv2 read paths. DSv2 performs the regular filter
+ * pushdown before variant extraction pushdown, so after variant extraction rewrites predicates into
+ * `v.`0`` struct accesses it runs a second Parquet-only predicate pushdown for those rewritten
+ * variant predicates. These tests assert skipping when the vectorized reader exposes the row-group
+ * count, and assert correctness on both vectorized and non-vectorized readers.
  *
  * The central correctness concern is soundness under fallback: shredding is per-row and per-file
  * best-effort, so values that don't fit the shredded type (overflow / type mismatch) or that are
@@ -102,7 +101,6 @@ class VariantShreddingFilterPushdownSuite extends QueryTest with ParquetTest
   }
 
   // Run `block` with pushdown enabled, across the {DSv1, DSv2} x {vectorized, non-vectorized} grid.
-  // `dsv1` is passed so a test can assert row-group skipping only on the DSv1 path.
   private def forEachReader(block: (Boolean, Boolean) => Unit): Unit = {
     Seq("parquet" -> true, "" -> false).foreach { case (useV1, dsv1) =>
       Seq(true, false).foreach { vectorized =>
@@ -145,11 +143,11 @@ class VariantShreddingFilterPushdownSuite extends QueryTest with ParquetTest
       val expected = baseline(read)
       assert(expected == Seq(Row(1500L)), s"baseline should return the fallback row, got $expected")
 
-      forEachReader { (dsv1, vectorized) =>
+      forEachReader { (_, vectorized) =>
         // The row group's only match is in the residual with a NULL leaf, so the guard must keep
         // it: results include the fallback row and the row group is not skipped.
         checkAnswer(read, expected)
-        if (dsv1 && vectorized) {
+        if (vectorized) {
           val all = spark.read.parquet(dir.getAbsolutePath)
             .selectExpr("try_variant_get(v, '$.a', 'bigint') AS a")
           assert(countRowGroupsRead(read) == countRowGroupsRead(all),
@@ -233,22 +231,22 @@ class VariantShreddingFilterPushdownSuite extends QueryTest with ParquetTest
     }
   }
 
-  test("residual-null happy path: a row group is skipped (DSv1) and results are correct") {
+  test("residual-null happy path: a row group is skipped and results are correct") {
     withTempDir { dir =>
       // Homogeneous typed data across two row groups. All values shred cleanly (residuals all
       // NULL), so the optimization fires and one row group is skipped on DSv1.
       val jsonExpr = "'{\"a\":' || id || '}'"
       // Small block size -> at least two row groups: [0,999] and [1000,1999].
       writeShredded(dir, "a bigint", jsonExpr, numRows = 2000, blockSize = 512)
 
-      forEachReader { (dsv1, vectorized) =>
+      forEachReader { (_, vectorized) =>
         val filtered = spark.read.parquet(dir.getAbsolutePath)
           .selectExpr("variant_get(v, '$.a', 'bigint') AS a")
           .where("a > 999")
         val all = spark.read.parquet(dir.getAbsolutePath)
           .selectExpr("variant_get(v, '$.a', 'bigint') AS a")
         checkAnswer(filtered.orderBy("a"), (1000L to 1999L).map(Row(_)))
-        if (dsv1 && vectorized) {
+        if (vectorized) {
           assert(countRowGroupsRead(filtered) < countRowGroupsRead(all),
             "Expected at least one row group to be skipped by the shredded leaf statistics")
         }
@@ -266,35 +264,35 @@ class VariantShreddingFilterPushdownSuite extends QueryTest with ParquetTest
       val jsonExpr = "'{\"a\":' || id || ', \"z\":\"outside\"}'"
       writeShredded(dir, "a bigint", jsonExpr, numRows = 2000, blockSize = 512)
 
-      forEachReader { (dsv1, vectorized) =>
+      forEachReader { (_, vectorized) =>
         val filtered = spark.read.parquet(dir.getAbsolutePath)
           .selectExpr("variant_get(v, '$.a', 'bigint') AS a").where("a > 999")
         val all = spark.read.parquet(dir.getAbsolutePath)
           .selectExpr("variant_get(v, '$.a', 'bigint') AS a")
         checkAnswer(filtered.orderBy("a"), (1000L to 1999L).map(Row(_)))
-        if (dsv1 && vectorized) {
+        if (vectorized) {
           assert(countRowGroupsRead(filtered) < countRowGroupsRead(all),
             "Expected skipping despite a non-null top-level residual (partial object)")
         }
       }
     }
   }
 
-  test("multi-level $.a.b: skip fires on DSv1 and results are correct") {
+  test("multi-level $.a.b: skip fires and results are correct") {
     withTempDir { dir =>
       // `a` shredded as struct<b bigint>. Homogeneous nested typed data across two row groups so
       // the skip fires on the nested leaf `v.typed_value.a.typed_value.b.typed_value`.
       val typedJson = "'{\"a\":{\"b\":' || id || '}}'"
       writeShredded(dir, "a struct<b bigint>", typedJson, numRows = 2000, blockSize = 512)
 
-      forEachReader { (dsv1, vectorized) =>
+      forEachReader { (_, vectorized) =>
         val filtered = spark.read.parquet(dir.getAbsolutePath)
           .selectExpr("variant_get(v, '$.a.b', 'bigint') AS b")
           .where("b > 999")
         val all = spark.read.parquet(dir.getAbsolutePath)
           .selectExpr("variant_get(v, '$.a.b', 'bigint') AS b")
         checkAnswer(filtered.orderBy("b"), (1000L to 1999L).map(Row(_)))
-        if (dsv1 && vectorized) {
+        if (vectorized) {
           assert(countRowGroupsRead(filtered) < countRowGroupsRead(all),
             "Expected a row group to be skipped by the nested shredded leaf statistics")
         }
@@ -319,9 +317,9 @@ class VariantShreddingFilterPushdownSuite extends QueryTest with ParquetTest
         .where("b > 999")
       assert(baseline(read) == Seq(Row(1500L)), "baseline should return the nested fallback row")
 
-      forEachReader { (dsv1, vectorized) =>
+      forEachReader { (_, vectorized) =>
         checkAnswer(read, Seq(Row(1500L)))
-        if (dsv1 && vectorized) {
+        if (vectorized) {
           val all = spark.read.parquet(dir.getAbsolutePath)
             .selectExpr("try_variant_get(v, '$.a.b', 'bigint') AS b")
           assert(countRowGroupsRead(read) == countRowGroupsRead(all),
@@ -338,13 +336,13 @@ class VariantShreddingFilterPushdownSuite extends QueryTest with ParquetTest
     withTempDir { dir =>
       writeShredded(dir, "a bigint", "'{\"a\":' || id || '}'", numRows = 2000, blockSize = 512,
         annotate = false)
-      forEachReader { (dsv1, vectorized) =>
+      forEachReader { (_, vectorized) =>
         val filtered = spark.read.parquet(dir.getAbsolutePath)
           .selectExpr("variant_get(v, '$.a', 'bigint') AS a").where("a > 999")
         val all = spark.read.parquet(dir.getAbsolutePath)
           .selectExpr("variant_get(v, '$.a', 'bigint') AS a")
         checkAnswer(filtered.orderBy("a"), (1000L to 1999L).map(Row(_)))
-        if (dsv1 && vectorized) {
+        if (vectorized) {
           assert(countRowGroupsRead(filtered) < countRowGroupsRead(all),
             "Expected a row group to be skipped on the unannotated layout")
         }