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

[qlong]

What changes were proposed in this pull request?

Three fixes in pushVariantExtractions (called byV2ScaqqnRelationPushDown.pushDownVariants):

1. Guard against double-visit: Add a pushedVariants.isEmpty
   sentinel check so the inner ScanBuilderHolder leaf visit (caused by
   transformDown recursing into the child after returning the plan
   unchanged) returns immediately. This ensures
   builder.pushVariantExtractions is called exactly once per holder.
2. Eager column pruning: While projectList and filters are in
   scope, call builder.pruneColumns(requiredSchema) for builders
   implementing SupportsPushDownRequiredColumns and trim
   sHolder.output to the required columns. By the time
   buildScanWithPushedVariants calls build(), the builder already
   has the correct pruned schema. This is similiar to how
   buildScanWithPushedAggregate works.
3. Keep whole-variant reads raw: pushdown is for extractions, not
   whole-variant reads. A bare variant reference -- SELECT v, or a
   column lifted to feed a variant_get above a Join/Sort/Aggregate
   barrier the local rewrite cannot see -- is recorded as fullVariant
   (path $), meaning "the entire value." Shredding that to a lone
   full-variant slot saves no I/O and is mishandled: the Parquet reader
   collapses it to a boolean placeholder, and above a barrier the
   re-exposed GetStructField AS v#orig alias is dropped by
   RemoveRedundantAliases, giving wrong results or an invalid plan. So
   when fullVariant is a variant's only requested field, leave the
   column raw; when it coexists with real extractions (e.g. SELECT v, variant_get(v, '$.a')) the >=2-slot struct is not collapsed and
   keeps its pushdown. This subsumes the join-key case and shreds
   automatically once barrier-aware pushdown makes the variant_get
   visible as a typed path.

Jira: https://issues.apache.org/jira/browse/SPARK-57499

Why are the changes needed?

Three bugs on the accepted variant pushdown path:

Issue 1 (Performance) — column pruning is skipped. buildScanWithPushedVariants calls
builder.build() and replaces the ScanBuilderHolder with a DataSourceV2ScanRelation.
The subsequent pruneColumns rule matches only ScanBuilderHolder nodes, so it is a
no-op and builder.pruneColumns() is never called. The scan reads the full table schema
including unreferenced columns. For unreferenced VARIANT columns this is especially
costly — each is fully reconstructed from its shredded Parquet tree on every row.

Issue 2 (Correctness) — invalid plan / crash on tables with >=2 VARIANT columns
pushDownVariants uses transformDown, which recurses into the child ScanBuilderHolder after returning the plan unchanged. The bare ScanBuilderHolder matches PhysicalOperation a second time, collecting an unreferenced sibling VARIANT column as a full-variant request and pushing it to the builder again. ParquetScanBuilder overwrites its state on every call, so the second push clobbers the correct extraction from the first. The rewritten scan then emits a fresh ExprId for the variant while the projection still references the original, and binding fails.

This affects any extraction shape — projection, ORDER BY, aggregate, join — on a table with two or more VARIANT columns. Single-VARIANT tables are unaffected.

Reproduce on stock spark-4.1.x (path-based views force DSv2):

SET spark.sql.sources.useV1SourceList = "";

CREATE TABLE t (a INT, v1 VARIANT, v2 VARIANT) USING PARQUET LOCATION '/tmp/vt';
INSERT INTO t VALUES
  (1, parse_json('{"x":1,"price":3,"name":"x"}'), parse_json('{"y":2}')),
  (2, parse_json('{"x":9,"price":1,"name":"z"}'), parse_json('{"y":8}'));
CREATE OR REPLACE TEMPORARY VIEW tv    USING parquet OPTIONS (path '/tmp/vt');
CREATE OR REPLACE TEMPORARY VIEW codes USING parquet OPTIONS (path '/tmp/vt');

-- All four crash with: [INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND] Could not find v1#NN in [...]
SELECT variant_get(v1, '$.x', 'int') FROM tv;        
SELECT variant_get(v1,'$.name','string') AS nm
  FROM tv ORDER BY variant_get(v1,'$.price','int');                   
SELECT max(variant_get(v1, '$.price', 'int')) FROM tv;                     
SELECT l.a FROM tv l JOIN codes r
  ON variant_get(l.v1,'$.x','int') = variant_get(r.v1,'$.x','int');  

[INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND] Could not find v1#21 in [a#33,v1#34,v2#35]. SQLSTATE: XX000
  at org.apache.spark.sql.catalyst.expressions.BindReferences$.attributeNotFoundException(BoundAttribute.scala:109)
  ...

Issue 3 (Correctness) -- wrong results / crash when a whole-variant read is shredded. A bare
variant reference (a plain SELECT v, or a column lifted to feed a variant_get
above a Join/Sort/Aggregate barrier the local rewrite cannot see) is recorded as a
full-variant request (path "$"). Shredding it to a lone full-variant slot is both
useless (the whole value is read regardless) and mishandled:

• The Parquet reader collapses a lone VariantType slot to a boolean placeholder,
  so ORDER BY variant_get(v, '$.price') sorts on the placeholder and silently
  returns the wrong order, and max(variant_get(v, '$.price')) fails to codegen.
