diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala index 2b291bf3a4db..42afb513f406 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala +++ b/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. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala index d6a9cfc94e8c..67bee55c2275 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala @@ -472,6 +472,17 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant super.sparkConf.set(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key, vectorizedReaderEnabled.toString) + // Locate the single DataSourceV2ScanRelation in an optimized plan, regardless of how many + // Project/Filter nodes wrap it. This keeps scan-content assertions robust against optimizer + // rules (e.g. CollapseProject) that may or may not collapse the identity outputProjection + // that buildScanWithPushedVariants inserts. + protected def findScanRelation(plan: LogicalPlan): DataSourceV2ScanRelation = { + val scans = plan.collect { case s: DataSourceV2ScanRelation => s } + assert(scans.length == 1, + s"Expected exactly one DataSourceV2ScanRelation but found ${scans.length}:\n$plan") + scans.head + } + test(s"V2 test - basic variant field extraction ($readerName)") { withTempPath { dir => val path = dir.getCanonicalPath @@ -794,6 +805,893 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant } } + test(s"V2 test - column pruning: only referenced variant column in scan ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v1 variant, v2 variant, v3 variant) using PARQUET " + + s"location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'), parse_json('{\"c\": 3}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + val query = "select variant_get(v1, '$.a', 'int') from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + sql(query).queryExecution.optimizedPlan match { + case Project(_, scanRelation: DataSourceV2ScanRelation) => + // Only v1 should be in the scan; v2 and v3 must be pruned. + assert(scanRelation.output.map(_.name) == Seq("v1"), + s"Expected scan output [v1] but got ${scanRelation.output.map(_.name)}") + assert(scanRelation.output(0).dataType.isInstanceOf[StructType], + "Expected v1 to be rewritten to struct type after extraction pushdown") + case other => + fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other") + } + } + } + } + + test(s"V2 test - column pruning: mixed variant/scalar, only variant referenced ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, s string, i int) using PARQUET location '$path'") + sql("insert into temp_v1 values (parse_json('{\"a\": 42}'), 'hello', 7)") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + val query = "select variant_get(v, '$.a', 'int') from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + sql(query).queryExecution.optimizedPlan match { + case Project(_, scanRelation: DataSourceV2ScanRelation) => + // s and i are unreferenced; only v should appear in the scan. + assert(scanRelation.output.map(_.name) == Seq("v"), + s"Expected scan output [v] but got ${scanRelation.output.map(_.name)}") + // v is referenced via variant_get, so it is rewritten to an extraction struct + // (slot per path) rather than read as a whole variant. + assert(scanRelation.output(0).dataType.isInstanceOf[StructType], + s"Expected v rewritten to struct, but got ${scanRelation.output(0).dataType}") + case other => + fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other") + } + } + } + } + + test(s"V2 test - column pruning: referenced scalar survives alongside variant ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, s string) using PARQUET location '$path'") + sql("insert into temp_v1 values (parse_json('{\"a\": 5}'), 'keep_me')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + val query = "select variant_get(v, '$.a', 'int'), s from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + sql(query).queryExecution.optimizedPlan match { + case Project(_, scanRelation: DataSourceV2ScanRelation) => + // Both v and s are referenced; neither should be pruned. + val names = scanRelation.output.map(_.name).toSet + assert(names == Set("v", "s"), + s"Expected scan output {v, s} but got $names") + // v is rewritten to an extraction struct; the scalar s keeps its original type. + val byName = scanRelation.output.map(a => a.name -> a.dataType).toMap + assert(byName("v").isInstanceOf[StructType], + s"Expected v rewritten to struct, but got ${byName("v")}") + assert(byName("s") == StringType, + s"Expected s to remain StringType, but got ${byName("s")}") + case other => + fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other") + } + } + } + } + + test(s"V2 test - column pruning: filter on second variant keeps it in scan ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v1 variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 10}'), parse_json('{\"b\": 1}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + val query = + "select variant_get(v1, '$.a', 'int') from T_V2 " + + "where variant_get(v2, '$.b', 'int') > 0" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + // v1 is in SELECT, v2 is in WHERE -- both must survive column pruning. + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + val names = scanRelation.output.map(_.name).toSet + assert(names == Set("v1", "v2"), + s"Expected scan output {v1, v2} but got $names") + // Both are referenced via variant_get, so both are rewritten to extraction structs. + scanRelation.output.foreach { a => + assert(a.dataType.isInstanceOf[StructType], + s"Expected ${a.name} rewritten to struct, but got ${a.dataType}") + } + } + } + } + + test(s"V2 test - isnotnull on variant with sibling variant column ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v1 variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values (parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + val query = "select 1 from T_V2 where isnotnull(v1)" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + // Only v1 should be in the scan; v2 is unreferenced and must be pruned. + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + assert(scanRelation.output.map(_.name) == Seq("v1"), + s"Expected scan output [v1] but got ${scanRelation.output.map(_.name)}") + assert(scanRelation.output(0).dataType.isInstanceOf[StructType], + "Expected v1 to be rewritten to a placeholder struct after pushdown") + } + } + } + + // Aggregate barrier: an aggregate over one variant column (v1) with a filter on the same + // column. The Aggregate node terminates the PhysicalOperation match, so variant extraction + // pushdown only fires on the Filter -> Project -> ScanBuilderHolder subtree below it. The + // $.price path inside max() lives above the barrier and is NOT pushed (v1 stays a full + // variant), but the unreferenced sibling v2 is never in the subtree's projectList/filters, + // so top-level column pruning still drops it from the scan. + test(s"V2 test - column pruning: aggregate barrier prunes unreferenced variant ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v1 variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"price\": 10, \"country\": \"china\"}'), parse_json('{\"b\": 2}')), " + + "(parse_json('{\"price\": 20, \"country\": \"japan\"}'), parse_json('{\"b\": 3}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + val query = + "select max(variant_get(v1, '$.price', 'long')) from T_V2 " + + "where variant_get(v1, '$.country', 'string') = 'china'" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + // v2 is unreferenced anywhere in the query -- it must be pruned out of the scan even + // though the aggregate barrier prevents the $.price extraction from being pushed. + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + assert(scanRelation.output.map(_.name) == Seq("v1"), + s"Expected scan output [v1] but got ${scanRelation.output.map(_.name)}") + } + } + } + + test(s"V2 test - column pruning: group by + variant_get aggregate, sibling pruned " + + s"($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant, name string) " + + s"using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"price\": 10, \"country\": \"china\"}'), " + + " parse_json('{\"junk\": 1}'), 'widget'), " + + "(parse_json('{\"price\": 20, \"country\": \"china\"}'), " + + " parse_json('{\"junk\": 2}'), 'widget'), " + + "(parse_json('{\"price\": 30, \"country\": \"japan\"}'), " + + " parse_json('{\"junk\": 3}'), 'gadget')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + val query = + "select name, max(variant_get(v, '$.price', 'long')) from T_V2 " + + "where variant_get(v, '$.country', 'string') = 'china' group by name" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + // v2 is unreferenced and must be pruned. v is referenced via variant_get in the WHERE + // (local, $.country) and via max(variant_get(v, '$.price')) above the aggregate barrier + // (lifted in as a bare reference -> fullVariant). Because the local $.country extraction is + // also present, fullVariant is not the column's only field, so v is still shredded (a + // multi-slot struct does not collapse to a placeholder); it is not kept raw. + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + assert(scanRelation.output.map(_.name).toSet == Set("v", "name"), + s"Expected scan output {v, name} but got ${scanRelation.output.map(_.name)}") + val vAttr = scanRelation.output.find(_.name == "v").get + assert(vAttr.dataType.isInstanceOf[StructType], + s"Expected v rewritten to struct, but got ${vAttr.dataType}") + } + } + } + + test(s"V2 test - order by variant_get: correct ordering, variant read raw, sibling pruned " + + s"($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant, name string) " + + s"using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"price\": 3}'), parse_json('{\"junk\": 1}'), 'x'), " + + "(parse_json('{\"price\": 1}'), parse_json('{\"junk\": 2}'), 'z'), " + + "(parse_json('{\"price\": 2}'), parse_json('{\"junk\": 3}'), 'y')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // The sort key variant_get(v, '$.price') lives in Sort.order, above the scan window, so v + // is lifted in only as a bare reference -> a whole-variant request -> v stays raw and the + // sort evaluates on the real variant. Compare ORDER-SENSITIVELY: a shredded full-variant + // slot would collapse to a boolean placeholder and silently mis-order, which the + // order-insensitive checkAnswer would not catch. + val query = "select name from T_V2 order by variant_get(v, '$.price', 'int')" + val expectedOrder = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect().map(_.getString(0)).toList + } + assert(expectedOrder == List("z", "y", "x"), + s"baseline sanity: expected z,y,x but got $expectedOrder") + assert(sql(query).collect().map(_.getString(0)).toList == expectedOrder, + "ORDER BY variant_get produced the wrong row order") + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + assert(scanRelation.output.map(_.name).toSet == Set("v", "name"), + s"Expected scan output {v, name} but got ${scanRelation.output.map(_.name)}") + val vAttr = scanRelation.output.find(_.name == "v").get + assert(vAttr.dataType.isInstanceOf[VariantType], + s"Expected v left as raw VariantType, but got ${vAttr.dataType}") + } + } + } + + test(s"V2 test - aggregate max(variant_get) with no local filter: no codegen crash, variant " + + s"read raw ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant, name string) " + + s"using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"price\": 10}'), parse_json('{\"junk\": 1}'), 'a'), " + + "(parse_json('{\"price\": 30}'), parse_json('{\"junk\": 2}'), 'a'), " + + "(parse_json('{\"price\": 20}'), parse_json('{\"junk\": 3}'), 'b')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // variant_get(v, '$.price') is inside an aggregate function, above the aggregate barrier, + // with no local filter/projection on v, so v is lifted in only as a bare reference -> a + // whole-variant request. A whole-variant read is kept raw: shredding it to a lone + // full-variant slot would collapse to a boolean placeholder, and max(variant_get(, + // ...)) would fail to codegen. Keeping v raw yields correct aggregates with a valid plan. + val query = + "select name, max(variant_get(v, '$.price', 'int')) as mx from T_V2 group by name" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) // a -> 30, b -> 20 + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + val vAttr = scanRelation.output.find(_.name == "v").get + assert(vAttr.dataType.isInstanceOf[VariantType], + s"Expected v left as raw VariantType, but got ${vAttr.dataType}") + assert(!scanRelation.output.exists(_.name == "v2"), + s"Expected v2 pruned but got ${scanRelation.output.map(_.name)}") + } + } + } + + test(s"V2 test - join on variant_get key: no crash, variant read raw, sibling pruned " + + s"($readerName)") { + withTempPath { dir => + val itemsPath = dir.getCanonicalPath + "/items" + val countriesPath = dir.getCanonicalPath + "/countries" + withTable("temp_items", "temp_countries") { + sql(s"create table temp_items (v variant, v2 variant, name string) " + + s"using PARQUET location '$itemsPath'") + sql("insert into temp_items values " + + "(parse_json('{\"country_code\": \"CN\"}'), parse_json('{\"junk\": 1}'), 'widget'), " + + "(parse_json('{\"country_code\": \"JP\"}'), parse_json('{\"junk\": 2}'), 'gadget')") + sql(s"create table temp_countries (code string, country_name string) " + + s"using PARQUET location '$countriesPath'") + sql("insert into temp_countries values ('CN', 'China'), ('JP', 'Japan')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(itemsPath).createOrReplaceTempView("ITEMS_V2") + spark.read.parquet(countriesPath).createOrReplaceTempView("COUNTRIES_V2") + val query = + "select i.name, c.country_name from ITEMS_V2 i join COUNTRIES_V2 c " + + "on variant_get(i.v, '$.country_code', 'string') = c.code" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + // A variant_get join key produces a valid plan (no attribute-binding failure) and + // correct results; the key is read raw and the extraction is evaluated above the scan. + checkAnswer(sql(query), expectedRows) + val scans = sql(query).queryExecution.optimizedPlan.collect { + case s: DataSourceV2ScanRelation => s + } + val itemsScan = scans.find(_.output.exists(_.name == "v")).getOrElse( + fail(s"Could not find the items scan in:\n${sql(query).queryExecution.optimizedPlan}")) + // v is the join key (referenced only via variant_get in the join condition): it must be + // left as a raw variant, not shredded. v2 is unreferenced and must be pruned. + assert(itemsScan.output.map(_.name).toSet == Set("v", "name"), + s"Expected items scan output {v, name} but got ${itemsScan.output.map(_.name)}") + val vAttr = itemsScan.output.find(_.name == "v").get + assert(vAttr.dataType.isInstanceOf[VariantType], + s"Expected v left as raw VariantType, but got ${vAttr.dataType}") + } + } + } + + test(s"V2 test - join on variant_get key with local variant filter: no crash ($readerName)") { + withTempPath { dir => + val itemsPath = dir.getCanonicalPath + "/items" + val countriesPath = dir.getCanonicalPath + "/countries" + withTable("temp_items", "temp_countries") { + sql(s"create table temp_items (v variant, v2 variant) " + + s"using PARQUET location '$itemsPath'") + sql("insert into temp_items values " + + "(parse_json('{\"country_code\": \"CN\", \"qty\": 5}'), parse_json('{\"junk\": 1}')), " + + "(parse_json('{\"country_code\": \"JP\", \"qty\": 0}'), parse_json('{\"junk\": 2}'))") + sql(s"create table temp_countries (code string, country_name string) " + + s"using PARQUET location '$countriesPath'") + sql("insert into temp_countries values ('CN', 'China'), ('JP', 'Japan')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(itemsPath).createOrReplaceTempView("ITEMS_V2") + spark.read.parquet(countriesPath).createOrReplaceTempView("COUNTRIES_V2") + // v appears in both a join condition and a local WHERE extraction. The join reference + // forces v to stay raw; the query must still optimize without crashing and be correct. + val query = + "select c.country_name from ITEMS_V2 i join COUNTRIES_V2 c " + + "on variant_get(i.v, '$.country_code', 'string') = c.code " + + "where variant_get(i.v, '$.qty', 'int') > 0" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + } + } + } + + test(s"V2 test - join on variant_get key through an aliasing projection: no crash, no " + + s"mis-shred ($readerName)") { + withTempPath { dir => + val itemsPath = dir.getCanonicalPath + "/items" + val countriesPath = dir.getCanonicalPath + "/countries" + withTable("temp_items", "temp_countries") { + sql(s"create table temp_items (v variant, v2 variant, name string) " + + s"using PARQUET location '$itemsPath'") + sql("insert into temp_items values " + + "(parse_json('{\"country_code\": \"CN\"}'), parse_json('{\"junk\": 1}'), 'widget'), " + + "(parse_json('{\"country_code\": \"JP\"}'), parse_json('{\"junk\": 2}'), 'gadget')") + sql(s"create table temp_countries (code string, country_name string) " + + s"using PARQUET location '$countriesPath'") + sql("insert into temp_countries values ('CN', 'China'), ('JP', 'Japan')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(itemsPath).createOrReplaceTempView("ITEMS_V2") + spark.read.parquet(countriesPath).createOrReplaceTempView("COUNTRIES_V2") + // The variant column is aliased (v AS vw) in a subquery before the join condition reads + // it via variant_get. Whether the optimizer inlines the alias or keeps it, the query must + // optimize without crashing, return correct results, and not over-shred (v2 pruned, the + // variant join key not turned into a struct that yields wrong data). + val query = + "select c.country_name from " + + "(select v as vw, name as nm from ITEMS_V2) i join COUNTRIES_V2 c " + + "on variant_get(i.vw, '$.country_code', 'string') = c.code" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + // Must not crash during optimization and must be correct. + checkAnswer(sql(query), expectedRows) + val scans = sql(query).queryExecution.optimizedPlan.collect { + case s: DataSourceV2ScanRelation => s + } + val itemsScan = scans.find(_.output.exists(a => a.name == "v" || a.name == "vw")) + .getOrElse(fail( + s"Could not find the items scan in:\n${sql(query).queryExecution.optimizedPlan}")) + // v2 is unreferenced and must be pruned. + assert(!itemsScan.output.exists(_.name == "v2"), + s"Expected v2 pruned but items scan was ${itemsScan.output.map(_.name)}") + // The variant join key must not be shredded: it is read as a raw variant. + val vAttr = itemsScan.output.find(a => a.name == "v" || a.name == "vw").get + assert(vAttr.dataType.isInstanceOf[VariantType], + s"Expected the variant join key left raw, but got ${vAttr.dataType}") + } + } + } + + // Non-shredded variant: parse_json writes the variant blob without shredded typed_value + // columns, so the Parquet file contains only metadata+value columns. The logical plan + // rewrite (variant_get -> GetStructField) and column pruning both still apply; the Parquet + // reader handles the non-shredded data at physical read time. + // All of the tests above already use parse_json (non-shredded) data, so this comment + // documents that the non-shredded case is fully covered by the existing test suite. + + test(s"V2 test - sole full variant: select v directly stays raw ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}')), " + + "(parse_json('{\"a\": 9}'), parse_json('{\"b\": 8}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // Selecting the whole variant is the canonical sole-fullVariant case: its only requested + // field is the entire value, which is not an extraction. v must stay raw (shredding it to a + // lone full-variant slot would collapse to a boolean placeholder and corrupt the output). + // v2 is unreferenced and must be pruned. + val query = "select v from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + assert(scanRelation.output.map(_.name) == Seq("v"), + s"Expected scan output [v] but got ${scanRelation.output.map(_.name)}") + assert(scanRelation.output(0).dataType.isInstanceOf[VariantType], + s"Expected v left raw, but got ${scanRelation.output(0).dataType}") + } + } + } + + test(s"V2 test - sole full variant on one column, sibling extracted ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}')), " + + "(parse_json('{\"a\": 9}'), parse_json('{\"b\": 8}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // The guard is per-column: v is read whole (sole fullVariant -> raw) while v2 has a real + // extraction (-> shredded). Both must coexist in the same scan. + val query = "select v, variant_get(v2, '$.b', 'int') as b from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + val vAttr = scanRelation.output.find(_.name == "v").get + val v2Attr = scanRelation.output.find(_.name == "v2").get + assert(vAttr.dataType.isInstanceOf[VariantType], + s"Expected v left raw, but got ${vAttr.dataType}") + assert(v2Attr.dataType.isInstanceOf[StructType], + s"Expected v2 shredded to struct, but got ${v2Attr.dataType}") + } + } + } + + test(s"V2 test - sole 2-arg variant_get (VariantType target) stays raw ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": {\"x\": 1}}'), parse_json('{\"b\": 2}')), " + + "(parse_json('{\"a\": {\"x\": 9}}'), parse_json('{\"b\": 8}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // 2-arg variant_get has no type argument, so its target is VariantType. Its requested + // field is NOT the exact `fullVariant` sentinel (the path is "$.a", not "$"), but it still + // shreds to a lone variant-typed slot which the reader collapses to a boolean placeholder. + // The strip must catch it by target type, not by exact-value equality, and keep v raw. + val query = "select variant_get(v, '$.a') as a from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + // v is read whole (kept raw); the unreferenced sibling v2 is pruned. + assert(scanRelation.output.map(_.name) == Seq("v"), + s"Expected scan output [v] but got ${scanRelation.output.map(_.name)}") + assert(scanRelation.output(0).dataType.isInstanceOf[VariantType], + s"Expected v left raw, but got ${scanRelation.output(0).dataType}") + } + } + } + + test(s"V2 test - sole variant_get('$$') under non-UTC session stays raw ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}')), " + + "(parse_json('{\"a\": 9}'), parse_json('{\"b\": 8}'))") + } + // A non-UTC session makes variant_get carry a non-UTC timezone, so even the "$" path is not + // equal to the UTC `fullVariant` sentinel. The strip must still keep v raw -- matching by + // target type rather than by the exact sentinel value covers this case. + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Los_Angeles") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + val query = "select variant_get(v, '$') as whole from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + assert(scanRelation.output.map(_.name) == Seq("v"), + s"Expected scan output [v] but got ${scanRelation.output.map(_.name)}") + assert(scanRelation.output(0).dataType.isInstanceOf[VariantType], + s"Expected v left raw, but got ${scanRelation.output(0).dataType}") + } + } + } + + test(s"V2 test - mixed extractions with order by a pushed path: shredded, correct order " + + s"($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"name\": \"x\", \"price\": 3}'), parse_json('{\"j\": 1}')), " + + "(parse_json('{\"name\": \"z\", \"price\": 1}'), parse_json('{\"j\": 2}')), " + + "(parse_json('{\"name\": \"y\", \"price\": 2}'), parse_json('{\"j\": 3}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // Two real extractions ($.name, $.price) are in the local Project, and the ORDER BY lifts a + // bare v (-> fullVariant) for the sort. fullVariant is not the only field, so v stays + // shredded (a multi-slot struct with a real full-variant slot, no placeholder collapse) and + // the sort reads that real slot. Compare ORDER-SENSITIVELY. + val query = "select variant_get(v, '$.name', 'string') as nm, " + + "variant_get(v, '$.price', 'long') as pr " + + "from T_V2 order by variant_get(v, '$.price', 'long')" + val expectedOrder = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect().map(r => (r.getString(0), r.getLong(1))).toList + } + assert(expectedOrder == List(("z", 1L), ("y", 2L), ("x", 3L)), + s"baseline sanity: $expectedOrder") + val actualOrder = sql(query).collect().map(r => (r.getString(0), r.getLong(1))).toList + assert(actualOrder == expectedOrder, + "ORDER BY a pushed path produced the wrong row order") + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + val vAttr = scanRelation.output.find(_.name == "v").get + assert(vAttr.dataType.isInstanceOf[StructType], + s"Expected v shredded, but got ${vAttr.dataType}") + } + } + } + + test(s"V2 test - order by a pushed path not in the select list: correct order ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"name\": \"x\", \"price\": 3}'), parse_json('{\"j\": 1}')), " + + "(parse_json('{\"name\": \"z\", \"price\": 1}'), parse_json('{\"j\": 2}')), " + + "(parse_json('{\"name\": \"y\", \"price\": 2}'), parse_json('{\"j\": 3}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // Only $.name is selected (local); the sort key $.price is lifted as a bare v (full + // variant). v shreds to {$.name, full-variant}; the sort reads the real full-variant slot. + // The ORDER-SENSITIVE check guards against a placeholder mis-order. + val query = "select variant_get(v, '$.name', 'string') as nm " + + "from T_V2 order by variant_get(v, '$.price', 'long')" + val expectedOrder = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect().map(_.getString(0)).toList + } + assert(expectedOrder == List("z", "y", "x"), s"baseline sanity: $expectedOrder") + assert(sql(query).collect().map(_.getString(0)).toList == expectedOrder, + "ORDER BY a non-selected pushed path produced the wrong row order") + } + } + } + + test(s"V2 test - group by variant_get key: optimal shred, no full-variant slot ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"k\": \"a\"}'), parse_json('{\"j\": 1}')), " + + "(parse_json('{\"k\": \"a\"}'), parse_json('{\"j\": 2}')), " + + "(parse_json('{\"k\": \"b\"}'), parse_json('{\"j\": 3}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // The grouping expression variant_get(v, '$.k') is pulled into a Project below the + // Aggregate (PullOutGroupingExpressions), so it is visible to the rewrite and shreds to + // the $.k slot -- a real extraction, no full-variant slot and no bare-v lift. v2 is pruned. + val query = + "select variant_get(v, '$.k', 'string') as k, count(1) as n " + + "from T_V2 group by variant_get(v, '$.k', 'string')" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + assert(scanRelation.output.map(_.name) == Seq("v"), + s"Expected scan output [v] but got ${scanRelation.output.map(_.name)}") + val vStruct = scanRelation.output(0).dataType match { + case s: StructType => s + case other => fail(s"Expected v shredded to struct, but got $other") + } + assert(!vStruct.fields.exists(_.dataType.isInstanceOf[VariantType]), + s"Expected optimal extraction with no full-variant slot, but got $vStruct") + } + } + } + + test(s"V2 test - distinct variant_get: optimal shred ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"k\": \"a\"}'), parse_json('{\"j\": 1}')), " + + "(parse_json('{\"k\": \"a\"}'), parse_json('{\"j\": 2}')), " + + "(parse_json('{\"k\": \"b\"}'), parse_json('{\"j\": 3}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // DISTINCT becomes an Aggregate whose grouping expression is the distinct column; like + // GROUP BY, the variant_get is pulled below the Aggregate and shreds to the $.k slot. + val query = "select distinct variant_get(v, '$.k', 'string') as k from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + assert(scanRelation.output(0).dataType.isInstanceOf[StructType], + s"Expected v shredded, but got ${scanRelation.output(0).dataType}") + } + } + } + + test(s"V2 test - window order by variant_get: no crash, shredded ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"k\": \"a\", \"price\": 3}'), parse_json('{\"j\": 1}')), " + + "(parse_json('{\"k\": \"a\", \"price\": 1}'), parse_json('{\"j\": 2}')), " + + "(parse_json('{\"k\": \"b\", \"price\": 2}'), parse_json('{\"j\": 3}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // Window is a fourth barrier type. The OVER (ORDER BY variant_get(v, '$.price')) order key + // and the selected variant_get(v, '$.k') are both pushed; the query must optimize without + // crashing and return correct results. + val query = "select variant_get(v, '$.k', 'string') as k, " + + "row_number() over (order by variant_get(v, '$.price', 'long')) as rn from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + assert(scanRelation.output.find(_.name == "v").get.dataType.isInstanceOf[StructType], + "Expected v shredded for the pushed window extractions") + } + } + } + + test(s"V2 test - two variant columns each extracted: both shredded ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}')), " + + "(parse_json('{\"a\": 9}'), parse_json('{\"b\": 8}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // Both variant columns carry real extractions on different paths; each shreds independently + // and both remain in the scan. + val query = + "select variant_get(v, '$.a', 'int') as a, variant_get(v2, '$.b', 'int') as b from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + assert(scanRelation.output.map(_.name).toSet == Set("v", "v2"), + s"Expected scan output {v, v2} but got ${scanRelation.output.map(_.name)}") + assert(scanRelation.output.forall(_.dataType.isInstanceOf[StructType]), + s"Expected both v and v2 shredded, but got " + + s"${scanRelation.output.map(a => s"${a.name}:${a.dataType.simpleString}")}") + } + } + } + + test(s"V2 test - column pruning: no variant_get, normal pruning applies ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v1 variant, v2 variant, s string) using PARQUET " + + s"location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'), 'hello')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + // No variant_get: pushdown does not fire; normal pruneColumns applies. + val query = "select s from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + sql(query).queryExecution.optimizedPlan match { + case scanRelation: DataSourceV2ScanRelation => + assert(scanRelation.output.map(_.name) == Seq("s"), + s"Expected scan output [s] but got ${scanRelation.output.map(_.name)}") + case Project(_, scanRelation: DataSourceV2ScanRelation) => + assert(scanRelation.output.map(_.name) == Seq("s"), + s"Expected scan output [s] but got ${scanRelation.output.map(_.name)}") + case other => + fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other") + } + } + } + } + + test(s"V2 test - column pruning: 2 variant cols, only second referenced ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v1 variant, v2 variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 99}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + val query = "select variant_get(v2, '$.b', 'int') from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + sql(query).queryExecution.optimizedPlan match { + case Project(_, scanRelation: DataSourceV2ScanRelation) => + assert(scanRelation.output.map(_.name) == Seq("v2"), + s"Expected scan output [v2] but got ${scanRelation.output.map(_.name)}") + assert(scanRelation.output(0).dataType.isInstanceOf[StructType], + s"Expected v2 rewritten to struct, but got ${scanRelation.output(0).dataType}") + case other => + fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other") + } + } + } + } + + test(s"V2 test - column pruning: 3 variant cols, 2 referenced ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v1 variant, v2 variant, v3 variant) using PARQUET " + + s"location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'), parse_json('{\"c\": 3}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + val query = "select variant_get(v1, '$.a', 'int'), variant_get(v3, '$.c', 'int') from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + sql(query).queryExecution.optimizedPlan match { + case Project(_, scanRelation: DataSourceV2ScanRelation) => + val names = scanRelation.output.map(_.name).toSet + assert(names == Set("v1", "v3"), + s"Expected scan output {v1, v3} but got $names") + scanRelation.output.foreach { a => + assert(a.dataType.isInstanceOf[StructType], + s"Expected ${a.name} rewritten to struct, but got ${a.dataType}") + } + case other => + fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other") + } + } + } + } + + test(s"V2 test - column pruning: 3 variant cols, all 3 referenced ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v1 variant, v2 variant, v3 variant) using PARQUET " + + s"location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'), parse_json('{\"c\": 3}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + val query = "select variant_get(v1, '$.a', 'int'), variant_get(v2, '$.b', 'int'), " + + "variant_get(v3, '$.c', 'int') from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + sql(query).queryExecution.optimizedPlan match { + case Project(_, scanRelation: DataSourceV2ScanRelation) => + val names = scanRelation.output.map(_.name).toSet + assert(names == Set("v1", "v2", "v3"), + s"Expected all 3 columns in scan but got $names") + scanRelation.output.foreach { a => + assert(a.dataType.isInstanceOf[StructType], + s"Expected ${a.name} rewritten to struct, but got ${a.dataType}") + } + case other => + fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other") + } + } + } + } + + test(s"V2 test - column pruning: 3 variant cols, none referenced via variant_get ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v1 variant, v2 variant, v3 variant, s string) " + + s"using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'), parse_json('{\"c\": 3}'), 'hi')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val df = spark.read.parquet(path) + df.createOrReplaceTempView("T_V2") + // Only s is selected; no variant_get, so extraction pushdown does not fire. + // Normal pruneColumns removes v1/v2/v3. + val query = "select s from T_V2" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + sql(query).queryExecution.optimizedPlan match { + case scanRelation: DataSourceV2ScanRelation => + assert(scanRelation.output.map(_.name) == Seq("s"), + s"Expected scan output [s] but got ${scanRelation.output.map(_.name)}") + case Project(_, scanRelation: DataSourceV2ScanRelation) => + assert(scanRelation.output.map(_.name) == Seq("s"), + s"Expected scan output [s] but got ${scanRelation.output.map(_.name)}") + case other => + fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other") + } + } + } + } + test(s"V2 No push down for JSON ($readerName)") { withTempPath { dir => val path = dir.getCanonicalPath