Skip to content

feat: project Spark 4 VARIANT columns in native Parquet scans - #5407

Open
peterxcli wants to merge 1 commit into
apache:mainfrom
peterxcli:feat/native-variant-proj
Open

feat: project Spark 4 VARIANT columns in native Parquet scans#5407
peterxcli wants to merge 1 commit into
apache:mainfrom
peterxcli:feat/native-variant-proj

Conversation

@peterxcli

Copy link
Copy Markdown
Member

Which issue does this PR close?

This is Phase A of #4295.

It enables whole-value projection of top-level Spark 4 VariantType columns in ordinary native Parquet scans. It uses Arrow/Parquet's existing whole-value unshredding, but does not take on #3983's Parquet writer, subfield-pruning, or predicate-pushdown scope. Iceberg projection remains a later phase.

Rationale for this change

Spark 4 exposes semi-structured values as the atomic VariantType:

class VariantType private () extends AtomicType
case object VariantType extends VariantType

The logical type is backed by structured physical storage at the Arrow boundary. Spark's Arrow conversion produces a Struct with non-null Binary children in [value, metadata] order, and ColumnVector.getVariant hard-codes that child order:

return new VariantVal(getChild(0).getBinary(rowId), getChild(1).getBinary(rowId));

Arrow/Parquet identifies the same logical value with the Field-level extension name arrow.parquet.variant, whose storage type is a Struct. A Parquet reader can expose metadata, value, and optional shredded typed_value children in arbitrary order; Arrow-rs resolves metadata and value / typed_value by name, while unshred_variant reconstructs whole values and removes typed_value.

Before this PR, that logical identity could not survive Comet's native path:

Consequently, a query as small as SELECT v FROM parquet_table fell back rather than remaining an ordinary Comet native Parquet scan.

What changes are included in this PR?

The resulting path is:

Spark VariantType
  -> Comet protobuf VARIANT
  -> Arrow Field<Struct[value, metadata], arrow.parquet.variant>
  -> ordinary Parquet schema adapter
  -> VariantArray::try_new + unshred_variant
  -> Struct[value: Binary, metadata: Binary]
  -> Arrow C Data Interface Field
  -> CometStructVector (logical Spark VariantType)
  -> ColumnVector.getVariant

Preserve logical type identity

  • Appends protobuf VARIANT = 21 without renumbering existing values and serializes Spark 4 VariantType through the version shim (proto, Spark serialization).
  • Adds a protobuf-to-Arrow Field path that builds physical Struct<value: Binary, metadata: Binary> storage and attaches the canonical parent extension marker (native serde, Field construction).
  • Spark 4 shims recognize VariantType and recursively detect nested uses; Spark 3.x shims remain false/None, so Spark 3 behavior and compilation remain unchanged (Spark 4 shim, Spark 3 shim).

Normalize once at the ordinary Parquet boundary

The normalization is installed through the existing Parquet schema-adaptation path. The schema adapter pairs the logical target Field with the physical Parquet Field and installs CometCastColumnExpr; that expression normalizes the reader StructArray at evaluation. Marked Variant targets take this path even for otherwise-identity casts (schema adapter, identity-cast handling).

normalize_variant_array:

  1. constructs Arrow-rs VariantArray from the reader Struct;
  2. delegates whole-value reconstruction to Arrow-rs unshred_variant;
  3. converts BinaryView/LargeBinary output to ordinary Binary;
  4. rebuilds exactly two children in Spark's required [value, metadata] order; and
  5. preserves the parent null bitmap while the target Field retains its name, nullability, and extension metadata.

This handles both already-unshredded value + metadata input and shredded input containing typed_value, without adding another Variant dependency or a Comet-owned value decoder.

Preserve the Field through FFI and bridge it back to Spark

The shared FFI exporter now accepts the corresponding RecordBatch Field and builds FFI_ArrowSchema from that Field, leaving Arrow array export unchanged (export helper, batch export). Offset-normalized arrays still use the original output Field.

On import, Utils.fromArrowField maps only the explicit ARROW:extension:name = arrow.parquet.variant marker to the version-shim Variant type. Plain Structs keep their existing StructType behavior. No new vector is needed: the existing CometStructVector preserves child ordinals, so Spark's inherited getVariant consumes child 0 as value and child 1 as metadata.

Keep Phase A's fallback boundary explicit

Protobuf serialization here is schema transport, not general native Variant-expression support. The ordinary FileSource Parquet scan admits only a direct top-level VariantType field (scan gate). Direct Variant attributes are rejected by expression serde, and non-scan operators recursively reject Variant-bearing schemas (expression gate, operator gate).

