bojana-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3950160361
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -958,6 +981,125 @@ private void appendWithNullStrippingImpl(
}
}
+ // A node in the keep-tree built from `variant_pick`'s JSONPaths. It records
whether a path
+ // terminates here (keep everything below), else which object keys / array
indices to descend
+ // into. Holds only strings and ints, so it is `Serializable` for
whole-stage codegen.
+ public static final class PickNode implements java.io.Serializable {
+ // A path terminates here: keep the whole value, subsuming any deeper
paths under this node.
+ private boolean keepAll = false;
+ // Lazily created; a node may hold both maps when paths disagree on the
container type, and
+ // only the map matching the actual value is used.
+ private HashMap<String, PickNode> objectChildren = null;
+ private HashMap<Integer, PickNode> arrayChildren = null;
+
+ // Insert the path suffix `path[depth..]` under this node.
+ private void add(PathSegment[] path, int depth) {
+ // A broader path already keeps everything here, so any deeper path is
subsumed.
+ if (keepAll) {
+ return;
+ }
+ if (depth == path.length) {
+ keepAll = true;
+ // Drop children from narrower paths added earlier; keepAll subsumes
them.
+ objectChildren = null;
+ arrayChildren = null;
+ return;
+ }
+ PathSegment seg = path[depth];
+ PickNode child;
+ if (seg instanceof ObjectKeySegment) {
+ if (objectChildren == null) {
+ objectChildren = new HashMap<>();
+ }
+ child = objectChildren.computeIfAbsent(((ObjectKeySegment) seg).key, k
-> new PickNode());
+ } else {
+ if (arrayChildren == null) {
+ arrayChildren = new HashMap<>();
+ }
+ child = arrayChildren.computeIfAbsent(((ArrayIndexSegment) seg).index,
k -> new PickNode());
+ }
+ child.add(path, depth + 1);
+ }
+ }
+
+ // Top-level entry for `pickAtPaths`: an object or array input yields a
(possibly empty) object
+ // or array; a scalar, variant null, or root `$` (`keepAll`) keeps the value
unchanged.
+ private void pickImplTopLevel(byte[] value, byte[] metadata, int pos,
PickNode root) {
+ checkIndex(pos, value.length);
+ int basicType = value[pos] & BASIC_TYPE_MASK;
+ if (root.keepAll || (basicType != OBJECT && basicType != ARRAY)) {
+ appendVariantImpl(value, metadata, pos);
+ } else {
+ pickImpl(value, metadata, pos, root);
+ }
+ }
+
+ // Append the substructures of the value at `pos` selected by `node`, and
return whether anything
+ // was appended. A caller drops a field or element whose pick produced
nothing by resetting
+ // `writePos`, so unmatched paths (missing keys, out-of-range indices, or
type mismatches) leave
+ // no trace. A dictionary key is registered only for a field that is
actually kept.
+ private boolean pickImpl(byte[] value, byte[] metadata, int pos, PickNode
node) {
+ checkIndex(pos, value.length);
+ if (node.keepAll) {
+ appendVariantImpl(value, metadata, pos);
+ return true;
+ }
+ int basicType = value[pos] & BASIC_TYPE_MASK;
+ if (basicType == OBJECT) {
+ return handleObject(
+ value, pos, (size, idSize, offsetSize, idStart, offsetStart,
dataStart) -> {
+ ArrayList<FieldEntry> fields = new ArrayList<>();
+ int start = writePos;
+ // No object-key children here: skip the whole scan and its per-field
key lookups.
+ if (node.objectChildren != null) {
+ for (int i = 0; i < size; ++i) {
+ int id = readUnsigned(value, idStart + idSize * i, idSize);
+ String fieldKey = getMetadataKey(metadata, id);
+ PickNode child = node.objectChildren.get(fieldKey);
+ if (child != null) {
+ int offset = readUnsigned(value, offsetStart + offsetSize * i,
offsetSize);
+ int fieldStart = writePos;
+ int fieldOffset = writePos - start;
+ if (pickImpl(value, metadata, dataStart + offset, child)) {
+ fields.add(new FieldEntry(fieldKey, addKey(fieldKey),
fieldOffset));
+ } else {
+ writePos = fieldStart;
Review Comment:
Good catch, addressed.
--
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]
