[SPARK-59151][SQL] Add a VariantBuilder.canonicalize primitive for Variant aggregation by aleksandr-chernousov-db · Pull Request #58455 · apache/spark

aleksandr-chernousov-db · 2026-09-01T16:23:53Z

What changes were proposed in this pull request?

This adds a canonicalization primitive to VariantBuilder (in common/variant).

public static Variant canonicalize(Variant v) returns a byte-canonical Variant, so two semantically-equal Variants produce byte-identical value/metadata. It:
- rebuilds the metadata dictionary sorted by UTF-8 key bytes, deduped, with unreferenced keys stripped and field ids remapped;
- normalizes scalars: smallest integer width, integer-valued decimals promoted to the integer encoding, decimal trailing zeros stripped, -0.0 -> +0.0, non-canonical NaN collapsed to the canonical bit pattern, short strings encoded as SHORT_STR. Types with no canonical rule (null/boolean/date/timestamp/timestamp_ntz/binary/uuid) pass through unchanged.
- A top-level (pos == 0) value that is already canonical is returned unchanged (fast path, below); a sub-variant (pos != 0, a view into a parent's shared arrays) is rebuilt into its standalone canonical form.
public static boolean isCanonical(byte[] value, byte[] metadata) is the read-side predicate whose true result guarantees canonicalize is a no-op. It is used as the fast path inside canonicalize.
VariantUtil gains Long.MIN_VALUE/Long.MAX_VALUE BigDecimal constants used by the decimal-to-long promotion check (hoisted out of the per-call path).

Why are the changes needed?

Hash-based grouping and comparison (GROUP BY / DISTINCT / equi-join) bucket by byte-equality, but two semantically-equal Variants can have different byte representations:

{"a": 1, "b": 2} vs {"b": 2, "a": 1}
different incoming object-key order
different metadata dictionary order
non-minimal integer/decimal encodings
-0.0 vs +0.0
non-canonical NaN bit patterns
short string stored as LONG_STR.
Without a canonical form, such values hash and compare unequal and land in different buckets, so grouping/deduplication on Variant would produce wrong results. canonicalize maps semantically-equal Variants to identical bytes, which is the building block for supporting Variant in aggregation/join keys. This PR introduces only the primitive and its tests. The expression and analyzer changes that use it will follow in subsequent PRs.

Does this PR introduce any user-facing change?

No

How was this patch tested?

New tests are introduced in VariantCanonicalizeSuite in this PR to cover all of the new code

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

Generated-By: Claude Opus 4.8

canonicalize(v) returns a byte-canonical Variant so semantically-equal Variants are byte-equal (enabling hash-agg bucketing / hash partitioning by value). isCanonical(value, metadata) is the read-side fast path: a true result guarantees canonicalize(v) is a no-op. Ported from the DBR POC and adapted to current apache conventions: - object/dictionary key ordering uses UTF-8 bytes (encodeKey/compareKeys), matching finishWritingObject and getFieldByKey; - isCanonical requires the exact metadata header result() emits (so a set sorted-strings bit is rebuilt, not wrongly accepted as canonical); - Long.MIN/MAX decimal bounds hoisted to VariantUtil static finals. Co-authored-by: Isaac <no-reply@databricks.com>

32 tests: structural canonicalization (key order, unused-key stripping, id remap), scalar normalization (minimal int width, integer-valued decimal promotion, trailing-zero strip, negative-zero to positive, canonical NaN, short-string re-encode), pass-through types, sub-variants, the isCanonical checks, a soundness oracle (isCanonical(v) implies bytesEqual(v, canonicalize(v))), and non-ASCII UTF-8-vs-UTF-16 key-ordering regression tests. Co-authored-by: Isaac <no-reply@databricks.com>

Co-authored-by: Isaac <no-reply@databricks.com>

scalastyle's NonASCII check decodes \u escapes, so the U+E000 / U+1F600 test keys tripped it; build them via Character.toChars(codePoint) so the source stays ASCII. Also add the missing space after /* in the inline /* allowDuplicateKeys = */ comment on the parse helper. Co-authored-by: Isaac <no-reply@databricks.com>

uros-b

@bojana-db @harshmotw-db Could you please help with reviewing this PR?

Also, @aleksandr-chernousov-db please update the PR title according to Spark conventions.

bojana-db

Made the first pass. Thank you for the PR!

bojana-db · 2026-09-02T15:04:11Z

  }

  private void appendVariantImpl(byte[] value, byte[] metadata, int pos) {
+    appendVariantImpl(value, metadata, pos, /* needNormalization */ false);


We should be consistent with the terminology, lets use canonicalization always.

Changed to "canonicalize"

bojana-db · 2026-09-02T15:23:19Z

+    ArrayList<String> keys = new ArrayList<>();
+    collectAllObjectKeys(value, metadata, pos, keys);
+    keys.sort((a, b) -> compareKeys(encodeKey(a), encodeKey(b)));
+    keys = new ArrayList<>(new LinkedHashSet<>(keys));


Do we need this, isnt add key idempotent (already a hashset)?

Changed to bitset with already visited keys

bojana-db · 2026-09-02T15:28:00Z

+  private void buildCanonicalized(byte[] value, byte[] metadata, int pos) {
+    ArrayList<String> keys = new ArrayList<>();
+    collectAllObjectKeys(value, metadata, pos, keys);
+    keys.sort((a, b) -> compareKeys(encodeKey(a), encodeKey(b)));


Should we encode it beforehand so we have no encoding per each comparison?

Now storing encoded keys (actually, just don't decode them in the first place)

bojana-db · 2026-09-02T15:42:30Z

+    // Fast path: a top-level (pos == 0) input that is already canonical is returned unchanged. A
+    // sub-variant (pos != 0) is a view into a parent's shared value/metadata, so it always takes
+    // the slow path, which reads the element at v.pos and rebuilds a standalone canonical Variant.
+    if (v.pos == 0 && isCanonical(v.value, v.metadata)) {


Just curious, why do we need this position check? The way i understand it, canonicalize will be called on top level variant only.

Sub-variants carry parents metadata and value, they just move pos
Since it has extra data, we need to recanonicalize it (and get rid of that extra data)

bojana-db · 2026-09-02T15:59:06Z

+          for (int i = 0; i < size; ++i) {
+            int id = readUnsigned(value, idStart + idSize * i, idSize);
+            int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize);
+            int elementPos = dataStart + offset;
+            keys.add(getMetadataKey(metadata, id));
+            collectAllObjectKeys(value, metadata, elementPos, keys);


Can we just take ids only and then after we recursed through the whole variant see the strings in the original metadata? This way we will have only one string decode per unique string.

Added bitmap (seen) to keep track of already encountered keys + postponed decoding to "after this cycle"

bojana-db · 2026-09-03T14:22:40Z

It would we good to have some fuzz tests testing legal but not canonical variants.

bojana-db · 2026-09-03T14:32:21Z

+    }
+  }
+
+  test("isCanonical rejects a non-minimal array offset width") {


Do we test anywhere larger widths? We should cover 2/3 bytes.

Added tests

bojana-db · 2026-09-03T14:35:13Z

+    assert(!bytesEqual(canon(parse("1.5")), canon(parse("2"))), "1.5 must not collapse to 2")
+  }
+
+  test("integer decimal too large for a long is not promoted (stays a decimal)") {


Do we test somewhere DECIMAL8/16 branches?

Added tests

bojana-db · 2026-09-03T14:40:06Z

+    assert(isCanon(parse("1")), "INT1(1) is canonical")
+  }
+
+  test("isCanonical rejects integer-valued and trailing-zero decimals") {


Do we test wider decimals somewhere also?

Added tests

bojana-db · 2026-09-03T14:44:41Z

+
+  // ----- isCanonical: object / array structure -----
+
+  test("isCanonical rejects a dictionary with an unused key") {


Do we cover the duplicate-field-id case?

Added tests

canonicalize() short-circuits on isCanonical, so the oracle's soundness assertion bytesEqual(v, canonicalize(v)) was trivially true whenever isCanonical(v) held. Extract doCanonicalize() -- the fast-path-free rebuild that canonicalize() now delegates to -- and have the oracle compare against it, so "isCanonical(v) => rebuild(v) == v" is actually verified. Addresses the review comment on the oracle assertion.

ZdravDim · 2026-09-07T09:25:07Z

+  }
+
+  // Smallest unsigned integer byte width that can hold `value`.
+  private static int minIntWidth(long value) {


minIntWidth(value) is a silent duplicate of the existing getIntegerSize(value). Code for these two functions is pretty similar, consider making getIntegerSize(value) do the assert and then call minIntWidth(value).

ZdravDim · 2026-09-07T09:48:26Z

+      build(_.appendBinary(Array[Byte](1, 2, 3, 4))),
+      build(_.appendUuid(new java.util.UUID(1L, 2L))))
+    for (v <- samples) {
+      assert(bytesEqual(v, canon(v)), "a pass-through scalar must be unchanged by canonicalize")


assert(bytesEqual(v, canon(v))) is a no-op here: these are already-canonical top-level values, so canonicalize returns v via the isCanonical fast path (appendCanonicalizedScalar never runs) and the assert compares v to itself. The isCanon(v) assert is fine; To actually cover this case, nest the value in a non-canonical structure so a rebuild is forced.

changed canon(v) to doCanonc(v) (which skips fast path)

- isCanonical now takes a Variant and returns false for a sub-variant (pos != 0): its value/metadata are a view into a parent, so it is never a standalone canonical variant. This moves the pos == 0 guard out of canonicalize and into isCanonical, making the predicate self-contained. - isCanonical now rejects a metadata dictionary whose offset[0] != 0 (padding before the keys) -- a form canonicalize never produces and the existing checks did not catch, which would have been a false positive. Addresses the review comments on offset[0] and the pos == 0 check.

canonicalize emits metadata of an exact length (header + dict-size field + (numKeys + 1) offsets + string bytes); isCanonical did not verify this, so trailing padding after the string region passed as canonical -- a false positive. Check metadata.length matches the expected size. Addresses the review comment about trailing bytes in the metadata.

The value analog of the metadata trailing-bytes check: isValueCanonical validates the value tree from position 0 but does not verify the value fills the whole array, so trailing padding after the value passed as canonical -- a false positive. Check valueSize(value, 0) == value.length. Addresses the review comment about trailing bytes after the value.

isValueCanonical checked the id-list and offset-list widths but not the size (element-count) field width. canonicalize emits a 1-byte size field unless size > U8_MAX (then 4 bytes; never 2 or 3), so a small object or array stored with a 4-byte size field is a canonicalize no-op only in appearance -- isCanonical wrongly returned true, a false positive that would split a GROUP BY bucket. Check the stored width in both the OBJECT and ARRAY branches, with two hand-crafted tests. Addresses the review comments about the object and array size-field width.

The width-minimality checks (metadata dict offset, object offset and id, array offset -- all via minIntWidth) were only ever exercised at width 1, because every functional test used tiny variants. Add positive coverage at widths 2 and 3: array and object offset widths (large element/field data), the object id-width 1->2 boundary (256 vs 257 keys), and the metadata dict offset width (many/long keys). Adds three width-reader test helpers. Addresses the review comment about covering 2/3-byte widths.

The decimal tests only covered DECIMAL4; DECIMAL8/16 were never asserted. Add three tests: canonicalize emits and isCanonical accepts DECIMAL8 and DECIMAL16; an over-wide decimal (trailing zeros inflating the stored precision) reduces to the minimal type; and isCanonical rejects a decimal stored in a wider type than its value needs (hand-crafted -- the type-width soundness check no parse input can reach), which canonicalize then reduces. Covers reductions DEC8->DEC4, DEC16->DEC4, and DEC16->DEC8. Test-only. Addresses the review comment about covering the DECIMAL8/16 branches.

isValueCanonical's OBJECT branch requires strictly ascending field ids (id <= prevId returns false), which also rejects duplicate keys -- but no test exercised it. The builder cannot produce a duplicate-field-id object (allowDuplicateKeys=true dedups, =false throws), so hand-craft one (id list [0, 0]) and assert isCanonical rejects it. Test-only. Addresses the review comment about the duplicate-field-id case.

isCanonical checked that the metadata dictionary is sorted by decoding each key to a String and re-encoding it (encodeKey(getMetadataKey(...))), about three allocations per key on a read-side fast path. Add VariantUtil.getMetadataKeyBytes, which returns a key's raw stored bytes with a single copyOfRange, and have the isCanonical sortedness loop compare those directly. Dictionary keys are stored as UTF-8, so the raw bytes equal encodeKey(getMetadataKey(...)) for every valid variant; the change is behavior-preserving and drops the decode + re-encode to one allocation per key. getMetadataKey is intentionally left unchanged rather than delegating to getMetadataKeyBytes: a delegating getMetadataKey would add a byte[] copy to each of its own callers, such as the getFieldByKey binary search. Add a test that getMetadataKeyBytes matches encodeKey(getMetadataKey) across keys of varying length. Addresses the review comment about skipping the decode/re-encode when comparing metadata keys in isCanonical.

buildCanonicalized collected the object keys as decoded Strings, sorted them by re-encoding both operands on every comparison, and removed duplicates with a LinkedHashSet before adding them to the dictionary. Collect the keys as their encoded (UTF-8) bytes instead (getMetadataKeyBytes), sort those bytes directly so the sort does no per-comparison encoding, and skip duplicates in the sorted list -- decoding a key to a String only once per unique key, for the addKey call. The separate LinkedHashSet dedup is dropped, since the adjacent skip already deduplicates. Behavior is unchanged: the same keys are collected, sorted into the same canonical (unsigned UTF-8) order, and unused dictionary keys are still stripped. Addresses the review comments about encoding keys on every sort comparison and the redundant dictionary dedup pass.

buildCanonicalized materialized a key's bytes at every object-field occurrence, then deduplicated after sorting -- so a key used in K objects was sliced K times and the sort ran over all occurrences. Mark each dictionary id in a `seen` bitmap as collectAllObjectKeys walks the value, and materialize a key's bytes only the first time its id is seen. Only the distinct keys are collected, so the sort runs over the unique keys and no separate dedup pass is needed. A small getMetadataNumKeys helper reads the dictionary size to allocate the bitmap. Behavior is unchanged: the same distinct keys are collected (unused keys are still never visited, so they stay stripped), sorted into the same canonical order, and a malformed id still throws MALFORMED_VARIANT. Addresses the review comment about taking ids only and decoding each unique key once.

The `appendVariantImpl` flag that gates scalar canonicalization was named `needNormalization` -- the only "normalization"-worded identifier in a feature otherwise named consistently (canonicalize, doCanonicalize, buildCanonicalized, appendCanonicalizedScalar, isCanonical). Rename it to `needCanonicalization` and fix a comment that referred to "normalization rules", so the terminology is uniform. One comment line is reflowed to stay within 100 columns. No behavior change -- this renames a private parameter and edits comments only. Addresses the review comment about using "canonicalization" naming consistently.

Note that numKeys can exceed lastOffset when the dictionary has an empty-string key (0 bytes), so the metadata offset width is sized for the max of the two. Comment-only, no behavior change. Addresses the review comment asking why lastOffset isn't always >= numKeys.

Generate a random logical value, render it to JSON two independent ways (object keys shuffled, and numbers spelled equivalently such as 5 / 5.0), parse both into byte-different but semantically equal Variants, and assert over 300 seeded iterations that canonicalize is: - consistent -- the two encodings canonicalize to identical bytes - complete -- isCanonical accepts canonicalize's output - sound -- doCanonicalize is a no-op on an already-canonical value The generator is self-contained (common/variant cannot depend on the sql/catalyst RandomDataGenerator) and small: an Obj/Arr/Scalar tree with recursion biased toward objects and a fixed seed for reproducibility. Addresses the review comment about adding fuzz tests for legal but not canonical variants.

getIntegerSize (write path) and minIntWidth (the isCanonical read path) computed the same smallest-unsigned-byte-width with identical thresholds. Make getIntegerSize keep its bounds assert and delegate to minIntWidth, so the width logic lives in one place and the writer and the isCanonical recognizer cannot drift. Addresses the review comment about minIntWidth duplicating getIntegerSize.

The test asserted bytesEqual(v, canonicalize(v)) for the date, timestamp, binary, and uuid scalars, but those are already canonical, so canonicalize short-circuits via the isCanonical fast path and the assertion compared v to itself. Compare against doCanonicalize (the fast-path-free rebuild) so the test actually exercises appendCanonicalizedScalar's pass-through handling. Addresses the review comment that this assertion was a no-op.

Collapse the aligned case arrows in gen/render to a single space before `=>` (Spark scalastyle requires exactly one), and wrap the fuzz soundness assertion in a block so no line exceeds 100 characters. Whitespace only, no behavior change.

aleksandr-chernousov-db and others added 4 commits September 1, 2026 15:03

remove bloated comments

f89e5d3

Co-authored-by: Isaac <no-reply@databricks.com>

uros-b reviewed Sep 2, 2026

View reviewed changes

aleksandr-chernousov-db changed the title ~~Spark 59151~~ [SPARK-59151][SQL] Add a VariantBuilder.canonicalize primitive for Variant aggregation Sep 2, 2026

fix linters again

1a2678f

bojana-db reviewed Sep 3, 2026

View reviewed changes

ZdravDim reviewed Sep 7, 2026

View reviewed changes

aleksandr-chernousov-db added 16 commits September 7, 2026 10:13

Reflow the fuzz test to satisfy scalastyle

18ddaf5

Collapse the aligned case arrows in gen/render to a single space before `=>` (Spark scalastyle requires exactly one), and wrap the fuzz soundness assertion in a block so no line exceeds 100 characters. Whitespace only, no behavior change.


		// ----- isCanonical: object / array structure -----

		test("isCanonical rejects a dictionary with an unused key") {

Uh oh!

Conversation

aleksandr-chernousov-db commented Sep 1, 2026

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

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

Uh oh!

uros-b left a comment

Choose a reason for hiding this comment

Uh oh!

bojana-db left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

bojana-db left a comment •

edited

Loading