fix: prevent concurrent access to non-thread-safe maps in SketchBufferAggregatorHelper

[zhang-arvin]

Description

Fixes #18040 - Runtime Exception when executing a query twice in a short duration.

Root Cause

The ParallelCombiner and StreamingMergeSortedGrouper are designed for concurrent read/write from different threads. However, SketchBufferAggregatorHelper used IdentityHashMap and Int2ObjectOpenHashMap, which are not thread-safe.

Concurrent access to these maps by:

• The writing thread (via aggregate() → getOrCreateUnion())
• The reading thread (via get())

could corrupt the internal map structure, causing ArrayIndexOutOfBoundsException during sketches aggregation (e.g., Index 180 out of bounds for length 129).

Fix

Replaced IdentityHashMap with ConcurrentHashMap and Int2ObjectOpenHashMap with ConcurrentHashMap to ensure thread safety when used with ParallelCombiner.

Changes

• extensions-core/datasketches/.../SketchBufferAggregatorHelper.java: Replace non-thread-safe maps with ConcurrentHashMap

Key Features/Changes

• Thread-safe access to Union cache and memory cache in SketchBufferAggregatorHelper
• Uses ConcurrentHashMap.computeIfAbsent for atomic map initialization

Verification

The existing tests for ParallelCombiner and StreamingMergeSortedGrouper cover the concurrent read/write pattern. The fix ensures that Datasketches-based aggregators work correctly under parallel combining.

[FrankChen021]

Severity	Findings
P0	0
P1	1
P2	0
P3	0
Total	1

Reviewed 1 of 1 changed files.

---

This is an automated review by Codex GPT-5.6-Luna(max)

[FrankChen021]

[P1] Do not key caches by mutable ByteBuffers

ConcurrentHashMap uses ByteBuffer.equals/hashCode, which depend on the buffer's remaining contents. Union updates mutate those contents in place, changing the key hash after insertion. Subsequent unions.get(buf) calls can miss and get() returns SketchHolder.EMPTY; distinct equal-content buffers can also alias. Preserve identity/stable-key semantics while making access concurrent.

[zhang-arvin]

Thanks @FrankChen021 — this is a real correctness bug and the fix is in 8576e12.

You are right that ConcurrentHashMap is the wrong structure here. It keys by ByteBuffer.equals/hashCode, which are derived from the buffer's remaining contents. Union updates mutate exactly those contents in place through the WritableMemory wrapping the same buffer, so an inserted key's hash changes as the aggregation progresses. The consequences are the ones you described: unions.get(buf) misses after the first mutation, get() returns SketchHolder.EMPTY, and two distinct buffers holding equal contents alias onto each other.

I went back to identity keying rather than trying to stabilise a content-derived key, because identity is the semantics this cache actually wants — the key is a specific arena instance, not a value:

private final IdentityHashMap<ByteBuffer, Int2ObjectMap<Union>> unions = new IdentityHashMap<>();
private final IdentityHashMap<ByteBuffer, WritableMemory> memCache = new IdentityHashMap<>();
private final Object lock = new Object();

IdentityHashMap is not thread-safe, so every access is now guarded by a monitor. That addresses the non-thread-safe access the previous change was intended to fix, without reintroducing the content-derived key. The critical sections are plain map operations on the query path, so a single monitor should not become a contention point — though if you'd prefer an RWMutex or per-buffer striping I'm happy to change it.

Also note createNewUnionLocked restores the get-then-put shape instead of computeIfAbsent: with IdentityHashMap there is no atomic primitive, hence the explicit lock and the Locked suffix on the helper to make the caller's obligation visible. relocate, getOrCreateUnion, createNewUnion and clear all hold the lock for their whole body so the multi-step invariants stay atomic.

Verification: mvn -pl extensions-core/datasketches -am package -DskipTests → BUILD SUCCESS (whole dependency chain, including the changed module).

[FrankChen021]

🟡 Changes recommended

The identity-keyed cache fix is correct for the reviewed grouper lifecycle: all IdentityHashMap and nested Int2ObjectMap access is serialized, while the existing buffer-slot ownership contract prevents getResult from racing an update on the same slot. One compatibility issue remains in the new raw-byte serialization shortcut: it is applied unconditionally to every ComplexMetricSerde, so byte[]-backed extension types can bypass their serializer and persist an invalid representation.

Reviewed 3 of 3 changed files: extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchBufferAggregatorHelper.java, processing/src/main/java/org/apache/druid/segment/data/ObjectStrategy.java, and processing/src/main/java/org/apache/druid/segment/serde/ComplexMetricSerde.java.

Severity	Findings
P0	0
P1	0
P2	1
P3	0
Total	1

---

This is an automated review by Codex GPT-5.6-Luna(max)

After addressing the findings or replying to the comments, you can request another review from me to trigger a new automated review.

[FrankChen021]

[P2] Preserve byte[] object-strategy semantics

Finding: This unconditional byte[] shortcut is now used by every ComplexMetricSerde, so a custom complex type whose legitimate intermediate object is byte[] (or whose strategy accepts byte[] and adds an envelope, compression, or type marker) silently skips its toBytes implementation. ComplexFieldWriter and ComplexFrameColumnWriter will persist the raw array, and the matching fromByteBuffer implementation can then fail or read a different value, breaking existing extension types.

Suggestion: Make the already-serialized path explicit or guard it using the strategy's declared object class, and add a regression test for a byte[]-backed strategy whose toBytes method transforms the input.