Skip to content

Native Celeborn shuffle: validate_remote_schema rejects dictionary shapes the shuffle writer can emit #5536

Description

@andygrove

Discovered during post-merge review of #5531.

Component: native/shuffle/src/remote_schema.rs:47

validate_remote_schema unwraps exactly one wire shape, a top-level
Dictionary(Int32, Utf8 | Binary). Every other dictionary is a hard mismatch, including
Dictionary(Int32, Int32), Dictionary(Int8, Utf8), and any dictionary nested inside a list, map,
or struct.

The local reader is far more permissive. unpack_dictionary (shuffle_scan.rs:254) casts any
top-level dictionary to its value type, and cast_and_stamp_schema absorbs nested ones. So the two
readers disagree about what a valid frame is, and the disagreement is asymmetric: a frame that reads
fine locally becomes, on Celeborn, a validate_remote_schema error, which the reader converts into
a FetchFailedException, which triggers a map-stage rerun that reproduces the same frame, ending in
job failure at spark.stage.maxConsecutiveAttempts.

This is not reachable today

I instrumented read_ipc_compressed_impl to log any decoded wire schema containing a Dictionary
at any nesting depth and ran 125 tests across CometFuzzTestSuite, CometShuffleSuite, and
CometNativeShuffleSuite. Zero hits. SchemaAlignExec aligns the writer input to the
catalyst-declared schema and ScanExec::unpack_dictionary_type strips top-level dictionaries before
that, so nothing dictionary-encoded currently reaches the wire.

Why it is still worth fixing

ShuffleBlockWriter demonstrably encodes nested dictionaries. Both frames below were produced by
ShuffleBlockWriter::write_batch, not hand-built:

wire type      = List(Dictionary(Int32, Utf8))
spark declares = List(Utf8)
validate       = Err(Execution("Shuffle block type mismatch at column 0:
                 got List(Dictionary(Int32, Utf8)) but expected List(Utf8)"))
local cast ok  = 3 rows

Dictionary(Int32,Int32) vs Int32
               = Err(Execution("Shuffle block type mismatch at column 0:
                 got Dictionary(Int32, Int32) but expected Int32"))

Nothing pins the "no dictionaries on the wire" invariant. The day something upstream stops stripping
them, local shuffle keeps working and Celeborn shuffle starts failing jobs, with no compile error
and no failing test in between.

Suggested fix

Either:

  1. Recurse the dictionary unwrap through the nested List / LargeList / FixedSizeList / Map /
    Struct arms of same_logical_type, so the validator accepts whatever the encoder can emit; or
  2. Add a test asserting the writer never emits a dictionary shape the reader would reject, so a
    drift in either direction breaks the build.

Option 1 alone still leaves Dictionary(Int8, Utf8) and Dictionary(Int32, Int32) rejected, so
option 2 is the one that actually closes the gap.

Verified at merge commit 98cd8c967.

Activity

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

Metadata

Metadata

Assignees

Labels

area:shuffleShuffle (JVM and native)bugSomething isn't workingcorrectnesspriority:mediumFunctional bugs, performance regressions, broken features

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions