From 1bb2a5681a292aff03d3799b97ccc6528e410dac Mon Sep 17 00:00:00 2001 From: Michael Simacek Date: Tue, 20 Oct 2026 19:31:15 +0200 Subject: [PATCH] Optimize creating dicts in JSON scanner --- .../src/tests/test_json.py | 33 ++++++++++++ .../modules/json/JSONScannerBuiltins.java | 33 ++++++++---- .../objects/common/EconomicMapStorage.java | 9 ++++ .../objects/common/ObjectHashMap.java | 51 +++++++++++++++++++ .../graal/python/util/ArrayBuilder.java | 12 +++++ 5 files changed, 129 insertions(+), 9 deletions(-) diff --git a/graalpython/com.oracle.graal.python.test/src/tests/test_json.py b/graalpython/com.oracle.graal.python.test/src/tests/test_json.py index 56ef337f6a..0c3cde0796 100644 --- a/graalpython/com.oracle.graal.python.test/src/tests/test_json.py +++ b/graalpython/com.oracle.graal.python.test/src/tests/test_json.py @@ -60,6 +60,39 @@ class JsonTest(unittest.TestCase): + def test_load_duplicate_nested_keys(self): + text = '{"a": {"old": 0}, "b": [{"x": 1}], "a": {"last": [{"y": 2}]}, "c": {}}' + result = json.loads(text) + assert list(result.items()) == [("a", {"last": [{"y": 2}]}), ("b", [{"x": 1}]), ("c", {})] + result["d"] = 3 + del result["b"] + assert list(result) == ["a", "c", "d"] + assert result.popitem() == ("d", 3) + assert json.loads(json.dumps(result)) == result + + def test_load_dict_index_sizes(self): + for size in (0, 1, 4, 8, 32, 130, 1024, 33000): + expected = {str(i): i for i in range(size)} + result = json.loads(json.dumps(expected)) + assert list(result.items()) == list(expected.items()) + result["extra"] = None + for i in range(0, size, 3): + assert result.pop(str(i)) == i + result.update(expected) + assert result == dict(expected, extra=None) + + def test_load_duplicate_colliding_keys(self): + # These strings have the same Java-style string hash. + keys = [""] + for _ in range(8): + keys = [key + suffix for key in keys for suffix in ("Aa", "BB")] + pairs = [(key, i) for i, key in enumerate(keys)] + pairs += [(key, -i) for i, key in enumerate(reversed(keys))] + text = "{" + ",".join(json.dumps(key) + ":" + str(value) for key, value in pairs) + "}" + result = json.loads(text) + assert list(result.items()) == list(dict(pairs).items()) + assert result == json.loads(text, object_pairs_hook=dict) + def test_callable_object_hook(self): called = [] class Hook: diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/json/JSONScannerBuiltins.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/json/JSONScannerBuiltins.java index 816e830db5..8395830585 100644 --- a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/json/JSONScannerBuiltins.java +++ b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/json/JSONScannerBuiltins.java @@ -459,6 +459,7 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall ObjectSequenceStorage currentListStorage = null; ObjectSequenceStorage currentPairsStorage = null; EconomicMapStorage currentDictStorage = null; + ArrayBuilder currentDictEntries = null; ScannerState state = ScannerState.initial; boolean commaSeen = false; int commaIndex = -1; @@ -482,7 +483,7 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall TruffleString parentKey = null; final Object dict; /* - * If no hooks are present, the stack contains only the parent dicts or + * If no hooks are present, the stack contains only the parent entry buffers or * lists the current object is nested in. If a objectHook or objectPairsHook * is present, the stack also stores the property keys of nested dicts. For * example, when parsing the innermost dict of '{"a":{"b":{"c":"d"}}}', the @@ -508,9 +509,10 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall assert ((PDict) topOfStack).getDictStorage() == currentDictStorage; dict = callObjectHook.executeObject(objectHook, topOfStack); } else { - assert topOfStack instanceof PDict; - assert ((PDict) topOfStack).getDictStorage() == currentDictStorage; - dict = topOfStack; + // Adopt the entry buffer and build the hash index only once, at }. + assert topOfStack == currentDictEntries; + EconomicMapStorage storage = EconomicMapStorage.fromStringEntries(currentDictEntries.getInternalArray(), currentDictEntries.size() / 2, hashCodeNode, equalNode); + dict = PFactory.createDict(language, storage); } if (stack.isEmpty()) { return dict; @@ -537,11 +539,13 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall } } else { Object parent = stack.peek(); - if (parent instanceof PDict parentDict) { - currentDictStorage = (EconomicMapStorage) parentDict.getDictStorage(); + if (parent instanceof ArrayBuilder entries) { + currentDictEntries = entries; + currentDictEntries.set(currentDictEntries.size() - 1, dict); nextState = ScannerState.dict; } else { currentListStorage = (ObjectSequenceStorage) ((PList) parent).getSequenceStorage(); + currentListStorage.setObjectItemNormalized(currentListStorage.length() - 1, dict); nextState = ScannerState.list; } } @@ -593,9 +597,12 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall } else if (hasPairsHook) { currentPairsStorage = (ObjectSequenceStorage) parent; nextState = ScannerState.dict; - } else { + } else if (hasObjectHook) { currentDictStorage = (EconomicMapStorage) ((PDict) parent).getDictStorage(); nextState = ScannerState.dict; + } else { + currentDictEntries = (ArrayBuilder) parent; + nextState = ScannerState.dict; } } } @@ -612,8 +619,10 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall idx++; if (hasPairsHook) { value = new ObjectSequenceStorage(4); - } else { + } else if (hasObjectHook) { value = PFactory.createDict(language, EconomicMapStorage.create()); + } else { + value = new ArrayBuilder<>(8); } nextState = ScannerState.dict; } else if (c == '[') { @@ -639,9 +648,12 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall assert propertyKey != null; if (hasPairsHook) { currentPairsStorage.appendItem(PFactory.createTuple(language, new Object[]{propertyKey, value})); - } else { + } else if (hasObjectHook) { HashingStorage newStorage = hashingStorageSetItem.execute(inliningTarget, currentDictStorage, propertyKey, value); assert newStorage == currentDictStorage; + } else { + currentDictEntries.add(propertyKey); + currentDictEntries.add(value); } } else if (state == ScannerState.list) { assert propertyKey == null; @@ -667,6 +679,9 @@ private Object scanOnceUnicode(VirtualFrame frame, BoundaryCallData boundaryCall */ stack.add(propertyKey); } + } else if (!hasObjectHook) { + assert nextState == ScannerState.dict; + currentDictEntries = (ArrayBuilder) value; } else { assert nextState == ScannerState.dict; currentDictStorage = (EconomicMapStorage) ((PDict) value).getDictStorage(); diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/EconomicMapStorage.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/EconomicMapStorage.java index 75f09f299f..69ac56a32c 100644 --- a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/EconomicMapStorage.java +++ b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/EconomicMapStorage.java @@ -83,6 +83,15 @@ private EconomicMapStorage(int initialCapacity) { super(initialCapacity); } + /** Takes ownership of the array; callers must not use it after this call. */ + public static EconomicMapStorage fromStringEntries(Object[] entries, int entryCount, TruffleString.HashCodeNode hashCodeNode, TruffleString.EqualNode equalNode) { + return new EconomicMapStorage(entries, entryCount, hashCodeNode, equalNode); + } + + private EconomicMapStorage(Object[] entries, int entryCount, TruffleString.HashCodeNode hashCodeNode, TruffleString.EqualNode equalNode) { + super(entries, entryCount, hashCodeNode, equalNode); + } + private EconomicMapStorage() { this(4); } diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/ObjectHashMap.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/ObjectHashMap.java index f07f0076a3..b74a2cdc89 100644 --- a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/ObjectHashMap.java +++ b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/ObjectHashMap.java @@ -40,10 +40,14 @@ */ package com.oracle.graal.python.builtins.objects.common; +import static com.oracle.graal.python.util.PythonUtils.TS_ENCODING; import static com.oracle.truffle.api.CompilerDirectives.SLOWPATH_PROBABILITY; +import java.util.Arrays; + import com.oracle.graal.python.builtins.objects.common.ObjectHashMapFactory.PutNodeGen; import com.oracle.graal.python.builtins.objects.common.ObjectHashMapFactory.RemoveNodeGen; +import com.oracle.graal.python.lib.PyObjectHashNode; import com.oracle.graal.python.lib.PyObjectRichCompareBool; import com.oracle.graal.python.runtime.sequence.storage.ObjectSequenceStorage; import com.oracle.graal.python.util.ArrayBuilder; @@ -65,6 +69,7 @@ import com.oracle.truffle.api.profiles.InlinedBranchProfile; import com.oracle.truffle.api.profiles.InlinedConditionProfile; import com.oracle.truffle.api.profiles.InlinedCountingConditionProfile; +import com.oracle.truffle.api.strings.TruffleString; /** * Generic dictionary/set backing storage implementation. @@ -227,6 +232,52 @@ public ObjectHashMap() { allocateData(INITIAL_INDICES_SIZE); } + /** + * Takes ownership of an alternating string key/value array and builds its hash index. Duplicate + * keys are compacted in place, retaining the first key and the last value. + */ + protected ObjectHashMap(Object[] entries, int entryCount, TruffleString.HashCodeNode hashCodeNode, TruffleString.EqualNode equalNode) { + assert (entries.length & 1) == 0 && entries.length > 0; + assert entryCount >= 0 && entryCount <= entries.length / 2; + int entryCapacity = entries.length / 2; + int bucketsCount = getMinBucketsCount(entryCapacity); + ensureArraySizesFit(bucketsCount, entryCapacity); + keysAndValues = entries; + metadata = createMetadata(bucketsCount, entryCapacity); + int indexByteSize = getIndexByteSize(entryCapacity); + int indicesOffset = getIndicesOffset(entryCapacity); + int collisionMask = getPhysicalCollisionMaskForIndexByteSize(indexByteSize); + for (int i = 0; i < entryCount; i++) { + TruffleString key = (TruffleString) entries[i * 2]; + Object value = entries[i * 2 + 1]; + long hash = PyObjectHashNode.hash(key, hashCodeNode); + int bucket = getIndex(bucketsCount, hash); + long perturb = hash; + while (true) { + int index = getIndex(metadata, indicesOffset, indexByteSize, bucket); + if (index == EMPTY_INDEX) { + setIndex(metadata, indicesOffset, indexByteSize, collisionMask, bucket, size + INDEX_OFFSET); + entries[size * 2] = key; + entries[size * 2 + 1] = value; + setHash(metadata, size, hash); + size++; + usedHashes++; + usedIndices++; + break; + } + int existingIndex = unwrapIndex(index); + if (getHash(metadata, existingIndex) == hash && equalNode.execute((TruffleString) getKey(existingIndex, entries), key, TS_ENCODING)) { + entries[existingIndex * 2 + 1] = value; + break; + } + markCollision(metadata, indicesOffset, indexByteSize, collisionMask, bucket); + perturb >>>= PERTURB_SHIFT; + bucket = nextIndex(bucketsCount, bucket, perturb); + } + } + Arrays.fill(entries, size * 2, entryCount * 2, null); + } + protected ObjectHashMap(ObjectHashMap original) { size = original.size; usedHashes = original.usedHashes; diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/util/ArrayBuilder.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/util/ArrayBuilder.java index 58d170bcae..755f5aaf59 100644 --- a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/util/ArrayBuilder.java +++ b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/util/ArrayBuilder.java @@ -83,6 +83,11 @@ public T get(int index) { return (T) data[index]; } + public void set(int index, Object item) { + assert index >= 0 && index < size; + data[index] = item; + } + @SuppressWarnings("unchecked") public T pop() { assert size > 0; @@ -129,6 +134,13 @@ public ObjectSequenceStorage toObjectSequenceStorage() { return new ObjectSequenceStorage(data, size); } + /** + * Returns the backing array, which may be overallocated. + */ + public Object[] getInternalArray() { + return data; + } + public int size() { return size; }