Skip to content

[SPARK-58277][SQL] Stream DataType.json to bound peak memory for large schemas#57454

Open
bhollis-dbx wants to merge 1 commit into
apache:masterfrom
bhollis-dbx:bhollis-dbx/bounded-datatype-json
Open

[SPARK-58277][SQL] Stream DataType.json to bound peak memory for large schemas#57454
bhollis-dbx wants to merge 1 commit into
apache:masterfrom
bhollis-dbx:bhollis-dbx/bounded-datatype-json

Conversation

@bhollis-dbx

Copy link
Copy Markdown

What changes were proposed in this pull request?

DataType.json builds the type's full jsonValue json4s AST and then renders it to a string. Building the whole AST holds every node live at once, so peak memory scales with the node count of the schema rather than the output size.

This PR adds DataTypeJsonStreaming and routes DataType.json through it. It emits the type through a single Jackson JsonGenerator, recursing through structs, arrays, and maps so it holds at most one leaf field's AST at a time; peak memory then scales with the largest single field that must be rendered whole, not the whole schema.

StructField.jsonValue does more than emit name/type/nullable/metadata: for a string type with special properties (including one nested in an array or map) it moves those properties off the type and into the field's metadata. A field's type is streamed directly only when that move cannot apply — a struct-typed field always qualifies, since the move stops at a struct boundary, and any other type qualifies when it contains no collated string. Every other field is emitted by delegating to StructField.jsonValue, which remains the sole implementation of that move. To keep the output identical and avoid per-field overhead, the serializer reuses one SerializerProvider and json4s's JValueSerializer per call and writes to a SegmentedStringWriter (the same buffer-recycling writer ObjectMapper uses), which avoids reallocating an ever-larger buffer as output grows.

Why are the changes needed?

DataType.json holds the whole schema's json4s AST live while rendering, so peak memory scales with the total number of nodes in the schema, not with the output size. For a large schema this both spikes the driver's peak heap (enough to OOM it in the extreme) and generates a lot of short-lived garbage. Streaming the output bounds peak memory and cuts the transient allocation, with byte-identical results.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

A new DataTypeSuite case asserts DataTypeJsonStreaming.json is byte-identical to DataType.json (and round-trips via DataType.fromJson) across leaf, wide (200-field), deeply nested, struct-under-array/map, collation-under-array/map, and mixed-collated schemas — covering both the streaming path and the whole-field fallback. Existing DataTypeSuite JSON ser/de tests continue to pass.

Measured against the previous compact(render(...)) implementation across width, depth, and nested-union schema shapes: throughput is at parity or better (small schemas roughly on par, a large schema noticeably faster), and peak retained heap on the large union is meaningfully lower while serializing.

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

Generated-by: Claude (Claude Code)

…e schemas

DataType.json builds the type's full json4s AST via jsonValue and then renders
it to a string. That AST is much larger than the resulting string, so for a very
wide or deeply nested schema (e.g. a protobuf oneof with 100+ branches
reprojected into nested structs) it can exhaust the driver heap even though the
JSON output is comparatively small.

Add DataTypeJsonStreaming and route DataType.json through it. It emits the type
through a single Jackson generator, holding at most one field's AST at a time, so
peak memory scales with the largest single non-struct field rather than the whole
schema.

Only StructType.fields and struct-typed fields are written directly; every other
field is emitted via StructField.jsonValue, which remains the sole place collation
on a nested string/array/map type is moved onto the field metadata. Recursing a
struct-typed field is safe because that move never crosses a struct boundary.

Output is byte-for-byte identical to the previous implementation. Verified against
the old output (and fromJson round-trip) over wide, deeply nested, and
mixed-collated array/map schemas. Throughput is at parity or better across schema
shapes, and a 175-branch nested union retains roughly a third less live heap while
serializing.

@uros-b uros-b 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.

@bhollis-dbx Please note that another PR has already been opened for the same issue: #57453. cc @marcuslin123 to align on the path forward

@bhollis-dbx

Copy link
Copy Markdown
Author

@uros-b @marcuslin123 awesome, happy to close this in favor of that PR. Feel free to take whatever you like from this.

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