xiangfu0 commented on code in PR #19825:
URL: https://github.com/apache/pinot/pull/19825#discussion_r4252839055
##########
pinot-core/src/main/java/org/apache/pinot/core/data/table/IndexedTable.java:
##########
@@ -108,14 +109,57 @@ public boolean upsert(Record record) {
return upsert(new Key(keyValues), record);
}
+ /// Upserts an already-extracted result into a table with exactly one
aggregation function.
+ /// The table may retain the key and result; callers must not modify them
after this call.
+ /// A null result follows the aggregation function's existing merge
semantics.
+ /// Final inputs must be Comparable, as required by the aggregation
function's final-result merge.
+ /// Returns true even when new groups are rejected; existing groups can
still accept updates.
+ public boolean upsertSingleAggregationResult(Key key, @Nullable Object
result) {
+ Preconditions.checkState(_aggregationFunctions.length == 1,
+ "upsertSingleAggregationResult requires exactly one aggregation
function");
+ Preconditions.checkArgument(!_hasFinalInput || result == null || result
instanceof Comparable,
+ "Final aggregation inputs must be Comparable");
+ if (_hasOrderBy) {
+ return upsert(key, createSingleAggregationRecord(key, result));
+ }
+ upsertWithoutOrderBy(key, null, result);
+ return true;
+ }
+
+ /// Applies the table's unordered admission policy to either an eager record
or a single aggregation result.
+ /// A null record selects the single-result input, which may itself be null.
+ protected abstract void upsertWithoutOrderBy(Key key, @Nullable Record
record, @Nullable Object result);
+
+ protected Record createSingleAggregationRecord(Key key, @Nullable Object
result) {
+ Object[] values = Arrays.copyOf(key.getValues(), _numKeyColumns + 1);
+ values[_numKeyColumns] = result;
+ return new Record(values);
+ }
+
+ protected Record updateSingleAggregationRecord(Record record, @Nullable
Object result) {
+ Object[] values = record.getValues();
+ values[_numKeyColumns] = mergeValue(0, values[_numKeyColumns], result,
_hasFinalInput);
+ return record;
+ }
+
/// Adds a record with new key or updates a record with existing key.
protected void addOrUpdateRecord(Key key, Record newRecord) {
- _lookupMap.compute(key, (k, v) -> v == null ? newRecord : updateRecord(v,
newRecord));
+ addOrUpdateRecord(key, newRecord, null);
+ }
+
+ protected void addOrUpdateRecord(Key key, @Nullable Record newRecord,
@Nullable Object result) {
+ _lookupMap.compute(key, (k, existing) -> {
+ if (existing == null) {
+ return newRecord != null ? newRecord :
createSingleAggregationRecord(k, result);
+ }
+ return newRecord != null ? updateRecord(existing, newRecord) :
updateSingleAggregationRecord(existing, result);
+ });
}
/// Updates a record with existing key. Record with new key will be ignored.
- protected void updateExistingRecord(Key key, Record newRecord) {
- _lookupMap.computeIfPresent(key, (k, v) -> updateRecord(v, newRecord));
+ protected void updateExistingRecord(Key key, @Nullable Record newRecord,
@Nullable Object result) {
+ _lookupMap.computeIfPresent(key, (k, existing) -> newRecord != null
Review Comment:
Added the rationale beside computeIfPresent: it must remain unconditional to
wait for a pending insertion instead of discarding its update. Applied in
[954b068598](https://github.com/apache/pinot/commit/954b068598c0bb3b7b3a9a240be81366db3dde77).
##########
pinot-core/src/test/java/org/apache/pinot/core/data/table/UnboundedConcurrentIndexedTableTest.java:
##########
@@ -0,0 +1,476 @@
+/**
+ * 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.data.table;
+
+import com.clearspring.analytics.stream.cardinality.HyperLogLog;
+import it.unimi.dsi.fastutil.ints.IntArrayList;
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.segment.local.customobject.AvgPair;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+
+
+/// Verifies unordered group insertion and atomic updates using independent
concurrent writers.
+public class UnboundedConcurrentIndexedTableTest {
+ @DataProvider(name = "inputModes")
+ public Object[][] inputModes() {
+ return new Object[][]{{false, false}, {true, false}, {false, true}, {true,
true}};
+ }
+
+ @DataProvider(name = "concurrentInputModes")
+ public Object[][] concurrentInputModes() {
+ return new Object[][]{{false, false, false}, {true, false, false}, {false,
true, false}, {true, true, false},
+ {false, false, true}, {true, false, true}, {false, true, true}, {true,
true, true}};
+ }
+
+ @DataProvider(name = "scalarInputs")
+ public Object[][] scalarInputs() {
+ return new Object[][]{{false}, {true}};
+ }
+
+ @DataProvider(name = "singleAggregations")
+ public Object[][] singleAggregations() {
+ Supplier<Object[]> numericInputs = () -> new Object[]{3d, 5d};
+ Supplier<Object[]> avgInputs = () -> new Object[]{new AvgPair(9, 3), new
AvgPair(11, 1)};
+ Supplier<Object[]> distinctInputs = () -> new Object[]{new
IntOpenHashSet(new int[]{1, 2}),
+ new IntOpenHashSet(new int[]{2, 3})};
+ Supplier<Object[]> finalDistinctInputs = () -> new Object[]{3, 5};
+ Supplier<Object[]> hllInputs = () -> {
+ HyperLogLog first = new HyperLogLog(8);
+ first.offer(1);
+ first.offer(2);
+ HyperLogLog second = new HyperLogLog(8);
+ second.offer(2);
+ second.offer(3);
+ return new Object[]{first, second};
+ };
+ Supplier<Object[]> arrayInputs = () -> new Object[]{new IntArrayList(new
int[]{1, 2}),
+ new IntArrayList(new int[]{2, 3})};
+ return new Object[][]{
+ {"SUM(metric)", false, numericInputs, 8d, null},
+ {"SUM(metric)", true, numericInputs, 8d, null},
+ {"MIN(metric)", false, numericInputs, 3d, null},
+ {"MIN(metric)", true, numericInputs, 3d, null},
+ {"MAX(metric)", false, numericInputs, 5d, null},
+ {"MAX(metric)", true, numericInputs, 5d, null},
+ {"AVG(metric)", false, avgInputs, 5d, null},
+ {"DISTINCTCOUNT(metric)", false, distinctInputs, 3, 0},
+ {"DISTINCTCOUNT(metric)", true, finalDistinctInputs, 8, null},
+ {"DISTINCTCOUNTHLL(metric)", false, hllInputs, 3L, 0L},
+ {"ARRAYAGG(metric, 'INT')", false, arrayInputs, new IntArrayList(new
int[]{1, 2, 2, 3}), new IntArrayList()}
+ };
+ }
+
+ @Test(dataProvider = "singleAggregations")
+ public void testSingleAggregationParity(String expression, boolean
hasFinalInput, Supplier<Object[]> inputSupplier,
+ Object mergedResult, Object emptyResult)
+ throws IOException {
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ IndexedTable eagerTable = singleAggregationTable(expression,
hasFinalInput, inputSupplier, false, executor);
+ IndexedTable scalarTable = singleAggregationTable(expression,
hasFinalInput, inputSupplier, true, executor);
+ assertEquals(scalarTable._lookupMap.keySet(),
eagerTable._lookupMap.keySet());
+ if (!hasFinalInput) {
+ for (Map.Entry<Key, Record> entry : eagerTable._lookupMap.entrySet()) {
+
assertEquivalentResult(scalarTable._lookupMap.get(entry.getKey()).getValues()[2],
+ entry.getValue().getValues()[2], expression + ": intermediate "
+ entry.getKey());
+ }
+ }
+ Map<Key, Object> eagerResults = finishSingleAggregation(eagerTable);
+ Map<Key, Object> scalarResults = finishSingleAggregation(scalarTable);
+ assertEquals(eagerResults.size(), 3);
+ assertEquals(scalarResults.keySet(), eagerResults.keySet());
+ for (String group : List.of("null-first", "null-last", "all-null")) {
+ Key key = new Key(new Object[]{group, 1d});
+ assertTrue(eagerResults.containsKey(key));
+ Object expectedResult = group.equals("all-null") ? emptyResult :
mergedResult;
+ assertTrue(Objects.deepEquals(eagerResults.get(key), expectedResult),
expression + ": " + group);
+ assertTrue(Objects.deepEquals(scalarResults.get(key),
eagerResults.get(key)), expression + ": " + group);
+ }
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ @Test(dataProvider = "concurrentInputModes", timeOut = 30_000)
+ public void testConcurrentCountsWithHashCollisions(boolean scalarInput,
boolean hasFinalInput,
+ boolean deterministicTable)
+ throws InterruptedException, ExecutionException, TimeoutException {
+ int numThreads = 4;
+ int iterations = 1_000;
+ assertEquals(new Key(new Object[]{"Aa", 1d}).hashCode(), new Key(new
Object[]{"BB", 1d}).hashCode());
+ ExecutorService executor = Executors.newFixedThreadPool(numThreads);
+ CountDownLatch ready = new CountDownLatch(numThreads);
+ CountDownLatch start = new CountDownLatch(1);
+ try {
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(
+ "SELECT COUNT(*) FROM testTable GROUP BY tv_id, weights");
+ IndexedTable table = deterministicTable
+ ? new DeterministicConcurrentIndexedTable(newDataSchema(),
hasFinalInput, queryContext, Integer.MAX_VALUE,
+ Integer.MAX_VALUE, Integer.MAX_VALUE, 1, executor)
+ : newTable(Integer.MAX_VALUE, hasFinalInput, executor);
+ List<Future<?>> futures = new ArrayList<>();
+ for (int threadIndex = 0; threadIndex < numThreads; threadIndex++) {
+ String distinctKey = "thread-" + threadIndex;
+ futures.add(executor.submit(() -> {
+ ready.countDown();
+ start.await();
+ for (int i = 0; i < iterations; i++) {
+ upsert(table, new Record(new Object[]{"Aa", 1d, 1L}), scalarInput);
+ upsert(table, new Record(new Object[]{"BB", 1d, 1L}), scalarInput);
+ upsert(table, new Record(new Object[]{distinctKey, 1d, 1L}),
scalarInput);
+ }
+ return null;
+ }));
+ }
+ assertTrue(ready.await(10, TimeUnit.SECONDS), "Writers did not become
ready");
+ start.countDown();
+ for (Future<?> future : futures) {
+ future.get(10, TimeUnit.SECONDS);
+ }
+ Map<Key, Long> counts = finishAndCount(table, hasFinalInput);
+ assertEquals(counts.size(), numThreads + 2);
+ assertEquals(counts.get(new Key(new Object[]{"Aa", 1d})),
Long.valueOf(numThreads * iterations));
+ assertEquals(counts.get(new Key(new Object[]{"BB", 1d})),
Long.valueOf(numThreads * iterations));
+ for (int threadIndex = 0; threadIndex < numThreads; threadIndex++) {
+ assertEquals(counts.get(new Key(new Object[]{"thread-" + threadIndex,
1d})), Long.valueOf(iterations));
+ }
+ } finally {
+ start.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test(dataProvider = "scalarInputs", timeOut = 30_000)
+ public void testConcurrentSaturatedMutableUpdates(boolean scalarInput)
+ throws InterruptedException, ExecutionException, TimeoutException {
+ int numThreads = 4;
+ int iterations = 1_000;
+ ExecutorService executor = Executors.newFixedThreadPool(numThreads);
+ CountDownLatch ready = new CountDownLatch(numThreads);
+ CountDownLatch start = new CountDownLatch(1);
+ try {
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(
+ "SELECT AVG(metric) FROM testTable GROUP BY tv_id, weights");
+ DataSchema dataSchema = new DataSchema(new String[]{"tv_id", "weights",
"avg(metric)"},
+ new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.DOUBLE,
ColumnDataType.OBJECT});
+ IndexedTable table = new UnboundedConcurrentIndexedTable(dataSchema,
false, queryContext, 1, 1, executor);
+ Key hotKey = new Key(new Object[]{"hot", 1d});
+ upsert(table, new Record(new Object[]{"hot", 1d, new AvgPair(9, 3)}),
scalarInput);
+ List<Future<?>> futures = new ArrayList<>();
+ for (int threadIndex = 0; threadIndex < numThreads; threadIndex++) {
+ String rejectedKey = "rejected-" + threadIndex;
+ futures.add(executor.submit(() -> {
+ ready.countDown();
+ start.await();
+ for (int i = 0; i < iterations; i++) {
+ upsert(table, new Record(new Object[]{"hot", 1d, new AvgPair(3,
1)}), scalarInput);
+ upsert(table, new Record(new Object[]{rejectedKey, 1d, new
AvgPair(100, 1)}), scalarInput);
+ }
+ return null;
+ }));
+ }
+ assertTrue(ready.await(10, TimeUnit.SECONDS), "Writers did not become
ready");
+ start.countDown();
+ for (Future<?> future : futures) {
+ future.get(10, TimeUnit.SECONDS);
+ }
+ assertEquals(table._lookupMap.keySet(), Set.of(hotKey));
+ AvgPair result = (AvgPair) table._lookupMap.get(hotKey).getValues()[2];
+ assertEquals(result.getCount(), 3L + numThreads * iterations);
+ assertEquals(result.getSum(), 9d + 3d * numThreads * iterations);
+ assertEquals(finishSingleAggregation(table), Map.of(hotKey, 3d));
+ } finally {
+ start.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test(dataProvider = "scalarInputs", timeOut = 30_000)
+ public void testPendingInsertionUpdateAfterGroupLimit(boolean scalarInput)
+ throws InterruptedException, ExecutionException, TimeoutException {
+ ExecutorService executor = Executors.newFixedThreadPool(3);
+ CountDownLatch insertionStarted = new CountDownLatch(1);
+ CountDownLatch releaseInsertion = new CountDownLatch(1);
+ CountDownLatch updateReady = new CountDownLatch(1);
+ AtomicReference<Thread> updateThread = new AtomicReference<>();
+ Key hotKey = new Key(new Object[]{"hot", 1d});
+ Key otherKey = new Key(new Object[]{"other-key", 1d});
+ int hotHash = hotKey.hashCode();
+ int otherHash = otherKey.hashCode();
+ // Different low spread-hash bits put the keys in separate buckets at
every map capacity above one.
+ assertTrue(((hotHash ^ (hotHash >>> 16)) & 1) != ((otherHash ^ (otherHash
>>> 16)) & 1));
+ try {
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(
+ "SELECT COUNT(*) FROM testTable GROUP BY tv_id, weights");
+ IndexedTable table = new
UnboundedConcurrentIndexedTable(newDataSchema(), false, queryContext, 1, 64,
executor) {
+ @Override
+ protected Record createSingleAggregationRecord(Key key, Object result)
{
+ if (key.equals(hotKey)) {
+ insertionStarted.countDown();
+ try {
+ assertTrue(releaseInsertion.await(10, TimeUnit.SECONDS),
"Pending insertion was not released");
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException(e);
+ }
+ }
+ return super.createSingleAggregationRecord(key, result);
+ }
+ };
+ Future<?> insertion = executor.submit(() ->
table.upsertSingleAggregationResult(hotKey, 1L));
+ assertTrue(insertionStarted.await(10, TimeUnit.SECONDS), "Insertion did
not reach the map computation");
+ Future<?> otherInsertion = executor.submit(() ->
table.upsertSingleAggregationResult(otherKey, 1L));
+ otherInsertion.get(10, TimeUnit.SECONDS);
+ assertEquals(table.size(), 1);
+ Future<?> update = executor.submit(() -> {
+ updateThread.set(Thread.currentThread());
+ updateReady.countDown();
+ upsert(table, new Record(new Object[]{"hot", 1d, 1L}), scalarInput);
+ });
+ assertTrue(updateReady.await(10, TimeUnit.SECONDS), "Update writer did
not become ready");
+ long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10);
+ while (!update.isDone() && updateThread.get().getState() !=
Thread.State.BLOCKED
+ && System.nanoTime() < deadline) {
+ Thread.onSpinWait();
Review Comment:
Replaced the deadline and busy-wait loop with TestUtils.waitForCondition,
using the same completed-or-BLOCKED predicate and 10-second bound. Both
controlled eager/scalar regression cases pass. Applied in
[954b068598](https://github.com/apache/pinot/commit/954b068598c0bb3b7b3a9a240be81366db3dde77).
##########
pinot-core/src/main/java/org/apache/pinot/core/data/table/IndexedTable.java:
##########
@@ -108,14 +109,57 @@ public boolean upsert(Record record) {
return upsert(new Key(keyValues), record);
}
+ /// Upserts an already-extracted result into a table with exactly one
aggregation function.
+ /// The table may retain the key and result; callers must not modify them
after this call.
+ /// A null result follows the aggregation function's existing merge
semantics.
+ /// Final inputs must be Comparable, as required by the aggregation
function's final-result merge.
+ /// Returns true even when new groups are rejected; existing groups can
still accept updates.
+ public boolean upsertSingleAggregationResult(Key key, @Nullable Object
result) {
+ Preconditions.checkState(_aggregationFunctions.length == 1,
+ "upsertSingleAggregationResult requires exactly one aggregation
function");
+ Preconditions.checkArgument(!_hasFinalInput || result == null || result
instanceof Comparable,
+ "Final aggregation inputs must be Comparable");
+ if (_hasOrderBy) {
+ return upsert(key, createSingleAggregationRecord(key, result));
+ }
+ upsertWithoutOrderBy(key, null, result);
+ return true;
+ }
+
+ /// Applies the table's unordered admission policy to either an eager record
or a single aggregation result.
+ /// A null record selects the single-result input, which may itself be null.
+ protected abstract void upsertWithoutOrderBy(Key key, @Nullable Record
record, @Nullable Object result);
+
+ protected Record createSingleAggregationRecord(Key key, @Nullable Object
result) {
+ Object[] values = Arrays.copyOf(key.getValues(), _numKeyColumns + 1);
+ values[_numKeyColumns] = result;
+ return new Record(values);
+ }
+
+ protected Record updateSingleAggregationRecord(Record record, @Nullable
Object result) {
Review Comment:
Made updateSingleAggregationRecord private after confirming there are no
subclass callers or overrides, and annotated mergeValue’s nullable return.
Applied in
[954b068598](https://github.com/apache/pinot/commit/954b068598c0bb3b7b3a9a240be81366db3dde77).
--
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]
