gortiz commented on code in PR #19601:
URL: https://github.com/apache/pinot/pull/19601#discussion_r4062273235
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java:
##########
@@ -446,16 +446,27 @@ private int[] generateGroupByKeys(List<Object[]> rows) {
}
private int[] generateGroupByKeys(DataBlock dataBlock) {
- Object[] keys;
- if (_groupKeyIds.length == 1) {
- keys = DataBlockExtractUtils.extractKey(dataBlock, _groupKeyIds[0]);
- } else {
- keys = DataBlockExtractUtils.extractKeys(dataBlock, _groupKeyIds);
- }
- int numRows = keys.length;
+ int numRows = dataBlock.getNumberOfRows();
int[] intKeys = new int[numRows];
- for (int i = 0; i < numRows; i++) {
- intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);
+ int numKeys = _groupKeyIds.length;
+ if (numKeys == 1) {
+ Object[] keys = DataBlockExtractUtils.extractKey(dataBlock,
_groupKeyIds[0]);
+ for (int i = 0; i < numRows; i++) {
+ intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);
+ }
+ } else {
+ Object[][] columns = new Object[numKeys][];
Review 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.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java:
##########
@@ -446,16 +446,27 @@ private int[] generateGroupByKeys(List<Object[]> rows) {
}
private int[] generateGroupByKeys(DataBlock dataBlock) {
- Object[] keys;
- if (_groupKeyIds.length == 1) {
- keys = DataBlockExtractUtils.extractKey(dataBlock, _groupKeyIds[0]);
- } else {
- keys = DataBlockExtractUtils.extractKeys(dataBlock, _groupKeyIds);
- }
- int numRows = keys.length;
+ int numRows = dataBlock.getNumberOfRows();
int[] intKeys = new int[numRows];
- for (int i = 0; i < numRows; i++) {
- intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);
+ int numKeys = _groupKeyIds.length;
+ if (numKeys == 1) {
+ Object[] keys = DataBlockExtractUtils.extractKey(dataBlock,
_groupKeyIds[0]);
+ for (int i = 0; i < numRows; i++) {
+ intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);
+ }
+ } else {
+ Object[][] columns = new Object[numKeys][];
+ for (int i = 0; i < numKeys; i++) {
+ 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];
+ for (int rowId = 0; rowId < numRows; rowId++) {
+ for (int i = 0; i < numKeys; i++) {
+ key[i] = columns[i][rowId];
+ }
+ intKeys[rowId] = _groupIdGenerator.getGroupId(key);
Review 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:
```java
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.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java:
##########
@@ -372,7 +372,7 @@ private void processMerge(MseBlock.Data block) {
int[] groupByKeys = generateGroupByKeys(block);
int numRows = groupByKeys.length;
int numFunctions = _aggFunctions.length;
- Object[][] intermediateResults = new Object[numFunctions][numRows];
+ Object[][] intermediateResults = new Object[numFunctions][];
Review 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.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/IntToIdMap.java:
##########
@@ -36,9 +36,10 @@ public IntToIdMap() {
@Override
public int put(int value) {
int numValues = _valueToIdMap.size();
- int id = _valueToIdMap.computeIfAbsent(value, k -> numValues);
- if (id == numValues) {
+ int id = _valueToIdMap.putIfAbsent(value, numValues);
+ if (id == INVALID_KEY) {
Review 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?
##########
pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/utils/NumericToIdMapTest.java:
##########
@@ -0,0 +1,76 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.query.aggregation.groupby.utils;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+
+/// Verifies numeric key identity and contiguous IDs across insertion,
repetition, and map growth.
+public class NumericToIdMapTest {
+ @DataProvider(name = "numericMaps")
+ public Object[][] numericMaps() {
+ return new Object[][]{
+ {new IntToIdMap(), new Object[]{Integer.MIN_VALUE, -1, 0, 1,
Integer.MAX_VALUE}},
+ {new LongToIdMap(), new Object[]{Long.MIN_VALUE, -1L, 0L, 1L,
Long.MAX_VALUE, 9007199254740993L}},
+ {new FloatToIdMap(), new Object[]{Float.NEGATIVE_INFINITY,
-Float.MAX_VALUE, -0.0f, 0.0f, Float.MIN_VALUE,
+ 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)}},
+ {new DoubleToIdMap(), new
Object[]{Double.longBitsToDouble(0x7ff8000000000001L)}}
+ };
+ }
+
+ @Test(dataProvider = "numericMaps")
+ public void testNumericKeys(ValueToIdMap map, Object[] values) {
Review 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.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java:
##########
@@ -446,16 +446,27 @@ private int[] generateGroupByKeys(List<Object[]> rows) {
}
private int[] generateGroupByKeys(DataBlock dataBlock) {
- Object[] keys;
- if (_groupKeyIds.length == 1) {
- keys = DataBlockExtractUtils.extractKey(dataBlock, _groupKeyIds[0]);
- } else {
- keys = DataBlockExtractUtils.extractKeys(dataBlock, _groupKeyIds);
- }
- int numRows = keys.length;
+ int numRows = dataBlock.getNumberOfRows();
int[] intKeys = new int[numRows];
- for (int i = 0; i < numRows; i++) {
- intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);
+ int numKeys = _groupKeyIds.length;
+ if (numKeys == 1) {
+ Object[] keys = DataBlockExtractUtils.extractKey(dataBlock,
_groupKeyIds[0]);
+ for (int i = 0; i < numRows; i++) {
+ intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);
+ }
+ } else {
+ Object[][] columns = new Object[numKeys][];
+ for (int i = 0; i < numKeys; i++) {
+ 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];
Review 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.getGroupId` reads `keyValues[0]`/`[1]`, maps them
to ints and packs them into a `long` before touching `_groupIdMap`. Nothing
keeps the array.
* `MultiKeysGroupIdGenerator.getGroupId` copies into a fresh `int[] keyIds`
and stores `new 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.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/DoubleToIdMap.java:
##########
@@ -36,9 +36,10 @@ public DoubleToIdMap() {
@Override
public int put(double value) {
int numValues = _valueToIdMap.size();
- int id = _valueToIdMap.computeIfAbsent(value, k -> numValues);
- if (id == numValues) {
+ int id = _valueToIdMap.putIfAbsent(value, numValues);
Review 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.
##########
pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/utils/NumericToIdMapTest.java:
##########
@@ -0,0 +1,76 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.query.aggregation.groupby.utils;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+
+/// Verifies numeric key identity and contiguous IDs across insertion,
repetition, and map growth.
+public class NumericToIdMapTest {
+ @DataProvider(name = "numericMaps")
+ public Object[][] numericMaps() {
+ return new Object[][]{
+ {new IntToIdMap(), new Object[]{Integer.MIN_VALUE, -1, 0, 1,
Integer.MAX_VALUE}},
+ {new LongToIdMap(), new Object[]{Long.MIN_VALUE, -1L, 0L, 1L,
Long.MAX_VALUE, 9007199254740993L}},
+ {new FloatToIdMap(), new Object[]{Float.NEGATIVE_INFINITY,
-Float.MAX_VALUE, -0.0f, 0.0f, Float.MIN_VALUE,
+ 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)}},
Review 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.)
##########
pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutorTest.java:
##########
@@ -0,0 +1,130 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.runtime.operator;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+import
org.apache.pinot.core.query.aggregation.function.CountAggregationFunction;
+import org.apache.pinot.query.planner.plannode.AggregateNode.AggType;
+import org.apache.pinot.query.runtime.blocks.MseBlock;
+import
org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+
+
+/// Exercises serialized composite keys and merge state across blocks. Each
test owns its executor and input blocks.
+public class MultistageGroupByExecutorTest {
+ private static final DataSchema INPUT_SCHEMA = new DataSchema(new
String[]{"weight", "count", "tv", "tag"},
+ new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.LONG,
ColumnDataType.INT, ColumnDataType.STRING});
+
+ @DataProvider
+ public Object[][] mergeModes() {
+ return new Object[][]{{2, false}, {2, true}, {3, false}, {3, true}};
+ }
+
+ @Test(dataProvider = "mergeModes")
+ public void testSerializedKeysAcrossBlocks(int numKeys, boolean
leafReturnFinalResult) {
Review 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.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/TypeUtils.java:
##########
@@ -39,13 +39,13 @@ private TypeUtils() {
public static Object convert(Object value, ColumnDataType storedType) {
switch (storedType) {
case INT:
- return ((Number) value).intValue();
+ return value instanceof Integer ? value : ((Number) value).intValue();
Review 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.
--
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]