The following remain deliberate Spark fallbacks:

  • PushVariantIntoScan / annotated VariantStruct output;
  • variant_get, predicates, casts, Variant functions, and Parquet writing;
  • native columnar-to-row, sort/limit and other native operators, shuffle, and spill;
  • nested Variant inside ARRAY/MAP/STRUCT; and
  • Iceberg Variant projection and equality deletes.

This builds on #5377: an unread Variant root is still pruned and its scan ordinals are still rebased instead of decoding it.

How are these changes tested?

The focused Spark test covers both unshredded and forced-shredded Parquet input, SELECT v, SELECT id, v, tail, object/array/scalar values, JSON null, SQL null, nullable parents, native scan retention, fallback from CometNativeColumnarToRowExec to JVM CometColumnarToRowExec, the logical Spark type, parent Field name/nullability/extension, exact [value, metadata] child order, Binary child types, and getVariant consumption (test). The SQL regression file also checks full-value projection plus the fallback matrix and #5377 pruning behavior.

Unshredded input uses Comet's byte-exact Spark answer checker. The shredded comparison is semantic because Spark's builder range-compresses integral values, while Arrow-rs selects its unshred builder from typed_value's Arrow primitive type and appends that typed value; Spark's VariantVal.equals compares the raw value and metadata byte arrays. The test separately verifies that Spark consumes the returned bytes through getVariant and that the reconstructed JSON values match.

Commands run:

make core

cd native
cargo fmt --all -- --check
env DYLD_LIBRARY_PATH="$JAVA_HOME/lib/server" \
  cargo test -p datafusion-comet test_normalize_shredded_variant_for_spark -- --nocapture
cargo clippy -p datafusion-comet --lib --tests -- -D warnings
cd ..

mvn -o -ntp -Pspark-4.0 -Dtest=none \
  '-Dsuites=org.apache.comet.CometSqlFileTestSuite variant' test
mvn -o -ntp -Pspark-4.0 -Dtest=none \
  '-Dsuites=org.apache.comet.parquet.ParquetReadV1Suite native scan projects Variant' test
mvn -o -ntp -Pspark-4.1 -Dtest=none \
  '-Dsuites=org.apache.comet.parquet.ParquetReadV1Suite native scan projects Variant' test
mvn -o -ntp -Pspark-4.1 -Dtest=none \
  '-Dsuites=org.apache.comet.CometIcebergNativeSuite variant' test
mvn -o -ntp -Pspark-4.0 -Dtest=none \
  '-Dsuites=org.apache.comet.rules.CometScanRuleSuite,org.apache.comet.rules.CometScanContribSuite' test
mvn -o -ntp -Pspark-3.5 -DskipTests test-compile

git diff --check upstream/main...HEAD

All commands passed. The focused results were 1/1 Rust normalization test, 1/1 Spark 4.0 SQL suite, 1/1 Spark 4.0 vector test, 1/1 Spark 4.1 vector test, 4/4 Spark 4.1 Iceberg fallback tests, and 16/16 Spark 4.0 scan-rule/contrib tests.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

I reviewed 45a0ed44ed9c58ede31410de077a0882e72fd4f8 against 92954d7884091d2c6fa3e109d11fc1b8cb4a7325, including the full 18-file diff and five independent review scopes. The scan-only design is a reasonable Phase A boundary: preserve Variant identity, reconstruct whole values at the Parquet boundary, and let Spark handle unsupported consumers. Seven P2 findings survived verification: six runtime compatibility regressions and one cross-version SQL-test failure confirmed in CI. Two runtime cases produce incorrect values, and four make previously valid operations fail.

Prior state and problem

Previously, a requested Spark Variant could not be serialized through Comet's type protobuf, and a native Arrow Struct could not be recovered as Spark's logical Variant type. Exporting an Arrow DataType rather than its Field also discarded the parent extension metadata needed to distinguish Variant storage from an ordinary Struct.

The existing pruning path already allowed scans to avoid unread Variant-bearing roots and retain supported siblings. This change extends that path to requested top-level Variant values, so it also makes previously unreachable default-value and downstream-consumer paths relevant.

Design approach

The PR appends VARIANT = 21, uses version shims to keep Spark 3 behavior inert, and represents Variant as a marked Arrow Field with Binary value and metadata children. The Parquet schema adapter installs a normalization expression that delegates reconstruction to Arrow's Variant implementation, converts its output to Binary, and restores Spark's child order.

Both native export call sites now pass Fields through the C Data Interface. The JVM recognizes the parent extension marker and reuses CometStructVector, allowing Spark's inherited getVariant to consume the two children.

Correctness / compatibility analysis

The ordinary projection path, name-based child lookup, and parent-null handling are supported by the added tests and focused review. The surviving problems are outside those examples: Variant existence defaults shift later defaults, reconstructed Unicode object keys are incompatible with Spark's lookup order, dictionary metadata is rejected before conversion, write wrappers evade the operator guard, the later Python rewrite bypasses it entirely, and Field export panics on a valid top-level NUL-containing name.

I built the JVM code at the requested head and used the macOS CI native artifact after verifying that its synthetic-merge tree is identical to the head tree. Five focused Scala cases and a PySpark comparison reproduced the six inline issues on Spark 4.0.4. Separate byte/serializer probes also checked the Unicode and Pandas cases against Spark 4.1.3. The PR's existing focused Variant projection test also passed locally, 1/1. These are targeted checks, not a full local Spark/native test run.

At the final refresh, CI had 48 successful, 2 failed, 13 running, and 7 skipped checks. The Spark 4.1 expressions job and Spark 4.2 expressions job both fail the new variant.sql:50 native-plan assertion because default Variant pushdown produces the deliberately unsupported VariantStruct representation. I inspected both job logs. Native builds, Rust tests, and the scan matrix are green, but the full CI run is not complete.

Key design decisions

Using an explicit extension marker is preferable to recognizing Variant from Struct shape, because ordinary user Structs must retain their existing meaning. Keeping datatype serialization separate from expression support is also appropriate for the proposed scope.

However, Arrow-compatible storage is not sufficient for every Spark consumer. Spark's object lookup order and Pandas Variant marker need compatibility handling or fallback. Similarly, a guard in tryConvertToComet cannot cover wrappers with empty outputs or operators introduced by a later transition rule.

Implementation sketch

On the Scala side, CometScanRule admits direct Variant roots, QueryPlanSerde transports the type, and CometNativeScan retains the requested logical field. Native schema construction attaches the extension marker, while the ordinary schema adapter wraps the physical reader column in CometCastColumnExpr.

Normalization resolves the reader's children by name, calls unshred_variant, and produces the two Binary children. The output batch's Field then accompanies its array through JNI, and Utils.fromArrowField restores Variant identity before Spark reads the vector. This is a compact path, but its new admission needs to be checked against existing default reconstruction and all downstream transitions.

Behavioral changes worth calling out

Whole-value top-level Variant scans can now remain native, including shredded layouts supported by the normalizer. Nested Variant, pushed VariantStruct, Iceberg projection, Variant expressions, shuffle, and native row conversion are intended to keep their existing fallback boundaries.

The change also affects more than direct projection: default expressions are now processed for admitted Variant schemas, and opt-in native-write and Python paths can receive Variant scans. Exporting the original Field name changes the common FFI path for non-Variant columns as well, which is why the NUL-name regression is included here.

Suggested improvements

Please address the six runtime cases with focused regressions and align the new native-projection SQL assertions with the supported scan configuration. Keep default values paired with their indexes, decode accepted reader representations before constructing VariantArray, and preserve Spark lookup behavior for reconstructed objects. The Unicode regression should include at least 32 keys with both supplementary and high-BMP characters, because a small ASCII-only object does not exercise Spark's binary-search path.

For the scan-only scope, apply fallback to the actual write input beneath WriteFilesExec and to the post-columnar Python rewrite instead of implicitly enabling those consumers. Preserve extension metadata without passing unsupported raw names to the C-string exporter. The existing successful projection tests should remain alongside these negative and compatibility cases. The SQL file also needs the same pushdown setting as the focused vector test so its native-plan assertions run on the intended path in Spark 4.1 and later.

Comment on lines +966 to +968
val schemaSupported = scanExec.requiredSchema.fields.forall { field =>
isVariantType(field.dataType) ||
typeChecker.isTypeSupported(field.dataType, field.name, fallbackReasons)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Reject unsupported Variant defaults before admitting the scan

Could we keep this scan on Spark when a required Variant field has a non-null existence default? CometNativeScan drops failed default serializations with flatMap, but retains every default index, and CometLiteral still rejects Variant. I reproduced this on Spark 4.0.4:

CREATE TABLE t(v VARIANT DEFAULT parse_json('1')) USING parquet;
INSERT INTO t VALUES (parse_json('42'));
ALTER TABLE t ADD COLUMNS(n INT DEFAULT 7);
SELECT v, n FROM t;

Spark returns (42, 7), while this head's native scan returns (42, NULL). The remaining default 7 is zipped to index 0 (v), where it is ignored because that column exists physically, leaving n without its default. Please reject an unserializable default or preserve and validate each value/index pair before enabling the scan.

}

let variant = VariantArray::try_new(array.as_ref())?;
let unshredded = unshred_variant(&variant)?;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Preserve Spark lookup compatibility when rebuilding Unicode objects

Could we account for Spark's object-key ordering before returning these reconstructed bytes? Arrow sorts object keys in UTF-8 order, but the supported Spark versions use Java String.compareTo and switch to binary search at 32 fields. I wrote a shredded Parquet object with Spark containing k00 through k29, U+E000, and 😀. With pushVariantIntoScan=false and allowReadingShredded=true, variant_get(v, '$.😀', 'int') returns 531 on Spark but NULL with this native scan. The expression itself correctly falls back to Spark, but it consumes the incompatible reconstructed ordering. The byte-level mismatch also reproduces on Spark 4.1.3. Please normalize for the Spark consumer or retain fallback for affected values, with a 32-key Unicode regression.

));
}

let variant = VariantArray::try_new(array.as_ref())?;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Decode dictionary metadata before constructing VariantArray

Could we decode dictionary-encoded metadata before this call? The canonical Arrow Variant representation permits it, and an Arrow-written Parquet file can retain metadata: Dictionary(Int32, Binary) in its embedded ARROW:schema while storing ordinary required BINARY children physically. I reproduced a file containing 42, 43, 44: Spark 4.0.4 reads it successfully, but this head's native scan throws Illegal shredded value type: Dictionary(Int32, Binary). Arrow/Parquet 58.4.0 restores the nested dictionary, which VariantArray::try_new rejects, so the Binary cast below is never reached. Decoding the metadata child first makes the same values readable.

Comment on lines +735 to +737
if (!op.isInstanceOf[CometScanExec] &&
(op.output ++ op.children.flatMap(_.output)).exists(attr =>
containsVariantType(attr.dataType))) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Check the write input beneath WriteFilesExec

Could this guard inspect the same unwrapped data-producing children used by requiresNativeChildren below? For DataWritingCommandExec(WriteFilesExec(CometNativeScan[Variant])), both the command output and WriteFilesExec.output are empty, so the Variant check misses the input. With spark.comet.parquet.write.enabled=true and spark.comet.operator.DataWritingCommandExec.allowIncompatible=true, copying a Spark-written Variant Parquet column now selects CometNativeWriteExec and fails in CometArrowStream.inputObjects -> Utils.toArrowSchema with Unsupported data type: ... VariantType ... variant. The intended Spark write fallback succeeds. Please apply the Variant check after unwrapping WriteFilesExec so this scan-only change does not enable the unsupported writer.

Comment on lines +739 to +741
op,
"Native operators do not support schemas containing type VariantType")
return None

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Apply the Variant fallback to the later Python rewrite too

Could we apply this boundary to EliminateRedundantTransitions.EligibleMapInBatch as well? That later rule creates CometMapInBatchExec without passing through this guard. I reproduced a native Variant scan followed by df.mapInPandas(lambda batches: batches, df.schema): it succeeds with spark.comet.exec.pyarrowUDF.enabled=false, but fails with the flag enabled. The accelerated runner forwards the new Arrow schema, which lacks Spark's variant=true metadata on the metadata child. Spark's Pandas serializer therefore supplies a dict rather than VariantVal, and the identity result fails assert isinstance(variant, VariantVal) during output conversion. Keeping Variant-bearing inputs on the ordinary Spark Python path would preserve the intended fallback.

unsafe {
std::ptr::write(array_ptr, FFI_ArrowArray::new(self));
std::ptr::write(schema_ptr, FFI_ArrowSchema::try_from(self.data_type())?);
std::ptr::write(schema_ptr, FFI_ArrowSchema::try_from(field)?);

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Handle NUL-containing field names before C schema export

Could we preserve the Field metadata without passing an embedded-NUL name to Arrow's C-string exporter? Spark accepts a top-level Parquet column named v\u0000suffix. I wrote and read that ordinary BIGINT column successfully with Spark 4.0.4, but the native scan at this head fails with NulError. Arrow 58.4.0's FFI_ArrowSchema::try_from(field) calls CString::new(field.name()).unwrap(), whereas the previous datatype-only export did not serialize the parent name. This affects non-Variant columns too, and the unaligned branch has the same issue. Please use a safe exported name or an explicit pre-execution fallback while retaining the logical metadata.

Comment on lines +49 to +50
query
SELECT v FROM test_variant

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Pin Variant pushdown for the native projection SQL assertions

Could we set spark.sql.variant.pushVariantIntoScan=false for these native-projection cases, as the new Scala vector test does? Spark 4.1 and 4.2 enable that optimizer rule by default, so even SELECT v becomes the annotated VariantStruct representation that this PR deliberately keeps on Spark. The plain query assertion then requires a native plan that cannot be produced. This is the actual failure in both Spark 4.1 CI and Spark 4.2 CI: variant.sql:50 fails with Expected only Comet native operators, but found Project. Please configure the whole-value test path explicitly and keep a separate fallback assertion for pushed VariantStruct.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants