cloud-fan commented on code in PR #56703:
URL: https://github.com/apache/spark/pull/56703#discussion_r3529869743
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -550,6 +564,146 @@ private void appendWithDeletionImpl(
}
}
+ private void appendWithInsertionImpl(
+ byte[] value, byte[] metadata, int pos, PathSegment[] segments, int
depth, Variant val) {
+ checkIndex(pos, value.length);
+ PathSegment seg = segments[depth];
+ boolean isLast = depth == segments.length - 1;
+ int basicType = value[pos] & BASIC_TYPE_MASK;
+ if (seg instanceof ObjectKeySegment && basicType == OBJECT) {
+ String key = ((ObjectKeySegment) seg).key;
+ handleObject(value, pos, (size, idSize, offsetSize, idStart,
offsetStart, dataStart) -> {
+ ArrayList<FieldEntry> fields = new ArrayList<>(size + 1);
+ int start = writePos;
+ boolean found = false;
+ 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;
+ String fieldKey = getMetadataKey(metadata, id);
+ boolean isTarget = fieldKey.equals(key);
+ found |= isTarget;
+ int newId = addKey(fieldKey);
+ fields.add(new FieldEntry(fieldKey, newId, writePos - start));
+ if (isTarget && !isLast) {
+ appendWithInsertionImpl(value, metadata, elementPos, segments,
depth + 1, val);
+ } else {
+ appendVariantImpl(value, metadata, elementPos);
+ }
+ }
+ if (isLast) {
+ // Append the new field unconditionally; if `key` already exists,
the duplicate is
+ // detected by `finishWritingObject`, which raises
VARIANT_DUPLICATE_KEY.
+ int newId = addKey(key);
+ fields.add(new FieldEntry(key, newId, writePos - start));
+ appendVariant(val);
+ } else if (!found) {
+ // Intermediate object key is missing; create the rest of the path
under it.
+ int newId = addKey(key);
+ fields.add(new FieldEntry(key, newId, writePos - start));
+ appendNewPath(segments, depth + 1, val);
+ }
+ finishWritingObject(start, fields);
+ return null;
+ });
+ } else if (seg instanceof ArrayIndexSegment && basicType == ARRAY) {
+ int index = ((ArrayIndexSegment) seg).index;
+ handleArray(value, pos, (size, offsetSize, offsetStart, dataStart) -> {
+ ArrayList<Integer> offsets = new ArrayList<>(size + 1);
+ int start = writePos;
+ if (isLast) {
+ // Insert `val` at `index`, shifting existing elements right.
+ for (int i = 0; i < size; ++i) {
+ if (i == index) {
+ offsets.add(writePos - start);
+ appendVariant(val);
+ }
+ int offset = readUnsigned(value, offsetStart + offsetSize * i,
offsetSize);
+ offsets.add(writePos - start);
+ appendVariantImpl(value, metadata, dataStart + offset);
+ }
+ if (index >= size) {
+ // Pad with variant nulls up to `index`, then append `val`.
+ for (int i = size; i < index; ++i) {
+ offsets.add(writePos - start);
+ appendNull();
+ }
+ offsets.add(writePos - start);
+ appendVariant(val);
+ }
+ } else if (index < size) {
+ // Descend into the existing element at `index`.
+ for (int i = 0; i < size; ++i) {
+ int offset = readUnsigned(value, offsetStart + offsetSize * i,
offsetSize);
+ int elementPos = dataStart + offset;
+ offsets.add(writePos - start);
+ if (i == index) {
+ appendWithInsertionImpl(value, metadata, elementPos, segments,
depth + 1, val);
+ } else {
+ appendVariantImpl(value, metadata, elementPos);
+ }
+ }
+ } else {
+ // Intermediate index is past the end; copy existing elements, pad
with variant nulls up
+ // to `index`, then create the rest of the path at `index`.
+ for (int i = 0; i < size; ++i) {
+ int offset = readUnsigned(value, offsetStart + offsetSize * i,
offsetSize);
+ offsets.add(writePos - start);
+ appendVariantImpl(value, metadata, dataStart + offset);
+ }
+ for (int i = size; i < index; ++i) {
+ offsets.add(writePos - start);
+ appendNull();
+ }
+ offsets.add(writePos - start);
+ appendNewPath(segments, depth + 1, val);
+ }
+ finishWritingArray(start, offsets);
+ return null;
+ });
+ } else {
+ // The segment kind does not match the container at this path prefix.
+ throw new VariantPathTypeMismatchException(depth);
+ }
+ }
+
+ // Build a fresh chain of containers for `segments[depth..]`, terminating in
`val`. Used to
+ // materialize missing intermediate path segments during insertion. The kind
of each segment
+ // decides the container created: an object-key segment creates a
single-field object, while an
+ // array-index segment creates an array padded with variant nulls up to the
index.
+ private void appendNewPath(PathSegment[] segments, int depth, Variant val) {
+ PathSegment seg = segments[depth];
+ boolean isLast = depth == segments.length - 1;
+ if (seg instanceof ObjectKeySegment) {
+ String key = ((ObjectKeySegment) seg).key;
+ ArrayList<FieldEntry> fields = new ArrayList<>(1);
+ int start = writePos;
+ int id = addKey(key);
+ fields.add(new FieldEntry(key, id, writePos - start));
+ if (isLast) {
+ appendVariant(val);
+ } else {
+ appendNewPath(segments, depth + 1, val);
+ }
+ finishWritingObject(start, fields);
+ } else {
+ int index = ((ArrayIndexSegment) seg).index;
+ ArrayList<Integer> offsets = new ArrayList<>(index + 1);
Review Comment:
`index` comes straight from the JSONPath, and the parser accepts any value
up to `Int.MaxValue` (SPARK-57880 rejects only larger). When this branch
creates a *fresh* intermediate array, `index + 1` overflows to
`Integer.MIN_VALUE` at `Int.MaxValue` -> `IllegalArgumentException: Illegal
Capacity: -2147483648`, and a large-but-valid index presizes a multi-GB array
-> OOM. Neither is caught by `insertAtPath` (which maps only the type-mismatch
and size-limit exceptions), so both surface as internal errors instead of
`VARIANT_SIZE_LIMIT`. The direct-insert path above (`new ArrayList<>(size +
1)`, line 612) stays graceful because the padding loop hits `checkCapacity`.
Presize on the default capacity so this path does too:
```suggestion
ArrayList<Integer> offsets = new ArrayList<>();
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
