Add GPU support for try_variant_get [databricks]

[nartal1]

Contributes to #15220 and #15221.

Description

This PR adds GPU support for Spark 4.x Variant extraction from Parquet reads.

It accelerates non-strict try_variant_get when:

• The path is a literal object-field path such as $.field or $.nested.field.
• The target type is tinyint, smallint, int, bigint, or string.

The implementation uses the cuDF Java VariantUtils.getVariantFieldValue and VariantUtils.castVariantValue APIs added in rapidsai/cudf#23069. Exact string and signed-integer Variant values are decoded on GPU. Batches containing values that require Spark coercion semantics, such as string-to-integer conversion, are evaluated through the CPU bridge to preserve Spark-compatible results.

Strict variant_get, non-literal paths, array paths, quoted-key paths, and unsupported target types continue to fall back to CPU. Variant writes remain unsupported on GPU.

Changes

• Add GPU expression support for try_variant_get across Spark 4.0+ shims, including Databricks 17.3.
• Support tinyint, smallint, int, bigint, and string targets.
• Support literal top-level and nested object-field paths.
• Add Spark-compatible integer widening, range checks, and overflow-to-null behavior.
• Use the CPU bridge for runtime values requiring Spark coercion.
• Add Variant handling to Parquet schema checks and GPU/host column conversion.
• Support both LIST<UINT8> and STRING Variant children during host conversion.
• Update generated Spark 4.x support metadata.
• Add integration and Scala unit coverage for supported extraction, nulls, missing fields, boundaries, heterogeneous values, downstream operators, and fallback paths.

Configuration

Variant extraction is currently registered as an incompatible operation and requires:

spark.rapids.sql.incompatibleOps.enabled=true

The CPU bridge used for Spark-compatible runtime coercions is enabled by default.

Testing

• Spark 4.0/Scala 2.13 distribution build: passed.
• Spark 4.0 Variant integration tests in variant_test.py: 31 passed.
• Databricks 17.3 (400db173) build and focused tests: passed.
• CPU and GPU results matched for supported extraction, null and missing values, integer boundaries, heterogeneous values, filters, aggregates, and shuffle paths.

Performance

Tested on a Tesla T4 with OSS Spark 4.0.0 using 10 million rows, 128 Parquet partitions, local[4], and 6 GB driver memory. Each result is the average of three measured runs after one warm-up run.

Projection	CPU	GPU	Speedup
$.id as BIGINT	4.694 s	2.035 s	2.31x
$.name as STRING	5.066 s	1.670 s	3.03x
Two nested integer fields	9.811 s	1.626 s	6.03x
Mixed string, integer, and bigint projection	11.464 s	2.057 s	5.57x

Projection queries used Spark's noop sink.

Benchmark queries

SELECT try_variant_get(v, '$.id', 'bigint') AS id
FROM variant_data;

SELECT try_variant_get(v, '$.name', 'string') AS name
FROM variant_data;

SELECT
  try_variant_get(v, '$.observation.value.temperature', 'int') AS temperature,
  try_variant_get(v, '$.observation.value.humidity', 'int') AS humidity
FROM variant_data;

SELECT
  try_variant_get(v, '$.name', 'string') AS name,
  try_variant_get(v, '$.observation.value.temperature', 'int') AS temperature,
  try_variant_get(v, '$.id', 'bigint') AS id
FROM variant_data;

A nested integer filter followed by COUNT(*) achieved a 1.51x speedup. Direct SUM(try_variant_get(...)) did not improve performance because its partial aggregate remains on CPU when the input schema contains the raw VariantType; the Variant extraction and final aggregate still run on GPU.

Checklists

Documentation

• Updated for new or modified user-facing features or behaviors
• No user-facing change

Testing

• Added or modified tests to cover new code paths
• Covered by existing tests
• Not required

Performance

• Tests ran and results are added in the PR description
• Issue filed with a link in the PR description
• Not required

[nartal1]

build

[greptile-apps]

Greptile Summary

This PR adds GPU acceleration for Spark 4.x try_variant_get, supporting tinyint, smallint, int, bigint, and string targets on literal object-field paths. It introduces GpuVariantGet with a CPU bridge fallback for batches requiring Spark coercion semantics, and wires Variant type support through the Parquet schema checker, type-check framework, batch utilities, and host column vector layer.

• GpuVariantGet decodes Variant children directly via VariantUtils.getVariantFieldValue/castVariantValue; integer widening coalesces INT8→INT64 decoded results, then narrows with out-of-range nullification; rows not fully covered by GPU decoders are evaluated through GpuCpuBridgeExpression.
• Type infrastructure adds TypeEnum.VARIANT / TypeSig.VARIANT, splits additionalParquetSupportedTypes into read-only and write-only variants so VariantType is registered for Parquet reads but not writes.
• Parquet schema plumbing adds isVariantPhysicalType to validate the two-BINARY physical layout and bypasses primitive-compat checks for Variant columns.

Confidence Score: 5/5

Safe to merge; GPU/CPU results are validated by 31 integration tests and the incompat flag gates rollout.