• A join key is re-exposed above the join as GetStructField(v_new, i) AS v#orig;
  RemoveRedundantAliases collapses the alias and the condition references a
  dropped ExprId, failing plan validation.

This issue is masked on stock 4.1 by the Issue 2 binding crash, which fails first; it surfaces only after that issue is fixed, which is why it has no repro for stock 4.1. See new unit tests for join, order by, aggregrate.

Does this PR introduce any user-facing change?

No

How was this patch tested?

• Added new unit tests
• manual testing with spark-sql

Was this patch authored or co-authored using generative AI tooling?

Co-authored with Claude code (Sonnet 4.6)

[cloud-fan]

can you fix merge conflicts?

[qlong]

@cloud-fan rebased. I will keep an eye on the CI build

[cloud-fan]

1 blocking, 2 non-blocking, 1 nit.
Solid, well-tested fix; rule structure is correct and consistent with the aggregate path. One correctness gap: the whole-variant guard is narrower than the Parquet reader's placeholder-collapse condition.

Correctness (2)

• V2ScanRelationPushDown.scala:476: whole-variant strip matches only the exact fullVariant sentinel, so 2-arg variant_get (VariantType target) escapes it and the reader collapses it to a boolean placeholder → wrong results — see inline
• V2ScanRelationPushDown.scala:373: one early return (!supportsDeferCastError) omits the sentinel this comment promises — see inline

Suggestions (1)

• PushVariantIntoScanSuite.scala: no 2-arg variant_get (VariantType target) coverage — a case asserting the column is kept raw (e.g. variant_get(v, '$.a'), and variant_get(v, '$') under a non-UTC session) would catch and guard the blocking finding above.

Nits: 1 minor item (see inline comments).

Verification

Traced the variant-extraction rewrite for equivalence. Equivalent for: bare whole reads (collected as the UTC fullVariant → stripped → kept raw, timezone-robust), non-variant-target extractions, IsNull/IsNotNull-only columns (placeholder injection), and NULL / nondeterminism / ordering (variant_get is deterministic, evaluated once over the pushed slot). The one non-equivalent cell: a sole VariantType-target extraction that isn't the exact fullVariant — variant_get(v,'$.a') (2-arg) or variant_get(v,'$') non-UTC — which the reader collapses to Boolean while catalyst keeps VariantType; the top-level-count assert doesn't catch the inner-type mismatch. Output-attribute metadata (nullability/exprId) is consistent — the rewrite synthesizes a Project over the scan's own readSchema and preserves the original column nullability; no outer-join-style widening.

[cloud-fan]

pushVariantExtractions strips a column's mapping only when its sole requested field equals the exact fullVariant sentinel (VariantMetadata("$", failOnError=true, "UTC")). But a 2-arg variant_get builds VariantGet(v, path, VariantType, …) carrying the session timezone and the real path, so variant_get(v, '$.a') (any session) and variant_get(v, '$') (non-UTC session) don't match and escape the strip. They then shred to a lone VariantType slot, which ParquetScan.rewriteVariantPushdownSchema collapses to a BooleanType placeholder — while catalyst's VariantInRelation.rewriteType keeps it VariantType. So holder.output declares struct{VariantType} but the scan produces struct{Boolean}; the new top-level-count assert (1==1) doesn't catch the inner-type mismatch, and GetStructField(0) reads the placeholder → wrong results / invalid plan — Issue 3 via the function form. Aligning the guard with the reader's collapse predicate fixes all of these and subsumes the bare-read case:

Suggested change

	!(fields.size == 1 && fields.contains(RequestedVariantField.fullVariant))
	!(fields.size == 1 && fields.head._1.targetType.isInstanceOf[VariantType])

Worth adding a 2-arg variant_get test — none exists today, so CI is blind to this.

[qlong]

Fixed, targetType.isInstanceOf[VariantType]) is more direct, added tests for 2-arg.

[cloud-fan]

This comment says the sentinel is written by the outer visit on completion "successfully or not", but the early return at the hasCompanionExtraction && !builder.supportsDeferCastError() check below (~L554) returns originalPlan without setting sHolder.pushedVariants. That leaves the guard's isEmpty true, so the inner leaf re-visit (Issue 2) would fire. It's masked today — ParquetScanBuilder.supportsDeferCastError() returns true, so the branch is unreachable for the only OSS implementer, and the new full-variant strip would also neutralize a re-visit — but it's a latent reintroduction of the bug for any future builder returning false. Suggest setting sHolder.pushedVariants = Some(new VariantInRelation()) before that return, consistent with the other early returns.

[qlong]

Thanks for review. Changed to set sHolder.pushedVariants = Some(new VariantInRelation()) before retturn.

[cloud-fan]

Nit: this comment (and a couple of others, e.g. the "SPARK-57499 regression" note) is framed as PR history — "Before the fix this shredded…". Once merged that history is invisible to readers. Consider rephrasing in terms of the invariant the test guards now, e.g. "A whole-variant read is kept raw; shredding it to a lone full-variant slot would collapse to a boolean placeholder and corrupt the result."

[qlong]

fixed

[cloud-fan]

thanks, merging to master/4.x

[qlong]

Thanks for review and merge.