Jackie-Jiang commented on code in PR #19825:
URL: https://github.com/apache/pinot/pull/19825#discussion_r4252759255
##########
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:
Optional cleanup: add a short comment explaining why this computeIfPresent
must remain unconditional. A preliminary containsKey lookup can miss an
insertion in progress, whereas computeIfPresent waits for its publication.
Keeping that rationale beside the code would help prevent a future optimization
from reintroducing the race.
##########
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:
Optional cleanup: make updateSingleAggregationRecord private, since it has
no subclass callers or overrides. Also annotate mergeValue's return with
@Nullable: the delegated merge helpers can return null. This keeps the new
helper surface narrow and makes the null contract explicit.
##########
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:
Optional cleanup: use TestUtils.waitForCondition with the same
update.isDone() || updateThread.get().getState() == Thread.State.BLOCKED
predicate. The repository already uses this pattern, including
RealtimeTableDataManagerTest. It removes the hand-written deadline and
busy-wait loop while preserving the controlled interleaving and final count
assertion.
--
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]
