Skip to content

DictLayoutDecoder extension-type fallback still expands strings eagerly #341

Description

@dfa1

Problem

DictLayoutDecoder.decode has three exits. Two are lazy; the third still expands (reader/layout/DictLayoutDecoder.java:110-119):

// Non-Utf8, non-Primitive dict — e.g. extension types backed by VarBin. Fall through
// to the existing string expansion for compatibility.
MemorySegment codesSegFallback = codes.materialize(arena);
...
return expandDictStrings(VarBinArray.toOffsetMode((VarBinArray) values, arena),
        codesSegFallback, codesPType, dtype, n, arena);

expandDictStrings (:248-259) then does the full two-pass gather:

long totalBytes = 0L;
for (long i = 0; i < n; i++) {
    long code = readUnsigned(codesSegs, i, codesPType);
    totalBytes += readUnsigned(valOffsets, code + 1, valOffPType) - readUnsigned(valOffsets, code, valOffPType);
}
MemorySegment outBytes = arena.allocate(totalBytes > 0 ? totalBytes : 1);
MemorySegment outOffsets = arena.allocate((n + 1) * 4L, 4);

Five lines above the fallback, the Utf8 branch handles a structurally identical dictionary lazily via VarBinArray.ofDict(...) (:92-93). The fallback is reached by VarBin-backed extension types — vortex.uuid being the obvious one — which then pay full dictionary expansion where a plain Utf8 column of the same shape pays nothing.

Root cause

The comment says it: "fall through to the existing string expansion for compatibility". The lazy Utf8 branch was added later and gated on valuesData instanceof VarBinOffsetArray plus dtype instanceof DType.Primitive, leaving extension dtypes on the original path rather than reasoning about whether they could join it.

Fix

Determine whether the fallback can route to VarBinArray.ofDict as well. The physical representation is the same — pool bytes, pool offsets, per-row codes — so the question is purely whether anything downstream of an extension dtype requires a contiguous VarBinOffsetArray. ExtensionDecoder implementations and Chunk.as() are the places to check; if they go through the VarBinArray interface rather than the concrete offset record, the fallback collapses into the branch above it and expandDictStrings can be deleted entirely.

If some extension decoder genuinely needs the flat form, VarBinArray.toOffsetMode(...) gives it that on demand at the point of use, rather than eagerly for every extension dict column.

The bufferCodes < n zip-bomb guard (:113-117) stays either way.

Context

Found in a sweep for remaining eager materializations after #329 / 7e0d6e7. Lower confidence than its siblings — the fix depends on what the extension decoders actually require, which I did not trace. See the related issues for the vortex.runend string expansion, vortex.sequence, and the vortex.dict encoding primitive path (which is the same asymmetry, in the other direction: there the layout is lazy and the encoding is not).

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions