Skip to content

Broadcast serialization keeps 64-bit offsets for large_string and large_binary children of nested columns #6933

Open

Description

@andygrove

Describe the bug

Utils.serializeBatches writes a struct, list or map column as it stands, so a large_string or large_binary child keeps its 64-bit offsets. #6371 narrows top-level LargeVarCharVector and LargeVarBinaryVector columns to 32-bit offsets before they are written. getFieldVector accepts a StructVector, ListVector or MapVector without looking at its children, though, so a nested child is written with the offsets it arrived with.

That matters when the batches of one broadcast disagree on the width. It is the case #6371 fixes for top-level columns, one level down:

  • coalesceBroadcastBatches cannot append a struct<c: large_string> batch to a struct<c: string> one, in either order, so the broadcast is left uncoalesced.
  • CometArrowStream.reconcileStreamSchema takes the input stream's schema from the first batch, children included, and ColumnarBatchArrowReader loads every later batch's buffers under it. A later batch of the other width then has its offsets read at the wrong width. With the 64-bit batch first, reading the second batch's child from the reader's root fails with IndexOutOfBoundsException: index: 17179869184. With the 32-bit batch first, the second batch silently reads ["", "cccc"] where it holds ["cccc", "dddd"]. That root is what gets exported to the native plan.

When every batch agrees, the stream still carries Struct<c: LargeUtf8> where the plan expects Struct<c: Utf8>, so native ScanExec would cast the column in every task that reads the broadcast.

Like #5488, this is latent on main. CometBroadcastExchangeExec only broadcasts a plan whose children are all CometNativeExec, and CometMapInBatchExec, the operator that keeps a Python worker's vectors, is not one of them. The gap is at the serialization boundary, ready for when such output can reach it.

Steps to reproduce

Verified on main at 474b5ec (Spark 4.1) with a throwaway suite:

  1. Build two StructVector columns with one child each: a LargeVarCharVector holding aaaa and bbbb, and a VarCharVector holding cccc and dddd. Wrap each in a CometVector and a ColumnarBatch.
  2. Call Utils.serializeBatches on the two batches. Decoding the first with Utils.decodeBatches gives back a LargeVarCharVector child.
  3. Call Utils.coalesceBroadcastBatches on the two buffers. It returns a batch count of 0, the uncoalesced fallback.
  4. Decode both buffers and pass the batches through CometArrowStream.reconcileStreamSchema, with the schema Utils.toArrowSchema gives for st struct<c: string>, then through a ColumnarBatchArrowReader over the reconciled schema. Read the struct child of each batch the reader loads. The second batch fails or reads wrong values as above, depending on which width comes first.

Expected behavior

A nested large_string or large_binary child is written with 32-bit offsets, as #6371 does for top-level columns, so every batch of a broadcast has the same schema and native readers get the Utf8 or Binary types they plan for. A child holding more data than 32-bit offsets can address can still be refused.

Additional context

Raised by @parthchandra in review of #6371 (#6371 (comment)), which scopes its fix to top-level columns.

The cache serializer is not affected. CachedBatchIpc.matchesReaderLayout compares nested types too, and converts a batch whose layout does not match.

Two ways to close this:

  1. Narrow nested children in getBatchFieldVectorsWithProviders, rebuilding a struct, list or map vector around narrowed children and sharing every other buffer.
  2. Convert a batch whose nested types differ from the Spark schema's Arrow types, as the cache serializer does. serializeBatches is not given the schema today, so it would need passing in.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:ffiArrow FFI / JNI boundarybugSomething isn't workingpriority:lowMinor issues, test failures, tooling, cosmetic

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions