Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #19601 +/- ##
============================================
- Coverage 67.83% 67.81% -0.02%
Complexity 1450 1450
============================================
Files 3504 3504
Lines 226581 226594 +13
Branches 35804 35812 +8
============================================
- Hits 153699 153668 -31
- Misses 60754 60805 +51
+ Partials 12128 12121 -7
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
gortiz
left a comment
gortiz
left a comment
There was a problem hiding this comment.
Reviewed the four changes independently. Summary of what I verified and what I'd like changed before merge — none of it blocking, but the javadoc one matters most.
Verified correct
- The scratch-array reuse is safe against HEAD: both
TwoKeysGroupIdGeneratorandMultiKeysGroupIdGeneratorproject the key into ints (longpacking /FixedIntArray) before storing anything, so neither retains the passedObject[]. extractKeyper column andextractKeysproduce identical values, including null-bitmap handling, andnumRowsnow comes fromdataBlock.getNumberOfRows(), which is whatextractKeyused internally. No semantic change.- The
new Object[numFunctions][]change is a pure dead-store removal. TypeUtils.convertidentity preservation is unobservable for the four immutable wrapper types, and NPE-on-null behavior is unchanged.putIfAbsentis not slower:Int2IntOpenHashMapoverrides it with a single-probefind()+insert()(checked against fastutil 8.5.15 bytecode), so theidMapPutIntdelta in the description is noise.
Would like changed
- Move the "must not retain the key array" guarantee onto
GroupIdGenerator#getGroupIdjavadoc rather than a caller-side comment. This is the one change here that can silently produce wrong results if a future implementation breaks the contract. ObjectToIdMaphas the byte-identical pattern and wasn't updated — it's the impl used for STRING keys, the most common group-by key type.- The filtered twin
generateGroupByKeys(DataBlock, numMatchedRows, matchedBitmap)still pays the fullextractKeyscost, leaving two near-identical methods optimized differently.
Description nits
- The
INVALID_KEYchange is presented as a correctness fix, but the oldid == numValuestest looks unreachable by construction (ids are always< size()). Worth softening so reviewers don't hunt for a bug that isn't there. - Conversely, change #1's win is understated: it removes the whole
Object[][]materialization (~40 MB → ~16 MB for 2 keys × 1M rows), which escape analysis can't account for either way. TypeUtils.convert's hottest callers areLeafOperator:696/733(per row, per column), not group-by merge — the change is broader than the title suggests.
Details inline.
| columns[i] = DataBlockExtractUtils.extractKey(dataBlock, _groupKeyIds[i]); | ||
| } | ||
| // Multi-column generators retain dictionary IDs, not this array, just as in the row-heap path. | ||
| Object[] key = new Object[numKeys]; |
There was a problem hiding this comment.
Please move this guarantee onto the interface.
key is now a single scratch array handed to _groupIdGenerator.getGroupId(key) for every row. That is correct only while no GroupIdGenerator retains the array. I checked both implementations reachable when numKeys >= 2:
TwoKeysGroupIdGenerator.getGroupIdreadskeyValues[0]/[1], maps them to ints and packs them into alongbefore touching_groupIdMap. Nothing keeps the array.MultiKeysGroupIdGenerator.getGroupIdcopies into a freshint[] keyIdsand storesnew FixedIntArray(keyIds). Nothing keeps the array.
So the change is correct against HEAD, and it makes the serialized path match what the row-heap overload already does.
My concern is durability, not current correctness. A future generator that caches or stores the incoming Object[] — e.g. an Object2IntOpenHashMap<Object[]> with a custom hash strategy, which is a natural thing to write — would corrupt every group, and no existing test would obviously point at the cause. Could the guarantee go onto GroupIdGenerator#getGroupId as javadoc? Something like "the key array is scratch and may be mutated after this call returns; implementations must not retain it." That is where the next implementer will look; a comment on the caller is invisible to them.
| intKeys[i] = _groupIdGenerator.getGroupId(keys[i]); | ||
| } | ||
| } else { | ||
| Object[][] columns = new Object[numKeys][]; |
There was a problem hiding this comment.
The allocation win here is larger, and different in kind, from what the description claims — worth correcting because it currently undersells the change.
The description frames this as "reuse removes N-1 of N per-row arrays, hidden from JMH by escape analysis". That is only half of it. The old extractKeys(dataBlock, _groupKeyIds) materialized an Object[numRows][] outer array plus numRows inner Object[numKeys] arrays, all live at once. The new code materializes numKeys column arrays of numRows each.
For 2 keys over 1M rows, roughly:
- before:
1M * (16B header + 16B refs)+ 8MB outer ≈ 40 MB - after:
2 * (16B + 8MB)≈ 16 MB
That reduction is escape-analysis-proof and does not depend on the scratch-array argument at all — it is simply the removal of the Object[][] materialization. I would lead with that instead of the "JMH understates it" narrative.
Secondary bonus: extraction is now column-major, which suits extractValue's access pattern better than the old row-major nesting.
| for (int i = 0; i < numKeys; i++) { | ||
| key[i] = columns[i][rowId]; | ||
| } | ||
| intKeys[rowId] = _groupIdGenerator.getGroupId(key); |
There was a problem hiding this comment.
The filtered twin of this method was left out — deliberate, or an oversight?
generateGroupByKeys(DataBlock, int numMatchedRows, RoaringBitmap matchedBitmap) (a few methods below) is the same shape and still does:
keys = DataBlockExtractUtils.extractKeys(dataBlock, _groupKeyIds, numMatchedRows, matchedBitmap);
for (...) intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);Its row-heap sibling already uses the reused-scratch pattern, so the same reasoning applies verbatim, and extractKey(dataBlock, colId, numMatchedRows, matchedBitmap) already exists. That path runs for filtered aggregations (COUNT(*) FILTER (WHERE ...)), which is not a rare shape.
Either do it in this PR (~10 lines, same argument) or note in the description why it was deferred. Otherwise the file is left with two near-identical methods optimized differently, which is the kind of drift that later gets copied in the wrong direction.
| int numRows = groupByKeys.length; | ||
| int numFunctions = _aggFunctions.length; | ||
| Object[][] intermediateResults = new Object[numFunctions][numRows]; | ||
| Object[][] intermediateResults = new Object[numFunctions][]; |
There was a problem hiding this comment.
Safe dead-store removal, no behavior change — the inner arrays were unconditionally overwritten on the very next statement.
Two things worth stating explicitly for future readers: this is only safe because getIntermediateResults assigns every index and never returns null (it does), and nothing downstream depends on intermediateResults[j] having exactly numRows entries — the merge loops index by row and would have gone out of bounds under the old code too if the lengths disagreed. So the old allocation was not acting as a safety net.
| int id = _valueToIdMap.computeIfAbsent(value, k -> numValues); | ||
| if (id == numValues) { | ||
| int id = _valueToIdMap.putIfAbsent(value, numValues); | ||
| if (id == INVALID_KEY) { |
There was a problem hiding this comment.
Two points on this one.
1. I don't think the old code could actually misbehave, so I'd soften the correctness claim in the description.
The description says the old id == numValues test "could falsely report absent when an existing key legitimately mapped to id numValues". By construction that looks unreachable: ids are handed out as the map size at insert time and _idToValueMap grows in lockstep, so after n inserts the ids in the map are exactly {0..n-1} while numValues == size() == n. An existing key therefore always satisfies id < numValues. There is no concurrency either — one ValueToIdMap per key column per generator, single-threaded opchain.
putIfAbsent + INVALID_KEY is clearer and drops a capturing lambda, which is reason enough to make the change. But presenting it as a latent bug fix will send reviewers hunting for a bug that isn't there.
2. ObjectToIdMap has the identical pattern and was not updated.
Same package, same shape: computeIntIfAbsent(value, k -> numValues) followed by if (id == numValues). That is the implementation ValueToIdMapFactory returns for STRING, BYTES and BIG_DECIMAL — i.e. the most common group-by key type in practice. Whatever justifies the change here applies there at least as strongly. Could it be included, or the omission explained?
| int numValues = _valueToIdMap.size(); | ||
| int id = _valueToIdMap.computeIfAbsent(value, k -> numValues); | ||
| if (id == numValues) { | ||
| int id = _valueToIdMap.putIfAbsent(value, numValues); |
There was a problem hiding this comment.
FYI for anyone worried about the 912 -> 945 us in the description's idMapPutInt row: I checked, and the swap is not a slowdown.
I was suspicious because Int2IntMap.putIfAbsent has a default method that does get(), then containsKey(), then put() — two or three probes, where Int2IntOpenHashMap.computeIfAbsent is one. But the open-hash-map class overrides it. Disassembling fastutil 8.5.15:
public int putIfAbsent(int, int);
invokespecial find:(I)I
iflt -> insert(-pos-1, k, v); return defRetValue
else -> return value[pos]
One find, same as computeIfAbsent. The Float/Double/Long variants are generated from the same template.
So this change is allocation-neutral and cost-neutral, and the ~3.6% delta is run-to-run noise rather than a regression. Worth saying so in the description so nobody blocks on that table row.
| switch (storedType) { | ||
| case INT: | ||
| return ((Number) value).intValue(); | ||
| return value instanceof Integer ? value : ((Number) value).intValue(); |
There was a problem hiding this comment.
This is the strongest change in the PR, and its biggest beneficiary is not group-by merging — the title undersells the blast radius.
TypeUtils.convert has six callers. Two of them, LeafOperator:696 and LeafOperator:733, run it per row per column on every single-stage → multi-stage boundary crossing. That is a far hotter path than MultistageGroupByExecutor:269, which runs once per output group. So this helps every MSE query with a leaf stage. Widening the title or at least the description would help whoever bisects a behavior change here later.
Correctness looks fine: the four cases now preserve identity for an already-correctly-typed box, all four wrapper types are immutable so aliasing is unobservable except via ==, and convertRow mutates the row array rather than the values. NPE behavior on a null value is unchanged (the instanceof is false, the cast to Number then NPEs) and the new test pins that.
Non-blocking design thought: the fact that nearly all values on the INT/DOUBLE path are already the right type suggests the conversion call itself is the waste, not the boxing. A cheaper shape would be for convertRow to precompute per column whether any conversion is possible at all and skip the switch entirely — which would also drop the switch dispatch this keeps.
| } | ||
|
|
||
| @Test(dataProvider = "mergeModes") | ||
| public void testSerializedKeysAcrossBlocks(int numKeys, boolean leafReturnFinalResult) { |
There was a problem hiding this comment.
Good test, and aimed at the right surface: numKeys=2 routes to TwoKeysGroupIdGenerator and numKeys=3 to MultiKeysGroupIdGenerator, both over a serialized block, which is exactly the path the scratch-array reuse touches. If a generator ever started retaining the key array, every row after the first would collapse into a single group and this fails loudly. The multi-block, empty-block and null/NaN/-0.0 mix is a nice touch.
One gap: newExecutor always passes new int[]{-1} for filterArgIds and -1 for maxFilterArgId, so processAggregateWithFilter and the filtered generateGroupByKeys(DataBlock, numMatchedRows, matchedBitmap) are never reached. That is consistent with those staying untouched — but if you take the suggestion to optimize the filtered path too, this test will need a filtered variant.
| Float.MAX_VALUE, Float.POSITIVE_INFINITY, Float.NaN}}, | ||
| {new DoubleToIdMap(), new Object[]{Double.NEGATIVE_INFINITY, -Double.MAX_VALUE, -0.0d, 0.0d, Double.MIN_VALUE, | ||
| Double.MAX_VALUE, Double.POSITIVE_INFINITY, Double.NaN}}, | ||
| {new FloatToIdMap(), new Object[]{Float.intBitsToFloat(0x7fc00001)}}, |
There was a problem hiding this comment.
These two rows look like they pin NaN canonicalization, but as written they don't: each is a one-element map, so the only assertion is "the first put returns 0".
Float.intBitsToFloat(0x7fc00001) and Double.longBitsToDouble(0x7ff8000000000001L) are non-canonical NaN bit patterns. The interesting question they gesture at is whether the map treats them as the same key as Float.NaN/Double.NaN — i.e. whether fastutil keys on floatToIntBits (canonicalizing, matching Float.equals) or floatToRawIntBits (not canonicalizing). Because each lives in its own data-provider row with its own fresh map, the two NaNs never coexist and nothing is compared.
If pinning that is the intent, put the canonical and non-canonical NaN in the same values array and assert whichever behavior HEAD actually has. If it isn't the intent, I'd drop the two rows — they add runtime and a false sense of coverage.
(The -0.0d / 0.0d pair in the row above is meaningful, since those share a map.)
| } | ||
|
|
||
| @Test(dataProvider = "numericMaps") | ||
| public void testNumericKeys(ValueToIdMap map, Object[] values) { |
There was a problem hiding this comment.
The class is named NumericToIdMapTest, but ObjectToIdMap gets no coverage — and that is both the impl backing STRING/BYTES/BIG_DECIMAL keys and the one left on the old computeIntIfAbsent pattern.
The invariants this test checks — contiguous ids, stable id on repeat put, get(id) round-trip, behavior across map growth — apply to it unchanged. Adding a row is nearly free and would give the one untouched implementation the same regression net, whether or not you also switch it to putIfAbsent.
Minor: the if (map instanceof IntToIdMap) ... else if ... chain in the growth loop is a little brittle. Passing a boxing function alongside the map in the data provider would read better and scale as types are added.
PR flow
Reuse key array and intermediate arrays; optimize ID maps per type and TypeUtils to cut allocations.
AI-generated · Green: added · Yellow: modified · Red: removed · Gray: existing
Diff evidence
Summary
Serialized multi-stage GROUP BY merging (
MultistageGroupByExecutor.processMerge) allocated a separate composite-key array for every input row, eagerly allocated intermediate-result inner arrays that were immediately overwritten, and created temporary numeric wrapper objects during dictionary insertion and result conversion. These allocations grow with partial-row and group counts.This PR makes four independent changes:
generateGroupByKeys, multi-column path). Extract each key column once with the existingextractKeyhelper, then fill a single reusedObject[numKeys]per row instead of allocating a fresh array per row viaextractKeys. Safe because the composite-key group generators copy out dictionary IDs and do not retain the passed array (same contract as the single-value row-heap path).processMerge). Changenew Object[numFunctions][numRows]tonew Object[numFunctions][]; the inner arrays were pure garbage — overwritten on the next line bygetIntermediateResults(...).putIfAbsentin the numeric ID maps (IntToIdMap/LongToIdMap/FloatToIdMap/DoubleToIdMap). ReplacecomputeIfAbsent(v, k -> numValues)(a capturing lambda) withputIfAbsent(v, numValues), and switch the "was-absent" test from== numValuesto theINVALID_KEYsentinel — the old test could falsely report "absent" when an existing key legitimately mapped to idnumValues.TypeUtils.convert). Return the existing wrapper (value instanceof Integer ? value : ...) for INT/LONG/FLOAT/DOUBLE instead of always unboxing and re-boxing.Null handling, numeric conversions, floating-point key behavior and group limits are unchanged.
Per-change performance attribution
Measured with a focused JMH benchmark (
BenchmarkGroupByMergeAlloc, not committed) over 100K values, 10K-cardinality, JDK 25, 2 forks × 8 iterations,-prof gc. Same benchmark bytecode linked against base (3b6be9c) vs this PR (20357e2):TypeUtils.convertconvertIntSameTypeTypeUtils.convertconvertDoubleSameType*ToIdMap.putidMapPutInt*ToIdMap.putidMapPutLongInterpretation of each change's contribution:
TypeUtils.convert) is the dominant allocation win. For already-correctly-typed values it eliminates essentially all re-boxing garbage (~1.5 MB per 100K INT, ~2.4 MB per 100K DOUBLE → ~0) and is ~30% faster on that path. This is the change most directly responsible for the merge-level allocation reduction.*ToIdMap) contributes correctness and dispatch cleanup, not measurable heap savings. In the microbenchmark the allocation is dominated by fastutil map/array growth as the dictionary fills and is unchanged; the removed capturing lambda is JIT-hoisted or negligible here. Its concrete value is theINVALID_KEYsentinel fix (avoids a false-absent when a key maps to idnumValues) plus avoiding lambda dispatch on the hot insert path.Object[]is scalarized by escape analysis (both patterns measured ~0 B/op), so a microbenchmark cannot reproduce the win. In production,keyescapes intogetGroupId(...)across a call boundary escape analysis cannot defeat, so the per-row array does allocate and reuse removes N−1 of N such arrays. Measured only via the real-executor probe below.numFunctionswastedObject[numRows]allocations per merged block; small and constant per block, folded into the executor-level probe.Validation
NumericToIdMapTest,TypeUtilsTest,MultistageGroupByExecutorTest, andAggregateOperatorTest.No public API, wire format, configuration default or permanent benchmark is added.