The core extraction logic is well-tested across boundary values, null rows, missing fields, heterogeneous batches, and every declared fallback path. Resource management follows withResource/safeMap throughout. The two findings are minor defensive gaps that do not affect correctness or runtime behaviour of supported paths.

Files Needing Attention: The mixed-type child routing in GpuVariantGet.withCudfVariantView and the shim annotation in GpuColumnVectorVariantSuite are the only spots worth a second look before merge.

Important Files Changed

Filename	Overview
sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/GpuVariantGet.scala	New GPU implementation of try_variant_get. Resource management is correct using withResource/safeMap/try-finally. One edge case: mixed STRING/LIST Variant children in withCudfVariantView produces a misleading error.
sql-plugin/src/main/java/com/nvidia/spark/rapids/GpuColumnVector.java	Adds Variant type plumbing: isVariantType predicate, variantHostType schema builders, typeConversionAllowed for Variant struct. isVariantBinaryChild correctly handles both STRING and LIST children.
sql-plugin/src/main/java/com/nvidia/spark/rapids/RapidsHostColumnVectorCore.java	getBinary gains a fast path for DType.STRING for host-side Variant child handling; getChild properly dispatches Variant type to BinaryType per child.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/parquet/GpuParquetScan.scala	Adds Variant schema handling to convertToParquetNative and checkSchemaCompat. isVariantPhysicalType enforces two REQUIRED BINARY children.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/TypeChecks.scala	Adds TypeEnum.VARIANT and TypeSig.VARIANT; updates CastChecks sparkXxxSig signatures; conditionally filters VARIANT from docs/tools output on pre-4.x shims.
sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/GpuTypeShims.scala	New shim file for 400+ adding VARIANT to additionalParquetReadSupportedTypes and additionalCommonOperatorSupportedTypes; supportsVariantType returns true here and false in spark330 shim.
integration_tests/src/main/python/variant_test.py	31 integration tests using assert_gpu_and_cpu_are_equal_collect and assert_gpu_fallback_collect; covers all supported types, boundaries, nulls, heterogeneous batches, filter/aggregate/shuffle, and every declared fallback.
sql-plugin/src/test/spark400/scala/com/nvidia/spark/rapids/GpuColumnVectorVariantSuite.scala	Unit tests for typeConversionAllowed on Variant struct columns with proper withResource usage. Missing shim versions 404 and 413 from the annotation block.

Sequence Diagram

sequenceDiagram
    participant Exec as ProjectExec (GPU)
    participant Meta as GpuVariantGetMeta
    participant Expr as GpuVariantGet.doColumnar
    participant cuDF as VariantUtils (cuDF)
    participant Bridge as GpuCpuBridgeExpression

    Meta->>Meta: tagExprForGpu()
    Meta->>Expr: convertToGpu(lhs, rhs)
    Exec->>Expr: doColumnar(input: GpuColumnVector)
    Expr->>Expr: "getChildColumnView(0=value, 1=metadata)"
    Expr->>Expr: withCudfVariantView
    Expr->>cuDF: getVariantFieldValue(cudfVariant, path)
    cuDF-->>Expr: rawValue ColumnView
    alt StringType target
        Expr->>cuDF: castVariantValue(rawValue, STRING)
        Expr->>Expr: allRowsCovered?
        alt covered
            Expr-->>Exec: GPU column
        else needs coercion
            Expr->>Bridge: columnarEval
            Bridge-->>Exec: CPU column
        end
    else Integer targets
        Expr->>cuDF: castVariantValue x4 INT8/16/32/64
        Expr->>Expr: allRowsCovered?
        alt covered
            Expr->>Expr: normalizeIntegers + narrowInteger
            Expr-->>Exec: GPU column
        else needs coercion
            Expr->>Bridge: columnarEval
            Bridge-->>Exec: CPU column
        end
    end

Loading

_{Reviews (8): Last reviewed commit: "Merge branch 'main' into variant-extract..." | Re-trigger Greptile}

[greptile-apps]

Availability check validates unused methods

isVariantCudfAvailable verifies the existence of getVariantFieldValue and castVariantValue, but neither method is called anywhere in this file — only extractVariantField is actually invoked at runtime. If a future cuDF release renames or removes those two helper methods while keeping extractVariantField, the entire feature silently falls back to CPU even though the required API is present. The guard should only check the method the code actually calls.

[greptile-apps]

Strict REQUIRED check may reject valid Variant files

isVariantPhysicalType rejects any Parquet file where value or metadata is encoded with OPTIONAL repetition. Non-Spark Variant writers (e.g., Delta Lake connectors, Arrow-based tools) may write these fields as OPTIONAL BINARY — still spec-compliant Variant — which would fail this check and cause an incompatibility error or unexpected CPU fallback. Consider accepting OPTIONAL repetition as well, or at minimum document why REQUIRED is the only valid encoding here.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[nvauto]

NOTE: release/26.08 has been created from main. Please retarget your PR to release/26.08 if it should be included in the release.

[nartal1]

build

[nartal1]

@greptile-apps review the PR again.