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 across all Spark 4.x shims (including Databricks 17.3), targeting tinyint, smallint, int, bigint, and string extractions from literal object-field paths. Batches that require Spark cast semantics are routed through a CPU bridge, and unsupported configurations fall back to CPU.

• GpuVariantGet implements extraction using VariantUtils.getVariantFieldValue / castVariantValue, builds a cuDF-native variant view with [metadata, value] child ordering, and narrows integer results with range-based null-on-overflow logic.
• GpuColumnVector and RapidsHostColumnVectorCore are extended to handle VariantType in host-type conversion, schema building, and type-compatibility checks; Parquet read-path support is added in GpuParquetScan and ParquetSchemaUtils.
• TypeChecks and GpuOverrides are updated with a new VARIANT TypeSig, a read/write split of additionalParquetSupportedTypes, and VariantGet registered as an incompat expression.

Confidence Score: 5/5

Safe to merge; all findings are style/rule-conformance issues that do not affect correctness under normal conditions.

The core extraction logic is well-structured, resource management is correct in all practically reachable paths, and the integration tests cover the full range of supported/unsupported cases. The bare .close() calls in the accumulator loops are a rule-conformance issue rather than a real leak risk since cuDF close() does not throw, and the missing shim entries in the test annotation are a coverage gap with no functional impact.

Files Needing Attention: GpuVariantGet.scala (allRowsCovered and normalizeIntegers accumulator loops); GpuColumnVectorVariantSuite.scala (shim annotation missing 404 and 413)

Important Files Changed

Filename	Overview
sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/GpuVariantGet.scala	New GPU expression for try_variant_get with CPU bridge fallback; bare .close() calls in allRowsCovered and normalizeIntegers violate C1, and evaluateOnCpu leaks an incRefCount if ColumnarBatch construction fails
sql-plugin/src/main/java/com/nvidia/spark/rapids/GpuColumnVector.java	Adds VariantType handling to host-type conversion, schema building, and typeConversionAllowed; logic is correct and resource management uses try-with-resources throughout
sql-plugin/src/main/scala/com/nvidia/spark/rapids/parquet/GpuParquetScan.scala	Adds VariantType cases to convertToParquetNative and checkSchemaCompat; isVariantPhysicalType correctly validates the value/metadata BINARY children
sql-plugin/src/main/scala/com/nvidia/spark/rapids/TypeChecks.scala	Adds VARIANT TypeEnum, TypeSig constant, and CastChecks entries; allSupportedTypes correctly excludes VARIANT for pre-4.x shims
sql-plugin/src/test/spark400/scala/com/nvidia/spark/rapids/GpuColumnVectorVariantSuite.scala	New unit tests for Variant column conversion; missing shim versions 404 and 413 in the json-lines annotation
integration_tests/src/main/python/variant_test.py	Comprehensive integration tests covering string/integral extraction, nested paths, null/missing fields, integral boundaries, heterogeneous values, filters, aggregates, shuffle, and fallback paths

Sequence Diagram

sequenceDiagram
    participant PR as GpuProjectExec
    participant VG as GpuVariantGet.doColumnar
    participant CU as VariantUtils (cuDF)
    participant BR as GpuCpuBridgeExpression

    PR->>VG: "input GpuColumnVector (VariantType STRUCT)"
    VG->>VG: "build cuDF view [metadata, value]"
    VG->>CU: "getVariantFieldValue(cudfVariant, path)"
    CU-->>VG: "rawValue ColumnView"
    alt "all rows GPU-decodable"
        VG->>CU: "castVariantValue(rawValue, targetDType)"
        CU-->>VG: "decoded ColumnVector"
        VG->>VG: "normalizeIntegers / narrowInteger"
        VG-->>PR: "result ColumnVector"
    else "coercion required"
        VG->>BR: "columnarEval(batch)"
        BR-->>VG: "CPU result ColumnVector"
        VG-->>PR: "result ColumnVector"
    end

Loading

_{Reviews (9): 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.

[nartal1]

build