From 0b5e0eda04f7e6521cf338bee3cc620368b491f9 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 28 Sep 2026 15:50:46 -0700 Subject: [PATCH 1/9] Add per-exchange AUTO selection for ordered window sorting --- .../requesthandler/WindowSortAutoTuner.java | 317 ++++++++++++++++++ .../WindowSortAutoTunerTest.java | 293 ++++++++++++++++ .../utils/config/QueryOptionsUtils.java | 5 - pinot-common/src/main/proto/plan.proto | 1 + .../utils/config/QueryOptionsUtilsTest.java | 8 - .../core/util/DataBlockExtractUtils.java | 8 +- .../rel/logical/PinotLogicalSortExchange.java | 30 +- .../planner/logical/WindowSortAutoPlan.java | 34 ++ .../query/planner/plannode/ExchangeNode.java | 20 +- .../planner/plannode/MailboxReceiveNode.java | 22 +- .../planner/serde/PlanNodeDeserializer.java | 2 +- .../planner/serde/PlanNodeSerializer.java | 1 + .../operator/BaseMailboxReceiveOperator.java | 8 +- .../operator/MailboxReceiveOperator.java | 84 +++++ .../operator/MailboxReceiveOperatorTest.java | 92 +++++ 15 files changed, 901 insertions(+), 24 deletions(-) create mode 100644 pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java create mode 100644 pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java create mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlan.java diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java new file mode 100644 index 000000000000..97cdb5601010 --- /dev/null +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java @@ -0,0 +1,317 @@ +/** + * 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.broker.requesthandler; + +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; +import java.time.Duration; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import javax.annotation.Nullable; +import org.apache.pinot.common.datatable.StatMap; +import org.apache.pinot.query.planner.logical.WindowSortAutoPlan; +import org.apache.pinot.query.runtime.operator.BaseMailboxReceiveOperator; +import org.apache.pinot.query.runtime.operator.MultiStageOperator; +import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; + + +/// Learns when sender sorting is likely to beat the existing receiver sort for one ordered-window exchange. +/// The cache is bounded, expires observations, and stores only plan fingerprints, never SQL or result data. +/// A query with no sufficiently strong history uses the receiver-sort plan. Sessions are used once per query; +/// concurrent queries may safely read and update the shared cache. +public final class WindowSortAutoTuner { + private static final int MIN_SENDERS = 4; + private static final long MIN_ROWS = 400_000; + private static final long MIN_SAMPLED_ROWS = MIN_SENDERS * 64L; + private static final int REQUIRED_OBSERVATIONS = 2; + private static final int NEGATIVE_OBSERVATIONS = -1; + private static final int DISABLED_OBSERVATIONS = -2; + private static final int NEGATIVE_SKIP_QUERIES = 7; + private static final int RESAMPLE_EVERY = 8; + private static final int SLOWDOWN_PERCENT = 15; + + private final Cache _observations = CacheBuilder.newBuilder() + .maximumSize(1_024).expireAfterWrite(Duration.ofMinutes(30)).build(); + private final AtomicLong _nextGeneration = new AtomicLong(); + + public Session newSession(long queryHash) { + return new Session(queryHash); + } + + /// A single query's planner decisions and, on successful completion, its receiver-stage observations. + public final class Session implements WindowSortAutoPlan { + private final long _queryHash; + private final Map _decisions = new HashMap<>(); + private final AtomicBoolean _observed = new AtomicBoolean(); + + private Session(long queryHash) { + _queryHash = queryHash; + } + + @Override + public boolean useSenderSort(int receiverStageId, int senderStageId, int inputHash, int collationHash) { + return decisionFor(receiverStageId, senderStageId, inputHash, collationHash).senderSort(); + } + + @Override + public boolean shouldProfile(int receiverStageId, int senderStageId, int inputHash, int collationHash) { + return decisionFor(receiverStageId, senderStageId, inputHash, collationHash).profile(); + } + + private Decision decisionFor(int receiverStageId, int senderStageId, int inputHash, int collationHash) { + ExchangeKey key = new ExchangeKey(_queryHash, receiverStageId, senderStageId, inputHash, collationHash); + return _decisions.computeIfAbsent(key, ignored -> { + Evidence evidence = _observations.getIfPresent(key); + if (evidence == null) { + Evidence initial = new Evidence(0, 0, _nextGeneration.incrementAndGet(), 0); + Evidence existing = _observations.asMap().putIfAbsent(key, initial); + evidence = existing != null ? existing : initial; + } + if (evidence.pendingProbe() != 0) { + // A read does not refresh the TTL of an abandoned probe. + return new Decision(false, false, evidence.generation(), 0); + } + if (evidence.qualifyingObservations() == DISABLED_OBSERVATIONS) { + // Do not refresh the disabled entry's TTL; a new profile is possible only after it expires. + return new Decision(false, false, evidence.generation(), 0); + } + if (evidence.qualifyingObservations() >= 0 + && evidence.qualifyingObservations() < REQUIRED_OBSERVATIONS) { + return new Decision(false, true, evidence.generation(), 0); + } + AtomicReference selected = new AtomicReference<>(); + _observations.asMap().computeIfPresent(key, (unused, current) -> { + if (current.pendingProbe() != 0) { + selected.set(new Decision(false, false, current.generation(), 0)); + return current; + } + if (current.qualifyingObservations() == DISABLED_OBSERVATIONS) { + selected.set(new Decision(false, false, current.generation(), 0)); + return current; + } + if (current.qualifyingObservations() == NEGATIVE_OBSERVATIONS) { + if (current.queriesSinceSample() > 0) { + selected.set(new Decision(false, false, current.generation(), 0)); + return new Evidence(NEGATIVE_OBSERVATIONS, current.queriesSinceSample() - 1, + current.generation(), 0); + } + long generation = _nextGeneration.incrementAndGet(); + selected.set(new Decision(false, true, generation, generation)); + return new Evidence(NEGATIVE_OBSERVATIONS, 0, generation, generation); + } + if (current.qualifyingObservations() < REQUIRED_OBSERVATIONS) { + selected.set(new Decision(false, true, current.generation(), 0)); + return current; + } + int queriesSinceSample = current.queriesSinceSample() + 1; + if (queriesSinceSample == RESAMPLE_EVERY) { + // Do not keep the sender choice ready while the receiver probe is in flight. If it fails or never + // reports stats, subsequent queries remain on receiver sort until this entry expires. + long generation = _nextGeneration.incrementAndGet(); + selected.set(new Decision(false, true, generation, generation)); + return new Evidence(1, 0, generation, generation, current.baselineNanos(), 0); + } + selected.set(new Decision(true, false, current.generation(), 0)); + return new Evidence(REQUIRED_OBSERVATIONS, queriesSinceSample, current.generation(), 0, + current.baselineNanos(), current.slowSenderStreak()); + }); + return selected.get() != null ? selected.get() : new Decision(false, true, evidence.generation(), 0); + }); + } + + /// Call only after a successful query. Stats are stage-indexed; ambiguous or incomplete evidence resets the + /// candidate to the receiver-sort plan. A sender-sort query does not provide a new receiver-sort sample. + public void observe(@Nullable List stageStats) { + observe(stageStats, 0); + } + + /// The elapsed time spans dispatch through broker reduction, excluding planning and response construction. + public void observe(@Nullable List stageStats, long elapsedNanos) { + if (!_observed.compareAndSet(false, true)) { + return; + } + Map candidateCountByStage = new HashMap<>(); + for (ExchangeKey key : _decisions.keySet()) { + candidateCountByStage.merge(key.receiverStageId(), 1, Integer::sum); + } + for (Map.Entry entry : _decisions.entrySet()) { + Decision decision = entry.getValue(); + ExchangeKey key = entry.getKey(); + if (decision.senderSort()) { + // End-to-end latency cannot attribute a slowdown to one of several AUTO exchanges in a query. + if (_decisions.size() == 1) { + observeSender(key, decision, elapsedNanos); + } + continue; + } + if (!decision.profile()) { + continue; + } + ProfileResult result = candidateCountByStage.get(key.receiverStageId()) == 1 + ? profile(stageStats, key.receiverStageId()) + : ProfileResult.INCOMPLETE; + Evidence snapshot = _observations.getIfPresent(key); + if (snapshot == null || snapshot.generation() != decision.generation() + || (snapshot.pendingProbe() != 0 && snapshot.pendingProbe() != decision.probeToken()) + || (snapshot.pendingProbe() == 0 && decision.probeToken() != 0)) { + continue; + } + _observations.asMap().compute(key, (ignored, evidence) -> { + if (evidence == null || evidence.generation() != decision.generation()) { + return evidence; + } + if (evidence.pendingProbe() != 0) { + if (decision.probeToken() != evidence.pendingProbe()) { + return evidence; + } + if (result == ProfileResult.CANDIDATE) { + int observations = evidence.qualifyingObservations() == NEGATIVE_OBSERVATIONS + ? 1 + : REQUIRED_OBSERVATIONS; + return new Evidence(observations, 0, _nextGeneration.incrementAndGet(), 0, + baselineAfterCandidate(evidence, elapsedNanos, observations), 0); + } + return result == ProfileResult.NON_CANDIDATE + ? new Evidence(NEGATIVE_OBSERVATIONS, NEGATIVE_SKIP_QUERIES, _nextGeneration.incrementAndGet(), 0) + : new Evidence(0, 0, _nextGeneration.incrementAndGet(), 0); + } + if (decision.probeToken() != 0) { + return evidence; + } + if (result == ProfileResult.NON_CANDIDATE) { + return new Evidence(NEGATIVE_OBSERVATIONS, NEGATIVE_SKIP_QUERIES, _nextGeneration.incrementAndGet(), 0); + } + if (result == ProfileResult.INCOMPLETE) { + return new Evidence(0, 0, _nextGeneration.incrementAndGet(), 0); + } + int observations = Math.min(evidence.qualifyingObservations() + 1, REQUIRED_OBSERVATIONS); + return new Evidence(observations, evidence.queriesSinceSample(), evidence.generation(), 0, + baselineAfterCandidate(evidence, elapsedNanos, observations), 0); + }); + } + } + + private void observeSender(ExchangeKey key, Decision decision, long elapsedNanos) { + if (elapsedNanos <= 0) { + return; + } + Evidence snapshot = _observations.getIfPresent(key); + if (snapshot == null || snapshot.generation() != decision.generation() + || snapshot.qualifyingObservations() != REQUIRED_OBSERVATIONS || snapshot.pendingProbe() != 0 + || snapshot.baselineNanos() <= 0) { + return; + } + _observations.asMap().computeIfPresent(key, (ignored, evidence) -> { + if (evidence.generation() != decision.generation() + || evidence.qualifyingObservations() != REQUIRED_OBSERVATIONS || evidence.pendingProbe() != 0 + || evidence.baselineNanos() <= 0) { + return evidence; + } + long slowdownNanos = evidence.baselineNanos() * SLOWDOWN_PERCENT / 100; + int slowStreak = elapsedNanos > evidence.baselineNanos() + slowdownNanos + ? evidence.slowSenderStreak() + 1 + : 0; + if (slowStreak >= 2) { + return new Evidence(DISABLED_OBSERVATIONS, 0, _nextGeneration.incrementAndGet(), 0, + evidence.baselineNanos(), 0); + } + return new Evidence(REQUIRED_OBSERVATIONS, evidence.queriesSinceSample(), evidence.generation(), 0, + evidence.baselineNanos(), slowStreak); + }); + } + } + + private static long baselineAfterCandidate(Evidence evidence, long elapsedNanos, int observations) { + if (observations == 1) { + return Math.max(elapsedNanos, 0); + } + return evidence.baselineNanos() > 0 && elapsedNanos > 0 + ? Math.min(evidence.baselineNanos(), elapsedNanos) + : evidence.baselineNanos(); + } + + private static ProfileResult profile(@Nullable List stageStats, + int receiverStageId) { + if (stageStats == null || receiverStageId < 0 || receiverStageId >= stageStats.size()) { + return ProfileResult.INCOMPLETE; + } + MultiStageQueryStats.StageStats.Closed closed = stageStats.get(receiverStageId); + if (closed == null) { + return ProfileResult.INCOMPLETE; + } + StatMap receiveStats = null; + for (int i = 0; i <= closed.getLastOperatorIndex(); i++) { + if (closed.getOperatorType(i) != MultiStageOperator.Type.MAILBOX_RECEIVE) { + continue; + } + if (receiveStats != null) { + return ProfileResult.INCOMPLETE; + } + @SuppressWarnings("unchecked") + StatMap stats = + (StatMap) closed.getOperatorStats(i); + receiveStats = stats; + } + if (receiveStats == null) { + return ProfileResult.INCOMPLETE; + } + if (receiveStats.getInt(BaseMailboxReceiveOperator.StatKey.FAN_IN) != MIN_SENDERS + || receiveStats.getLong(BaseMailboxReceiveOperator.StatKey.EMITTED_ROWS) < MIN_ROWS) { + return ProfileResult.NON_CANDIDATE; + } + long sampledRows = receiveStats.getLong(BaseMailboxReceiveOperator.StatKey.AUTO_SAMPLED_ROWS); + if (sampledRows == 0) { + // A large zero-sample query is indistinguishable from stats reported by an older receiver. + return ProfileResult.INCOMPLETE; + } + if (receiveStats.getInt(BaseMailboxReceiveOperator.StatKey.AUTO_SAMPLE_STREAMS) != MIN_SENDERS + || sampledRows < MIN_SAMPLED_ROWS) { + // A successful receiver query with sampled rows but fewer than 65 rows on every sender cannot qualify. + // Cache that negative result instead of profiling every execution of a skewed or sparse exchange. + return ProfileResult.NON_CANDIDATE; + } + return receiveStats.getInt(BaseMailboxReceiveOperator.StatKey.AUTO_CANDIDATE_STREAMS) == MIN_SENDERS + ? ProfileResult.CANDIDATE + : ProfileResult.NON_CANDIDATE; + } + + private enum ProfileResult { + CANDIDATE, NON_CANDIDATE, INCOMPLETE + } + + private record ExchangeKey(long queryHash, int receiverStageId, int senderStageId, int inputHash, + int collationHash) { + } + + /// A negative observation uses -1 and counts down unprofiled receiver queries; -2 disables sender choice until TTL. + /// A qualified observation counts sender queries toward its next receiver probe. + private record Evidence(int qualifyingObservations, int queriesSinceSample, long generation, long pendingProbe, + long baselineNanos, int slowSenderStreak) { + private Evidence(int qualifyingObservations, int queriesSinceSample, long generation, long pendingProbe) { + this(qualifyingObservations, queriesSinceSample, generation, pendingProbe, 0, 0); + } + } + + private record Decision(boolean senderSort, boolean profile, long generation, long probeToken) { + } +} diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java new file mode 100644 index 000000000000..2dd2777aa392 --- /dev/null +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java @@ -0,0 +1,293 @@ +/** + * 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.broker.requesthandler; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.pinot.common.datatable.StatMap; +import org.apache.pinot.query.runtime.operator.BaseMailboxReceiveOperator; +import org.apache.pinot.query.runtime.operator.MultiStageOperator; +import org.apache.pinot.query.runtime.operator.OperatorTypeDescriptor; +import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +public class WindowSortAutoTunerTest { + @Test + public void testRequiresTwoStrongObservationsForTheSameExchange() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + List stats = stats(1, 400_000, 4, 4, 4, 256, false); + + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, stats)); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, stats)); + assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + WindowSortAutoTuner.Session senderSession = tuner.newSession(17L); + assertTrue(senderSession.useSenderSort(1, 2, 31, 41)); + senderSession.observe(null); + assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + + assertFalse(tuner.newSession(18L).useSenderSort(1, 2, 31, 41)); + assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 32, 41)); + assertFalse(tuner.newSession(17L).useSenderSort(1, 3, 31, 41)); + } + + @Test + public void testIncompleteOrChangedEvidenceResetsQualification() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + List strong = stats(1, 400_000, 4, 4, 4, 256, false); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, stats(1, 400_000, 4, 0, 0, 0, false))); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + + WindowSortAutoTuner missingStats = new WindowSortAutoTuner(); + assertFalse(observe(missingStats, 17L, 1, 2, 31, 41, strong)); + assertFalse(observe(missingStats, 17L, 1, 2, 31, 41, null)); + assertFalse(observe(missingStats, 17L, 1, 2, 31, 41, strong)); + assertFalse(missingStats.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertTrue(missingStats.newSession(17L).shouldProfile(1, 2, 31, 41)); + + WindowSortAutoTuner cold = new WindowSortAutoTuner(); + assertFalse(observe(cold, 17L, 1, 2, 31, 41, stats(1, 399_999, 4, 4, 4, 256, false))); + assertFalse(observe(cold, 17L, 1, 2, 31, 41, stats(1, 400_000, 3, 3, 3, 256, false))); + assertFalse(observe(cold, 17L, 1, 2, 31, 41, stats(1, 400_000, 4, 4, 3, 256, false))); + assertFalse(cold.newSession(17L).useSenderSort(1, 2, 31, 41)); + } + + @Test + public void testNegativeProfileCooldownAndReprobe() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + List presorted = stats(1, 400_000, 4, 4, 0, 256, false); + List candidate = stats(1, 400_000, 4, 4, 4, 256, false); + WindowSortAutoTuner.Session first = tuner.newSession(17L); + assertFalse(first.useSenderSort(1, 2, 31, 41)); + assertTrue(first.shouldProfile(1, 2, 31, 41)); + first.observe(presorted); + + for (int i = 0; i < 7; i++) { + WindowSortAutoTuner.Session skipped = tuner.newSession(17L); + assertFalse(skipped.useSenderSort(1, 2, 31, 41)); + assertFalse(skipped.shouldProfile(1, 2, 31, 41)); + skipped.observe(candidate); + } + WindowSortAutoTuner.Session probe = tuner.newSession(17L); + assertFalse(probe.useSenderSort(1, 2, 31, 41)); + assertTrue(probe.shouldProfile(1, 2, 31, 41)); + WindowSortAutoTuner.Session follower = tuner.newSession(17L); + assertFalse(follower.useSenderSort(1, 2, 31, 41)); + assertFalse(follower.shouldProfile(1, 2, 31, 41)); + follower.observe(candidate); + probe.observe(candidate); + WindowSortAutoTuner.Session second = tuner.newSession(17L); + assertFalse(second.useSenderSort(1, 2, 31, 41)); + assertTrue(second.shouldProfile(1, 2, 31, 41)); + second.observe(candidate); + assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + + WindowSortAutoTuner tiny = new WindowSortAutoTuner(); + assertFalse(observe(tiny, 17L, 1, 2, 31, 41, stats(1, 1_000, 4, 0, 0, 40, false))); + assertFalse(tiny.newSession(17L).shouldProfile(1, 2, 31, 41)); + + WindowSortAutoTuner skewed = new WindowSortAutoTuner(); + assertFalse(observe(skewed, 17L, 1, 2, 31, 41, stats(1, 400_000, 4, 1, 1, 68, false))); + assertFalse(skewed.newSession(17L).shouldProfile(1, 2, 31, 41)); + + WindowSortAutoTuner missingSample = new WindowSortAutoTuner(); + assertFalse(observe(missingSample, 17L, 1, 2, 31, 41, stats(1, 1_000, 4, 0, 0, 0, false))); + assertFalse(missingSample.newSession(17L).shouldProfile(1, 2, 31, 41)); + assertFalse(observe(missingSample, 18L, 1, 2, 31, 41, stats(1, 400_000, 4, 0, 0, 0, false))); + assertTrue(missingSample.newSession(18L).shouldProfile(1, 2, 31, 41)); + } + + @Test + public void testTwoConsecutiveSlowSenderRunsDisableChoice() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + List candidate = stats(1, 400_000, 4, 4, 4, 256, false); + WindowSortAutoTuner.Session first = tuner.newSession(17L); + assertFalse(first.useSenderSort(1, 2, 31, 41)); + first.observe(candidate, 1_000); + WindowSortAutoTuner.Session second = tuner.newSession(17L); + assertFalse(second.useSenderSort(1, 2, 31, 41)); + second.observe(candidate, 1_100); + + WindowSortAutoTuner.Session slow = tuner.newSession(17L); + assertTrue(slow.useSenderSort(1, 2, 31, 41)); + slow.observe(candidate, 1_151); + WindowSortAutoTuner.Session recovered = tuner.newSession(17L); + assertTrue(recovered.useSenderSort(1, 2, 31, 41)); + recovered.observe(candidate, 1_140); + WindowSortAutoTuner.Session inFlight = tuner.newSession(17L); + assertTrue(inFlight.useSenderSort(1, 2, 31, 41)); + WindowSortAutoTuner.Session slowAgain = tuner.newSession(17L); + assertTrue(slowAgain.useSenderSort(1, 2, 31, 41)); + slowAgain.observe(candidate, 1_160); + WindowSortAutoTuner.Session slowTwice = tuner.newSession(17L); + assertTrue(slowTwice.useSenderSort(1, 2, 31, 41)); + slowTwice.observe(candidate, 1_170); + + inFlight.observe(candidate, 900); + WindowSortAutoTuner.Session disabled = tuner.newSession(17L); + assertFalse(disabled.useSenderSort(1, 2, 31, 41)); + assertFalse(disabled.shouldProfile(1, 2, 31, 41)); + disabled.observe(candidate, 500); + assertFalse(tuner.newSession(17L).shouldProfile(1, 2, 31, 41)); + } + + @Test + public void testAmbiguousReceiverStatsDoNotTrainAnyExchange() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + List stats = stats(1, 400_000, 4, 4, 4, 256, false); + for (int i = 0; i < 2; i++) { + WindowSortAutoTuner.Session session = tuner.newSession(17L); + assertFalse(session.useSenderSort(1, 2, 31, 41)); + assertFalse(session.useSenderSort(1, 3, 32, 42)); + session.observe(stats); + } + assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertFalse(tuner.newSession(17L).useSenderSort(1, 3, 32, 42)); + + List twoReceives = stats(1, 400_000, 4, 4, 4, 256, true); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, twoReceives)); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, twoReceives)); + assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertTrue(tuner.newSession(17L).shouldProfile(1, 2, 31, 41)); + } + + @Test + public void testPeriodicReceiverResampleCanRevokeSenderChoice() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + List strong = stats(1, 400_000, 4, 4, 4, 256, false); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + + for (int i = 0; i < 7; i++) { + assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); + } + WindowSortAutoTuner.Session inFlightProbe = tuner.newSession(17L); + assertFalse(inFlightProbe.useSenderSort(1, 2, 31, 41)); + assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41), + "An in-flight receiver probe must not leave the sender plan ready"); + inFlightProbe.observe(strong); + assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); + + for (int i = 0; i < 6; i++) { + assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); + } + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, stats(1, 400_000, 4, 4, 0, 256, false))); + assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertFalse(tuner.newSession(17L).shouldProfile(1, 2, 31, 41)); + } + + @Test + public void testAbandonedReceiverProbeKeepsFallback() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + List strong = stats(1, 400_000, 4, 4, 4, 256, false); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + for (int i = 0; i < 7; i++) { + assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); + } + assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41), + "A later receiver query must not rearm an abandoned probe"); + } + + @Test + public void testTwoExchangesInOneQueryCanChooseDifferentPlans() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + List stats = new ArrayList<>(Collections.nCopies(4, null)); + stats.set(1, stats(1, 400_000, 4, 4, 4, 256, false).get(1)); + stats.set(3, stats(3, 400_000, 4, 4, 0, 256, false).get(3)); + for (int i = 0; i < 2; i++) { + WindowSortAutoTuner.Session session = tuner.newSession(17L); + assertFalse(session.useSenderSort(1, 2, 31, 41)); + assertFalse(session.useSenderSort(3, 4, 32, 42)); + session.observe(stats, 1_000); + } + WindowSortAutoTuner.Session session = tuner.newSession(17L); + assertTrue(session.useSenderSort(1, 2, 31, 41)); + assertFalse(session.useSenderSort(3, 4, 32, 42)); + session.observe(stats, 2_000); + WindowSortAutoTuner.Session next = tuner.newSession(17L); + assertTrue(next.useSenderSort(1, 2, 31, 41)); + assertFalse(next.useSenderSort(3, 4, 32, 42)); + next.observe(stats, 2_000); + assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41), + "Whole-query latency must not demote one exchange in a plan with multiple AUTO exchanges"); + } + + @Test + public void testOnlyTheSelectedProbeCanRearmAfterConcurrentReceiverQueries() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + List strong = stats(1, 400_000, 4, 4, 4, 256, false); + WindowSortAutoTuner.Session oldCold = tuner.newSession(17L); + assertFalse(oldCold.useSenderSort(1, 2, 31, 41)); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + for (int i = 0; i < 7; i++) { + assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); + } + WindowSortAutoTuner.Session probe = tuner.newSession(17L); + assertFalse(probe.useSenderSort(1, 2, 31, 41)); + oldCold.observe(strong); + assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); + assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + probe.observe(strong); + assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + } + + private static boolean observe(WindowSortAutoTuner tuner, long queryHash, int receiverStageId, int senderStageId, + int inputHash, int collationHash, List stats) { + WindowSortAutoTuner.Session session = tuner.newSession(queryHash); + boolean decision = session.useSenderSort(receiverStageId, senderStageId, inputHash, collationHash); + session.observe(stats); + return decision; + } + + private static List stats(int receiverStageId, long rows, int fanIn, + int sampledStreams, int candidateStreams, long sampledRows, boolean secondReceive) { + StatMap receive = new StatMap<>(BaseMailboxReceiveOperator.StatKey.class); + receive.merge(BaseMailboxReceiveOperator.StatKey.EMITTED_ROWS, rows); + receive.merge(BaseMailboxReceiveOperator.StatKey.FAN_IN, fanIn); + receive.merge(BaseMailboxReceiveOperator.StatKey.AUTO_SAMPLE_STREAMS, sampledStreams); + receive.merge(BaseMailboxReceiveOperator.StatKey.AUTO_CANDIDATE_STREAMS, candidateStreams); + receive.merge(BaseMailboxReceiveOperator.StatKey.AUTO_SAMPLED_ROWS, sampledRows); + + List types = new ArrayList<>(); + List> maps = new ArrayList<>(); + types.add(MultiStageOperator.Type.MAILBOX_RECEIVE); + maps.add(receive); + if (secondReceive) { + types.add(MultiStageOperator.Type.MAILBOX_RECEIVE); + maps.add(new StatMap<>(receive)); + } + List stages = + new ArrayList<>(Collections.nCopies(receiverStageId + 1, null)); + stages.set(receiverStageId, new MultiStageQueryStats.StageStats.Closed(types, maps)); + return stages; + } +} diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java index 3fcf70397b81..27663346a78a 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java @@ -761,11 +761,6 @@ public static boolean isUsePhysicalOptimizer(Map queryOptions, b return option != null ? Boolean.parseBoolean(option) : defaultValue; } - public static boolean isWindowSortOnSender(Map queryOptions, boolean defaultValue) { - String option = queryOptions.get(QueryOptionKey.WINDOW_SORT_ON_SENDER); - return option != null ? Boolean.parseBoolean(option) : defaultValue; - } - /// Reads the `streamStats` query option that opts a single query into the `SubmitWithStream` /// dispatch path. See [QueryOptionKey#STREAM_STATS]. public static boolean isStreamStats(Map queryOptions, boolean defaultValue) { diff --git a/pinot-common/src/main/proto/plan.proto b/pinot-common/src/main/proto/plan.proto index 55f2ff094c08..4a73f0b40cd1 100644 --- a/pinot-common/src/main/proto/plan.proto +++ b/pinot-common/src/main/proto/plan.proto @@ -200,6 +200,7 @@ message MailboxReceiveNode { repeated Collation collations = 5; bool sort = 6; bool sortedOnSender = 7; + bool autoProfile = 8; } // Each input stream is ordered by the logical plan. Older servers reject the unknown PlanNode oneof tag. diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java index 882ab9845126..73d248c5c09d 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java @@ -112,14 +112,6 @@ public void shouldReadIgnoreMissingSegmentsOption() { org.testng.Assert.assertFalse(QueryOptionsUtils.isIgnoreMissingSegments(optsMissing)); } - @Test - public void shouldResolveWindowSortOnSenderOptionOverDefault() { - assertFalse(QueryOptionsUtils.isWindowSortOnSender(Map.of(), false)); - assertTrue(QueryOptionsUtils.isWindowSortOnSender(Map.of(), true)); - assertTrue(QueryOptionsUtils.isWindowSortOnSender(Map.of(WINDOW_SORT_ON_SENDER, "true"), false)); - assertFalse(QueryOptionsUtils.isWindowSortOnSender(Map.of(WINDOW_SORT_ON_SENDER, "false"), true)); - } - @Test public void shouldReadInPredicatePruningThresholdOption() { // Any integer is accepted; a negative value means always attempt IN-predicate pruning diff --git a/pinot-core/src/main/java/org/apache/pinot/core/util/DataBlockExtractUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/util/DataBlockExtractUtils.java index 9914b1cd97d3..540375052588 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/util/DataBlockExtractUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/util/DataBlockExtractUtils.java @@ -42,6 +42,12 @@ private DataBlockExtractUtils() { } public static List extractRows(DataBlock dataBlock) { + return extractRows(dataBlock, dataBlock.getNumberOfRows()); + } + + /// Extracts at most the first {@code maxRows} rows without decoding the rest of the block. + public static List extractRows(DataBlock dataBlock, int maxRows) { + Preconditions.checkArgument(maxRows >= 0, "maxRows must be non-negative"); DataSchema dataSchema = dataBlock.getDataSchema(); ColumnDataType[] storedTypes = dataSchema.getStoredColumnDataTypes(); int numColumns = storedTypes.length; @@ -49,7 +55,7 @@ public static List extractRows(DataBlock dataBlock) { for (int colId = 0; colId < numColumns; colId++) { nullBitmaps[colId] = dataBlock.getNullRowIds(colId); } - int numRows = dataBlock.getNumberOfRows(); + int numRows = Math.min(dataBlock.getNumberOfRows(), maxRows); List rows = new ArrayList<>(numRows); for (int rowId = 0; rowId < numRows; rowId++) { Object[] row = new Object[numColumns]; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalSortExchange.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalSortExchange.java index 3683dcf4bfcd..d7881eff7802 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalSortExchange.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalSortExchange.java @@ -43,6 +43,7 @@ public class PinotLogicalSortExchange extends SortExchange { protected final boolean _isSortOnSender; protected final boolean _isSortOnReceiver; + protected final boolean _autoWindowSort; protected final PinotRelExchangeType _exchangeType; // Can be used to override the partitioning info calculated from the distribution trait. @Nullable @@ -50,11 +51,12 @@ public class PinotLogicalSortExchange extends SortExchange { private PinotLogicalSortExchange(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, RelDistribution distribution, PinotRelExchangeType exchangeType, RelCollation collation, boolean isSortOnSender, - boolean isSortOnReceiver, @Nullable Boolean prePartitioned) { + boolean isSortOnReceiver, boolean autoWindowSort, @Nullable Boolean prePartitioned) { super(cluster, traitSet, input, distribution, collation); _exchangeType = exchangeType; _isSortOnSender = isSortOnSender; _isSortOnReceiver = isSortOnReceiver; + _autoWindowSort = autoWindowSort; _prePartitioned = prePartitioned; } @@ -64,6 +66,7 @@ public PinotLogicalSortExchange(RelInput input) { _exchangeType = PinotRelExchangeType.STREAMING; _isSortOnSender = false; _isSortOnReceiver = true; + _autoWindowSort = false; _prePartitioned = null; } @@ -73,6 +76,13 @@ public static PinotLogicalSortExchange create(RelNode input, RelDistribution dis isSortOnReceiver, null); } + /// Defers the global-window ordering choice until stage ids are assigned by the fragmenter. + public static PinotLogicalSortExchange createWindowAuto(RelNode input, RelDistribution distribution, + RelCollation collation) { + return create(input, distribution, PinotRelExchangeType.getDefaultExchangeType(), collation, false, true, true, + null); + } + public static PinotLogicalSortExchange create(RelNode input, RelDistribution distribution, RelCollation collation, boolean isSortOnSender, boolean isSortOnReceiver, @Nullable Boolean prePartitioned) { return create(input, distribution, PinotRelExchangeType.getDefaultExchangeType(), collation, isSortOnSender, @@ -91,12 +101,19 @@ public static PinotLogicalSortExchange create(RelNode input, RelDistribution dis public static PinotLogicalSortExchange create(RelNode input, RelDistribution distribution, PinotRelExchangeType exchangeType, RelCollation collation, boolean isSortOnSender, boolean isSortOnReceiver, @Nullable Boolean prePartitioned) { + return create(input, distribution, exchangeType, collation, isSortOnSender, isSortOnReceiver, false, + prePartitioned); + } + + private static PinotLogicalSortExchange create(RelNode input, RelDistribution distribution, + PinotRelExchangeType exchangeType, RelCollation collation, boolean isSortOnSender, boolean isSortOnReceiver, + boolean autoWindowSort, @Nullable Boolean prePartitioned) { RelOptCluster cluster = input.getCluster(); collation = RelCollationTraitDef.INSTANCE.canonize(collation); distribution = RelDistributionTraitDef.INSTANCE.canonize(distribution); RelTraitSet traitSet = input.getTraitSet().replace(Convention.NONE).replace(distribution).replace(collation); return new PinotLogicalSortExchange(cluster, traitSet, input, distribution, exchangeType, collation, isSortOnSender, - isSortOnReceiver, prePartitioned); + isSortOnReceiver, autoWindowSort, prePartitioned); } //~ Methods ---------------------------------------------------------------- @@ -105,13 +122,16 @@ public static PinotLogicalSortExchange create(RelNode input, RelDistribution dis public SortExchange copy(RelTraitSet traitSet, RelNode newInput, RelDistribution newDistribution, RelCollation newCollation) { return new PinotLogicalSortExchange(this.getCluster(), traitSet, newInput, newDistribution, _exchangeType, - newCollation, _isSortOnSender, _isSortOnReceiver, _prePartitioned); + newCollation, _isSortOnSender, _isSortOnReceiver, _autoWindowSort, _prePartitioned); } @Override public RelWriter explainTerms(RelWriter pw) { RelWriter relWriter = super.explainTerms(pw).item("isSortOnSender", _isSortOnSender).item("isSortOnReceiver", _isSortOnReceiver); + if (_autoWindowSort) { + relWriter.item("autoWindowSort", true); + } if (_exchangeType != PinotRelExchangeType.getDefaultExchangeType()) { relWriter.item("relExchangeType", _exchangeType); } @@ -126,6 +146,10 @@ public boolean isSortOnReceiver() { return _isSortOnReceiver; } + public boolean isAutoWindowSort() { + return _autoWindowSort; + } + public PinotRelExchangeType getExchangeType() { return _exchangeType; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlan.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlan.java new file mode 100644 index 000000000000..5c61de5ecff6 --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlan.java @@ -0,0 +1,34 @@ +/** + * 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.planner.logical; + + +/// Selects one ordering strategy for every worker of a global window exchange. +/// +/// The receiver and sender stage identifiers are stable within one planned query. Implementations must return false +/// when evidence for this exchange is absent or ambiguous. A selector is shared only during planning and must not +/// mutate the physical plan after dispatch. +public interface WindowSortAutoPlan { + boolean useSenderSort(int receiverStageId, int senderStageId, int inputHash, int collationHash); + + /// Whether a receiver-sort execution should collect a fresh AUTO sample. The default preserves existing selectors. + default boolean shouldProfile(int receiverStageId, int senderStageId, int inputHash, int collationHash) { + return true; + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/ExchangeNode.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/ExchangeNode.java index 7a1eb15a5494..463f5131fd44 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/ExchangeNode.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/ExchangeNode.java @@ -41,6 +41,7 @@ public class ExchangeNode extends BasePlanNode { private final List _collations; private final boolean _sortOnSender; private final boolean _sortOnReceiver; + private final boolean _autoWindowSort; // Table names should be set for SUB_PLAN exchange type. private final Set _tableNames; @Nullable @@ -51,6 +52,15 @@ public ExchangeNode(int stageId, DataSchema dataSchema, List inputs, P RelDistribution.Type distributionType, @Nullable List keys, boolean prePartitioned, @Nullable List collations, boolean sortOnSender, boolean sortOnReceiver, @Nullable Set tableNames, ExchangeStrategy exchangeStrategy, String hashFunction) { + this(stageId, dataSchema, inputs, exchangeType, distributionType, keys, prePartitioned, collations, sortOnSender, + sortOnReceiver, false, tableNames, exchangeStrategy, hashFunction); + } + + public ExchangeNode(int stageId, DataSchema dataSchema, List inputs, PinotRelExchangeType exchangeType, + RelDistribution.Type distributionType, @Nullable List keys, boolean prePartitioned, + @Nullable List collations, boolean sortOnSender, boolean sortOnReceiver, + boolean autoWindowSort, @Nullable Set tableNames, ExchangeStrategy exchangeStrategy, + String hashFunction) { super(stageId, dataSchema, null, inputs); _exchangeType = exchangeType; _distributionType = distributionType; @@ -59,6 +69,7 @@ public ExchangeNode(int stageId, DataSchema dataSchema, List inputs, P _collations = collations; _sortOnSender = sortOnSender; _sortOnReceiver = sortOnReceiver; + _autoWindowSort = autoWindowSort; _tableNames = tableNames; _exchangeStrategy = exchangeStrategy; _hashFunction = hashFunction; @@ -94,6 +105,10 @@ public boolean isSortOnReceiver() { return _sortOnReceiver; } + public boolean isAutoWindowSort() { + return _autoWindowSort; + } + @Nullable public Set getTableNames() { return _tableNames; @@ -121,7 +136,7 @@ public T visit(PlanNodeVisitor visitor, C context) { @Override public PlanNode withInputs(List inputs) { return new ExchangeNode(_stageId, _dataSchema, inputs, _exchangeType, _distributionType, _keys, _prePartitioned, - _collations, _sortOnSender, _sortOnReceiver, _tableNames, null, _hashFunction); + _collations, _sortOnSender, _sortOnReceiver, _autoWindowSort, _tableNames, null, _hashFunction); } @Override @@ -137,6 +152,7 @@ public boolean equals(Object o) { } ExchangeNode that = (ExchangeNode) o; return _sortOnSender == that._sortOnSender && _sortOnReceiver == that._sortOnReceiver + && _autoWindowSort == that._autoWindowSort && _prePartitioned == that._prePartitioned && _exchangeType == that._exchangeType && _distributionType == that._distributionType && Objects.equals(_keys, that._keys) && Objects.equals(_collations, that._collations) && Objects.equals(_tableNames, that._tableNames) @@ -146,6 +162,6 @@ public boolean equals(Object o) { @Override public int hashCode() { return Objects.hash(super.hashCode(), _exchangeType, _distributionType, _keys, _sortOnSender, _sortOnReceiver, - _prePartitioned, _collations, _tableNames, _hashFunction); + _autoWindowSort, _prePartitioned, _collations, _tableNames, _hashFunction); } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/MailboxReceiveNode.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/MailboxReceiveNode.java index ebbdf7aab3a8..ade61b6239d7 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/MailboxReceiveNode.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/MailboxReceiveNode.java @@ -36,6 +36,7 @@ public class MailboxReceiveNode extends BasePlanNode { private final List _collations; private final boolean _sort; private final boolean _sortedOnSender; + private final boolean _autoProfile; // NOTE: This is only available during query planning, and should not be serialized. private transient MailboxSendNode _sender; @@ -45,6 +46,14 @@ public MailboxReceiveNode(int stageId, DataSchema dataSchema, int senderStageId, PinotRelExchangeType exchangeType, RelDistribution.Type distributionType, @Nullable List keys, @Nullable List collations, boolean sort, boolean sortedOnSender, @Nullable MailboxSendNode sender) { + this(stageId, dataSchema, senderStageId, exchangeType, distributionType, keys, collations, sort, + sortedOnSender, false, sender); + } + + public MailboxReceiveNode(int stageId, DataSchema dataSchema, int senderStageId, + PinotRelExchangeType exchangeType, RelDistribution.Type distributionType, @Nullable List keys, + @Nullable List collations, boolean sort, boolean sortedOnSender, boolean autoProfile, + @Nullable MailboxSendNode sender) { super(stageId, dataSchema, null, List.of()); _senderStageId = senderStageId; _exchangeType = exchangeType; @@ -53,6 +62,7 @@ public MailboxReceiveNode(int stageId, DataSchema dataSchema, int senderStageId, _collations = collations != null ? collations : List.of(); _sort = sort; _sortedOnSender = sortedOnSender; + _autoProfile = autoProfile; _sender = sender; } @@ -90,6 +100,11 @@ public boolean isSortedOnSender() { return _sortedOnSender; } + /// Whether this receiver gathers bounded order samples for the default AUTO strategy. + public boolean isAutoProfile() { + return _autoProfile; + } + public MailboxSendNode getSender() { assert _sender != null; return _sender; @@ -102,7 +117,7 @@ public void setSender(MailboxSendNode sender) { @Override public String explain() { - return "MAIL_RECEIVE(" + _distributionType + ")"; + return "MAIL_RECEIVE(" + _distributionType + ")" + (_autoProfile ? "[WINDOW_SORT_AUTO_PROFILE]" : ""); } @Override @@ -118,7 +133,7 @@ public PlanNode withInputs(List inputs) { public MailboxReceiveNode withSender(MailboxSendNode sender) { return new MailboxReceiveNode(_stageId, _dataSchema, _senderStageId, _exchangeType, _distributionType, _keys, - _collations, _sort, _sortedOnSender, sender); + _collations, _sort, _sortedOnSender, _autoProfile, sender); } @Override @@ -134,6 +149,7 @@ public boolean equals(Object o) { } MailboxReceiveNode that = (MailboxReceiveNode) o; return _senderStageId == that._senderStageId && _sort == that._sort && _sortedOnSender == that._sortedOnSender + && _autoProfile == that._autoProfile && _exchangeType == that._exchangeType && _distributionType == that._distributionType && Objects.equals(_keys, that._keys) && Objects.equals(_collations, that._collations); } @@ -141,6 +157,6 @@ public boolean equals(Object o) { @Override public int hashCode() { return Objects.hash(super.hashCode(), _senderStageId, _exchangeType, _distributionType, _keys, _collations, _sort, - _sortedOnSender); + _sortedOnSender, _autoProfile); } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java index f627972e5ee4..06c39dd8b213 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java @@ -169,7 +169,7 @@ private static MailboxReceiveNode deserializeMailboxReceiveNode(Plan.PlanNode pr protoMailboxReceiveNode.getSenderStageId(), convertExchangeType(protoMailboxReceiveNode.getExchangeType()), convertDistributionType(protoMailboxReceiveNode.getDistributionType()), protoMailboxReceiveNode.getKeysList(), convertCollations(protoMailboxReceiveNode.getCollationsList()), protoMailboxReceiveNode.getSort(), - protoMailboxReceiveNode.getSortedOnSender(), null); + protoMailboxReceiveNode.getSortedOnSender(), protoMailboxReceiveNode.getAutoProfile(), null); } private static MailboxSendNode deserializeMailboxSendNode(Plan.PlanNode protoNode) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java index 7c8fa01d42df..7f8cc605463a 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java @@ -195,6 +195,7 @@ public Void visitMailboxReceive(MailboxReceiveNode node, Plan.PlanNode.Builder b .addAllCollations(convertCollations(node.getCollations())) .setSort(node.isSort()) .setSortedOnSender(node.isSortedOnSender()) + .setAutoProfile(node.isAutoProfile()) .build(); builder.setMailboxReceiveNode(mailboxReceiveNode); return null; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java index da5afb98b9e9..4eb487eab065 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java @@ -278,7 +278,13 @@ public int merge(int value1, int value2) { /// can fail in two corner cases, where rows are read but never emitted: after the downstream operator has /// early terminated, and when a sorted receive buffers its rows and then ends in error. Such a worker is /// reported as idle despite having been given data. - NON_ACTIVE_WORKERS(StatMap.Type.INT); + NON_ACTIVE_WORKERS(StatMap.Type.INT), + /// Number of sender streams with 64 adjacent order-key comparisons in the AUTO baseline sample. + AUTO_SAMPLE_STREAMS(StatMap.Type.INT), + /// Sampled streams with 16 to 48 inversions and at most 4 ties among 64 adjacent comparisons. + AUTO_CANDIDATE_STREAMS(StatMap.Type.INT), + /// Number of rows inspected for AUTO order profiling, capped at 65 per sender stream. + AUTO_SAMPLED_ROWS(StatMap.Type.LONG); private final StatMap.Type _type; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java index 222785a0f7e7..4c9df1c1d70f 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java @@ -18,8 +18,17 @@ */ package org.apache.pinot.query.runtime.operator; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.core.util.DataBlockExtractUtils; +import org.apache.pinot.query.mailbox.ReceivingMailbox; import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.SerializedDataBlock; +import org.apache.pinot.query.runtime.operator.utils.AsyncStream; +import org.apache.pinot.query.runtime.operator.utils.SortUtils; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -30,9 +39,20 @@ public class MailboxReceiveOperator extends BaseMailboxReceiveOperator { private static final Logger LOGGER = LoggerFactory.getLogger(MailboxReceiveOperator.class); private static final String EXPLAIN_NAME = "MAILBOX_RECEIVE"; + private static final int AUTO_SAMPLE_COMPARISONS = 64; + private static final int AUTO_DISORDERED_INVERSIONS = 16; + private static final int AUTO_MAX_DISORDERED_INVERSIONS = 48; + private static final int AUTO_MAX_EQUAL_PAIRS = 4; + + @Nullable + private final SortUtils.SortComparator _autoComparator; + private final Map, StreamSample> _autoSamples; + private boolean _autoProfilingDisabled; public MailboxReceiveOperator(OpChainExecutionContext context, MailboxReceiveNode node) { super(context, node); + _autoComparator = node.isAutoProfile() ? new SortUtils.SortComparator(node.getCollations(), false) : null; + _autoSamples = _autoComparator != null ? new IdentityHashMap<>() : Map.of(); } @Override @@ -56,10 +76,74 @@ protected MseBlock getNextBlock() { block = _multiConsumer.readMseBlockBlocking(); } if (block.isData()) { + if (_autoComparator != null && !_autoProfilingDisabled) { + try { + sampleSenderOrder((MseBlock.Data) block); + } catch (RuntimeException e) { + // AUTO evidence must never turn a successful receiver-sort query into a failed one. + _autoProfilingDisabled = true; + _autoSamples.clear(); + LOGGER.debug("Disabling window AUTO order profiling for this receiver", e); + } + } checkTerminationAndSampleUsage(); } else { onEos(); } return block; } + + /// Profiles the first 64 adjacent pairs from each sender independently. This is diagnostic only: the original + /// block and its ordering are passed through unchanged. Serialized blocks decode no more than the sampled prefix. + private void sampleSenderOrder(MseBlock.Data block) { + AsyncStream stream = _multiConsumer.getLastReadStream(); + if (stream == null) { + return; + } + StreamSample sample = _autoSamples.computeIfAbsent(stream, ignored -> new StreamSample()); + int remainingRows = AUTO_SAMPLE_COMPARISONS + 1 - sample._sampledRows; + if (remainingRows == 0) { + return; + } + List rows; + if (block.isRowHeap()) { + rows = block.asRowHeap().getRows(); + } else if (block instanceof SerializedDataBlock) { + rows = DataBlockExtractUtils.extractRows(((SerializedDataBlock) block).getDataBlock(), remainingRows); + } else { + return; + } + int numSampled = Math.min(rows.size(), remainingRows); + for (int i = 0; i < numSampled; i++) { + Object[] row = rows.get(i); + if (sample._previousRow != null) { + int comparison = _autoComparator.compare(sample._previousRow, row); + if (comparison > 0) { + sample._inversions++; + } else if (comparison == 0) { + sample._equalPairs++; + } + } + sample._previousRow = row; + sample._sampledRows++; + } + _statMap.merge(StatKey.AUTO_SAMPLED_ROWS, numSampled); + if (sample._sampledRows == AUTO_SAMPLE_COMPARISONS + 1) { + _statMap.merge(StatKey.AUTO_SAMPLE_STREAMS, 1); + if (sample._inversions >= AUTO_DISORDERED_INVERSIONS + && sample._inversions <= AUTO_MAX_DISORDERED_INVERSIONS + && sample._equalPairs <= AUTO_MAX_EQUAL_PAIRS) { + _statMap.merge(StatKey.AUTO_CANDIDATE_STREAMS, 1); + } + sample._previousRow = null; + } + } + + private static final class StreamSample { + private int _sampledRows; + private int _inversions; + private int _equalPairs; + @Nullable + private Object[] _previousRow; + } } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java index b2efc3652df9..75370a3c61e0 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java @@ -21,13 +21,18 @@ import java.io.IOException; import java.lang.ref.ReferenceQueue; import java.lang.ref.WeakReference; +import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.calcite.rel.RelDistribution; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.calcite.rel.RelFieldCollation.Direction; +import org.apache.calcite.rel.RelFieldCollation.NullDirection; import org.apache.pinot.common.datatable.StatMap; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.datablock.DataBlockBuilder; import org.apache.pinot.query.mailbox.MailboxService; import org.apache.pinot.query.mailbox.ReceivingMailbox; import org.apache.pinot.query.planner.physical.MailboxIdUtils; @@ -39,6 +44,7 @@ import org.apache.pinot.query.routing.WorkerMetadata; import org.apache.pinot.query.runtime.blocks.ErrorMseBlock; import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.SerializedDataBlock; import org.apache.pinot.query.runtime.operator.MultiStageOperator.Type; import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; @@ -230,6 +236,82 @@ public void shouldReceiveMailboxFromTwoServers() { } } + @Test + public void shouldProfileEachSenderAcrossBlocksWithBoundedSerializedExtraction() + throws IOException { + Object[][] tieHeavyRows = new Object[70][]; + Object[][] disorderedRows = new Object[70][]; + for (int i = 0; i < 70; i++) { + tieHeavyRows[i] = new Object[]{i >= 64 ? null : (i % 3 == 0 ? 1 : 0), 1}; + disorderedRows[i] = new Object[]{i, 2}; + } + for (int i = 0; i < 16; i++) { + int index = i * 4; + Object[] row = disorderedRows[index]; + disorderedRows[index] = disorderedRows[index + 1]; + disorderedRows[index + 1] = row; + } + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + new ReceivingMailbox.MseBlockWithStats(new SerializedDataBlock( + DataBlockBuilder.buildFromRows(Arrays.asList(tieHeavyRows).subList(0, 32), DATA_SCHEMA)), List.of()), + new ReceivingMailbox.MseBlockWithStats(new SerializedDataBlock( + DataBlockBuilder.buildFromRows(Arrays.asList(tieHeavyRows).subList(32, 70), DATA_SCHEMA)), List.of()), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, Arrays.copyOfRange(disorderedRows, 0, 20)), + OperatorTestUtil.blockWithStats(DATA_SCHEMA, Arrays.copyOfRange(disorderedRows, 20, 70)), + OperatorTestUtil.eosWithEmptyStats()); + + try (MailboxReceiveOperator operator = getOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, Long.MAX_VALUE, true)) { + drain(operator); + StatMap stats = operator.copyStatMaps(); + assertEquals(stats.getLong(BaseMailboxReceiveOperator.StatKey.AUTO_SAMPLED_ROWS), 130L); + assertEquals(stats.getInt(BaseMailboxReceiveOperator.StatKey.AUTO_SAMPLE_STREAMS), 2); + assertEquals(stats.getInt(BaseMailboxReceiveOperator.StatKey.AUTO_CANDIDATE_STREAMS), 1); + } + } + + @Test + public void shouldNotTreatReverseSortedStreamAsDisordered() { + Object[][] reverseSortedRows = new Object[65][]; + for (int i = 0; i < reverseSortedRows.length; i++) { + reverseSortedRows[i] = new Object[]{65 - i, 1}; + } + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn(OperatorTestUtil.blockWithStats(DATA_SCHEMA, reverseSortedRows), + OperatorTestUtil.eosWithEmptyStats()); + + try (MailboxReceiveOperator operator = getOperator(_stageMetadata1, + RelDistribution.Type.SINGLETON, Long.MAX_VALUE, true)) { + drain(operator); + StatMap stats = operator.copyStatMaps(); + assertEquals(stats.getInt(BaseMailboxReceiveOperator.StatKey.AUTO_SAMPLE_STREAMS), 1); + assertEquals(stats.getInt(BaseMailboxReceiveOperator.StatKey.AUTO_CANDIDATE_STREAMS), 0); + } + } + + @Test + public void shouldPreserveBlocksWhenAutoProfilingCannotCompareRows() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn(OperatorTestUtil.blockWithStats(DATA_SCHEMA, + new Object[]{1, 1}, new Object[]{2, 2}), OperatorTestUtil.eosWithEmptyStats()); + OpChainExecutionContext context = OperatorTestUtil.getOpChainContext(_mailboxService, Long.MAX_VALUE, + _stageMetadata1); + MailboxReceiveNode node = mock(MailboxReceiveNode.class); + when(node.getDistributionType()).thenReturn(RelDistribution.Type.SINGLETON); + when(node.getSenderStageId()).thenReturn(1); + when(node.isAutoProfile()).thenReturn(true); + when(node.getCollations()).thenReturn(List.of(new RelFieldCollation(5))); + try (MailboxReceiveOperator operator = new MailboxReceiveOperator(context, node)) { + assertEquals(((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows().size(), 2); + assertTrue(operator.nextBlock().isSuccess()); + assertEquals(operator.copyStatMaps().getInt(BaseMailboxReceiveOperator.StatKey.AUTO_CANDIDATE_STREAMS), 0); + } + } + @Test public void shouldGetReceptionReceiveErrorMailbox() { when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); @@ -342,10 +424,20 @@ public void differentUpstreamStatsProduceEmptyStats() private MailboxReceiveOperator getOperator(StageMetadata stageMetadata, RelDistribution.Type distributionType, long deadlineMs) { + return getOperator(stageMetadata, distributionType, deadlineMs, false); + } + + private MailboxReceiveOperator getOperator(StageMetadata stageMetadata, RelDistribution.Type distributionType, + long deadlineMs, boolean autoProfile) { OpChainExecutionContext context = OperatorTestUtil.getOpChainContext(_mailboxService, deadlineMs, stageMetadata); MailboxReceiveNode node = mock(MailboxReceiveNode.class); when(node.getDistributionType()).thenReturn(distributionType); when(node.getSenderStageId()).thenReturn(1); + if (autoProfile) { + when(node.isAutoProfile()).thenReturn(true); + when(node.getCollations()).thenReturn( + List.of(new RelFieldCollation(0, Direction.ASCENDING, NullDirection.LAST))); + } return new MailboxReceiveOperator(context, node); } From cd7ec90ad24a347cd25a375cf5554aba31bfbb77 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 30 Sep 2026 13:30:59 +0530 Subject: [PATCH 2/9] Keep adaptive window sorting explicitly opt-in Preserve the boolean option and planner builder. Only allocate a tuning session for auto; the default keeps receiver sorting. --- .../pinot/common/utils/config/QueryOptionsUtils.java | 5 +++++ .../pinot/common/utils/config/QueryOptionsUtilsTest.java | 8 ++++++++ 2 files changed, 13 insertions(+) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java index 27663346a78a..3fcf70397b81 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java @@ -761,6 +761,11 @@ public static boolean isUsePhysicalOptimizer(Map queryOptions, b return option != null ? Boolean.parseBoolean(option) : defaultValue; } + public static boolean isWindowSortOnSender(Map queryOptions, boolean defaultValue) { + String option = queryOptions.get(QueryOptionKey.WINDOW_SORT_ON_SENDER); + return option != null ? Boolean.parseBoolean(option) : defaultValue; + } + /// Reads the `streamStats` query option that opts a single query into the `SubmitWithStream` /// dispatch path. See [QueryOptionKey#STREAM_STATS]. public static boolean isStreamStats(Map queryOptions, boolean defaultValue) { diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java index 73d248c5c09d..882ab9845126 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java @@ -112,6 +112,14 @@ public void shouldReadIgnoreMissingSegmentsOption() { org.testng.Assert.assertFalse(QueryOptionsUtils.isIgnoreMissingSegments(optsMissing)); } + @Test + public void shouldResolveWindowSortOnSenderOptionOverDefault() { + assertFalse(QueryOptionsUtils.isWindowSortOnSender(Map.of(), false)); + assertTrue(QueryOptionsUtils.isWindowSortOnSender(Map.of(), true)); + assertTrue(QueryOptionsUtils.isWindowSortOnSender(Map.of(WINDOW_SORT_ON_SENDER, "true"), false)); + assertFalse(QueryOptionsUtils.isWindowSortOnSender(Map.of(WINDOW_SORT_ON_SENDER, "false"), true)); + } + @Test public void shouldReadInPredicatePruningThresholdOption() { // Any integer is accepted; a negative value means always attempt IN-predicate pruning From f60387a44c070d1cabc54be48d122e6519eafcb1 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 30 Sep 2026 13:39:22 +0530 Subject: [PATCH 3/9] Release window profiling samples when receive stops Drop short-stream samples on terminal blocks, errors, close and cancellation while preserving output and stats. --- .../operator/MailboxReceiveOperator.java | 26 +++++++++++++++-- .../operator/MailboxReceiveOperatorTest.java | 29 +++++++++++++++++++ 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java index 4c9df1c1d70f..1678be73ade9 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java @@ -46,7 +46,7 @@ public class MailboxReceiveOperator extends BaseMailboxReceiveOperator { @Nullable private final SortUtils.SortComparator _autoComparator; - private final Map, StreamSample> _autoSamples; + private Map, StreamSample> _autoSamples; private boolean _autoProfilingDisabled; public MailboxReceiveOperator(OpChainExecutionContext context, MailboxReceiveNode node) { @@ -67,6 +67,15 @@ protected Logger logger() { @Override protected MseBlock getNextBlock() { + try { + return readNextBlock(); + } catch (RuntimeException e) { + releaseBuffers(); + throw e; + } + } + + private MseBlock readNextBlock() { MseBlock block = _multiConsumer.readMseBlockBlocking(); // When early termination flag is set, caller is expecting an EOS block to be returned, however since the 2 stages // between sending/receiving mailbox are setting early termination flag asynchronously, there's chances that the @@ -81,13 +90,13 @@ protected MseBlock getNextBlock() { sampleSenderOrder((MseBlock.Data) block); } catch (RuntimeException e) { // AUTO evidence must never turn a successful receiver-sort query into a failed one. - _autoProfilingDisabled = true; - _autoSamples.clear(); + releaseBuffers(); LOGGER.debug("Disabling window AUTO order profiling for this receiver", e); } } checkTerminationAndSampleUsage(); } else { + releaseBuffers(); onEos(); } return block; @@ -139,6 +148,17 @@ private void sampleSenderOrder(MseBlock.Data block) { } } + @Override + protected void releaseBuffers() { + _autoProfilingDisabled = true; + _autoSamples = Map.of(); + } + + @Override + protected boolean hasBufferedState() { + return !_autoSamples.isEmpty(); + } + private static final class StreamSample { private int _sampledRows; private int _inversions; diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java index 75370a3c61e0..d4ed6b7d8aa3 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java @@ -63,6 +63,7 @@ import static org.mockito.Mockito.when; import static org.mockito.MockitoAnnotations.openMocks; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; @@ -274,6 +275,34 @@ public void shouldProfileEachSenderAcrossBlocksWithBoundedSerializedExtraction() } } + @Test + public void shouldReleaseShortAutoSamplesOnTerminalAndTeardown() { + for (String termination : List.of("eos", "error", "close", "cancel")) { + Object[] row = new Object[]{1, 1}; + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn(OperatorTestUtil.blockWithStats(DATA_SCHEMA, row), + termination.equals("error") + ? OperatorTestUtil.errorWithEmptyStats(new RuntimeException("test error")) + : OperatorTestUtil.eosWithEmptyStats()); + try (MailboxReceiveOperator operator = getOperator(_stageMetadata1, + RelDistribution.Type.SINGLETON, Long.MAX_VALUE, true)) { + MseBlock.Data block = (MseBlock.Data) operator.nextBlock(); + assertTrue(operator.hasBufferedState()); + switch (termination) { + case "eos", "error" -> assertTrue(operator.nextBlock().isEos()); + case "close" -> operator.close(); + case "cancel" -> operator.cancel(new RuntimeException("test cancel")); + default -> throw new AssertionError(termination); + } + assertFalse(operator.hasBufferedState(), termination); + operator.close(); + assertFalse(operator.hasBufferedState(), "Repeated teardown must be safe"); + assertEquals(block.asRowHeap().getRows().get(0), row); + assertEquals(operator.copyStatMaps().getLong(BaseMailboxReceiveOperator.StatKey.AUTO_SAMPLED_ROWS), 1L); + } + } + } + @Test public void shouldNotTreatReverseSortedStreamAsDisordered() { Object[][] reverseSortedRows = new Object[65][]; From fde886fdf3a2213a901080ca06d21f8ad1a7a56d Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 30 Sep 2026 13:59:03 +0530 Subject: [PATCH 4/9] Keep AUTO serde coverage compatible with the join stack Use a distinct enum import and the existing assertion import so both independent follow-ups merge cleanly. From 2c262097a4935951b43794939376e07f81a69398 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 30 Sep 2026 19:41:58 +0530 Subject: [PATCH 5/9] Trim unused AUTO planner helpers Keep only the AUTO entry points used by planning and broker observations. Apply receiver observations only to existing cache entries, preserving generation and probe checks. --- .../requesthandler/WindowSortAutoTuner.java | 8 ++----- .../WindowSortAutoTunerTest.java | 22 +++++++++---------- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java index 97cdb5601010..e947045d3d4a 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java @@ -141,10 +141,6 @@ private Decision decisionFor(int receiverStageId, int senderStageId, int inputHa /// Call only after a successful query. Stats are stage-indexed; ambiguous or incomplete evidence resets the /// candidate to the receiver-sort plan. A sender-sort query does not provide a new receiver-sort sample. - public void observe(@Nullable List stageStats) { - observe(stageStats, 0); - } - /// The elapsed time spans dispatch through broker reduction, excluding planning and response construction. public void observe(@Nullable List stageStats, long elapsedNanos) { if (!_observed.compareAndSet(false, true)) { @@ -176,8 +172,8 @@ public void observe(@Nullable List stage || (snapshot.pendingProbe() == 0 && decision.probeToken() != 0)) { continue; } - _observations.asMap().compute(key, (ignored, evidence) -> { - if (evidence == null || evidence.generation() != decision.generation()) { + _observations.asMap().computeIfPresent(key, (ignored, evidence) -> { + if (evidence.generation() != decision.generation()) { return evidence; } if (evidence.pendingProbe() != 0) { diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java index 2dd2777aa392..64a5288429bd 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java @@ -43,7 +43,7 @@ public void testRequiresTwoStrongObservationsForTheSameExchange() { assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); WindowSortAutoTuner.Session senderSession = tuner.newSession(17L); assertTrue(senderSession.useSenderSort(1, 2, 31, 41)); - senderSession.observe(null); + senderSession.observe(null, 0); assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); assertFalse(tuner.newSession(18L).useSenderSort(1, 2, 31, 41)); @@ -85,13 +85,13 @@ public void testNegativeProfileCooldownAndReprobe() { WindowSortAutoTuner.Session first = tuner.newSession(17L); assertFalse(first.useSenderSort(1, 2, 31, 41)); assertTrue(first.shouldProfile(1, 2, 31, 41)); - first.observe(presorted); + first.observe(presorted, 0); for (int i = 0; i < 7; i++) { WindowSortAutoTuner.Session skipped = tuner.newSession(17L); assertFalse(skipped.useSenderSort(1, 2, 31, 41)); assertFalse(skipped.shouldProfile(1, 2, 31, 41)); - skipped.observe(candidate); + skipped.observe(candidate, 0); } WindowSortAutoTuner.Session probe = tuner.newSession(17L); assertFalse(probe.useSenderSort(1, 2, 31, 41)); @@ -99,12 +99,12 @@ public void testNegativeProfileCooldownAndReprobe() { WindowSortAutoTuner.Session follower = tuner.newSession(17L); assertFalse(follower.useSenderSort(1, 2, 31, 41)); assertFalse(follower.shouldProfile(1, 2, 31, 41)); - follower.observe(candidate); - probe.observe(candidate); + follower.observe(candidate, 0); + probe.observe(candidate, 0); WindowSortAutoTuner.Session second = tuner.newSession(17L); assertFalse(second.useSenderSort(1, 2, 31, 41)); assertTrue(second.shouldProfile(1, 2, 31, 41)); - second.observe(candidate); + second.observe(candidate, 0); assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); WindowSortAutoTuner tiny = new WindowSortAutoTuner(); @@ -164,7 +164,7 @@ public void testAmbiguousReceiverStatsDoNotTrainAnyExchange() { WindowSortAutoTuner.Session session = tuner.newSession(17L); assertFalse(session.useSenderSort(1, 2, 31, 41)); assertFalse(session.useSenderSort(1, 3, 32, 42)); - session.observe(stats); + session.observe(stats, 0); } assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); assertFalse(tuner.newSession(17L).useSenderSort(1, 3, 32, 42)); @@ -190,7 +190,7 @@ public void testPeriodicReceiverResampleCanRevokeSenderChoice() { assertFalse(inFlightProbe.useSenderSort(1, 2, 31, 41)); assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41), "An in-flight receiver probe must not leave the sender plan ready"); - inFlightProbe.observe(strong); + inFlightProbe.observe(strong, 0); assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); for (int i = 0; i < 6; i++) { @@ -253,10 +253,10 @@ public void testOnlyTheSelectedProbeCanRearmAfterConcurrentReceiverQueries() { } WindowSortAutoTuner.Session probe = tuner.newSession(17L); assertFalse(probe.useSenderSort(1, 2, 31, 41)); - oldCold.observe(strong); + oldCold.observe(strong, 0); assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); - probe.observe(strong); + probe.observe(strong, 0); assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); } @@ -264,7 +264,7 @@ private static boolean observe(WindowSortAutoTuner tuner, long queryHash, int re int inputHash, int collationHash, List stats) { WindowSortAutoTuner.Session session = tuner.newSession(queryHash); boolean decision = session.useSenderSort(receiverStageId, senderStageId, inputHash, collationHash); - session.observe(stats); + session.observe(stats, 0); return decision; } From 294ab693b505780c17786bbc50d2282b783333ba Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 4 Oct 2026 03:01:14 +0800 Subject: [PATCH 6/9] Declare annotations required for warning-enabled compilation From 5f8714bf396f4f1f627cd28d0f016bcf9332ecc4 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 4 Oct 2026 03:01:22 +0800 Subject: [PATCH 7/9] Resolve adaptive window sorting before stage fragmentation Use stable logical exchange identities and per-query stage bindings for AUTO evidence. Lower explicit sender Sort and k-way exchanges for warm decisions, preserve unbounded cold window input and fail-closed cluster capability, and port the plain transport merge contract with regression coverage. --- .../MultiStageBrokerRequestHandler.java | 39 +++- .../requesthandler/WindowSortAutoTuner.java | 42 +++-- .../WindowSortAutoTunerTest.java | 173 +++++++++++------ .../rel/logical/PinotLogicalSortExchange.java | 2 +- .../PinotWindowExchangeNodeInsertRule.java | 39 ++-- .../apache/pinot/query/QueryEnvironment.java | 65 +++++-- .../query/planner/explain/PlanNodeMerger.java | 3 +- .../logical/EquivalentStagesFinder.java | 1 + .../logical/PinotLogicalQueryPlanner.java | 20 +- .../query/planner/logical/PlanFragmenter.java | 25 ++- .../logical/RelToPlanNodeConverter.java | 19 ++ .../planner/logical/WindowSortAutoPlan.java | 21 ++- .../logical/WindowSortAutoPlanner.java | 99 ++++++++++ .../query/planner/plannode/ExchangeNode.java | 20 +- ...PinotWindowExchangeNodeInsertRuleTest.java | 2 +- .../pinot/query/QueryCompilationTest.java | 178 +++++++++++++++++- .../pinot/query/QueryEnvironmentTestBase.java | 13 +- .../planner/serde/PlanNodeSerDeTest.java | 15 ++ .../operator/WindowAggregateOperatorTest.java | 4 +- .../pinot/spi/utils/CommonConstants.java | 3 +- 20 files changed, 625 insertions(+), 158 deletions(-) create mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlanner.java diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java index 496c33834bee..d179f8e4ff06 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java @@ -31,6 +31,7 @@ import java.util.List; import java.util.Locale; import java.util.Map; +import java.util.Objects; import java.util.Set; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; @@ -83,6 +84,7 @@ import org.apache.pinot.query.QueryEnvironment; import org.apache.pinot.query.mailbox.MailboxService; import org.apache.pinot.query.planner.explain.AskingServerStageExplainer; +import org.apache.pinot.query.planner.logical.WindowSortAutoPlan; import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; import org.apache.pinot.query.planner.physical.DispatchableSubPlan; import org.apache.pinot.query.planner.plannode.PlanNode; @@ -138,6 +140,7 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler { private final WorkerManager _multiClusterWorkerManager; private final MailboxService _mailboxService; private final QueryDispatcher _queryDispatcher; + private final WindowSortAutoTuner _windowSortAutoTuner = new WindowSortAutoTuner(); @Nullable private final ServerRoutingStatsManager _serverRoutingStatsManager; private final boolean _explainAskingServerDefault; @@ -495,7 +498,16 @@ private QueryEnvironment.CompiledQuery compileQuery(long requestId, String query Map queryOptions = sqlNodeAndOptions.getOptions(); try { - ImmutableQueryEnvironment.Config queryEnvConf = getQueryEnvConf(httpHeaders, queryOptions, requestId); + String database = DatabaseUtils.extractDatabaseFromQueryRequest(queryOptions, httpHeaders); + // The profile key contains no SQL text. Stage/input/collation fingerprints further separate individual + // ordered-window exchanges, even when a query has several windows. + String windowSortMode = queryOptions.getOrDefault(CommonConstants.Broker.Request.QueryOptionKey + .WINDOW_SORT_ON_SENDER, _config.getProperty(CommonConstants.Broker.CONFIG_OF_WINDOW_SORT_ON_SENDER, + CommonConstants.Broker.DEFAULT_WINDOW_SORT_ON_SENDER_MODE)); + WindowSortAutoTuner.Session windowSortAutoSession = "auto".equalsIgnoreCase(windowSortMode) + ? _windowSortAutoTuner.newSession(Objects.hash(query, database, queryOptions)) : null; + ImmutableQueryEnvironment.Config queryEnvConf = getQueryEnvConf(httpHeaders, queryOptions, requestId, + windowSortAutoSession); QueryEnvironment queryEnv = new QueryEnvironment(queryEnvConf, _multiClusterRoutingContext); return callAsync(requestId, query, () -> queryEnv.compile(query, sqlNodeAndOptions), queryTimer); } catch (WebApplicationException e) { @@ -540,6 +552,11 @@ private void checkAuthorization(RequesterIdentity requesterIdentity, RequestCont private ImmutableQueryEnvironment.Config getQueryEnvConf(HttpHeaders httpHeaders, Map queryOptions, long requestId) { + return getQueryEnvConf(httpHeaders, queryOptions, requestId, null); + } + + private ImmutableQueryEnvironment.Config getQueryEnvConf(HttpHeaders httpHeaders, Map queryOptions, + long requestId, @Nullable WindowSortAutoPlan windowSortAutoPlan) { String database = DatabaseUtils.extractDatabaseFromQueryRequest(queryOptions, httpHeaders); boolean inferPartitionHint = _config.getProperty(CommonConstants.Broker.CONFIG_OF_INFER_PARTITION_HINT, CommonConstants.Broker.DEFAULT_INFER_PARTITION_HINT); @@ -590,9 +607,9 @@ private ImmutableQueryEnvironment.Config getQueryEnvConf(HttpHeaders httpHeaders int sealedInListThreshold = _config.getProperty( CommonConstants.Broker.CONFIG_OF_SEALED_IN_LIST_THRESHOLD, CommonConstants.Broker.DEFAULT_SEALED_IN_LIST_THRESHOLD); - boolean windowSortOnSender = _config.getProperty( + String windowSortOnSenderMode = _config.getProperty( CommonConstants.Broker.CONFIG_OF_WINDOW_SORT_ON_SENDER, - CommonConstants.Broker.DEFAULT_WINDOW_SORT_ON_SENDER); + CommonConstants.Broker.DEFAULT_WINDOW_SORT_ON_SENDER_MODE); boolean defaultUnnestColumnPruning = _config.getProperty( CommonConstants.Broker.CONFIG_OF_UNNEST_COLUMN_PRUNING, CommonConstants.Broker.DEFAULT_UNNEST_COLUMN_PRUNING); @@ -632,7 +649,8 @@ private ImmutableQueryEnvironment.Config getQueryEnvConf(HttpHeaders httpHeaders .defaultDisabledPlannerRules(_defaultDisabledPlannerRules) .defaultSortExchangeCopyLimit(sortExchangeCopyThreshold) .defaultSealedInListThreshold(sealedInListThreshold) - .defaultWindowSortOnSender(windowSortOnSender) + .defaultWindowSortOnSenderMode(windowSortOnSenderMode) + .windowSortAutoPlan(windowSortAutoPlan) .isKWayMergeSupported(!QueryOptionsUtils.isMultiClusterRoutingEnabled(queryOptions, false) && _kWayMergeSupported.getAsBoolean()) .build(); @@ -828,6 +846,7 @@ private BrokerResponse query(QueryEnvironment.CompiledQuery query, long requestI } } + long dispatchToReduceTimeNs = System.nanoTime() - executionStartTimeNs; BrokerResponseNativeV2 brokerResponse = new BrokerResponseNativeV2(); QueryProcessingException processingException = queryResults.getProcessingException(); @@ -844,9 +863,8 @@ private BrokerResponse query(QueryEnvironment.CompiledQuery query, long requestI requestContext.setErrorCode(errorCode); } else { brokerResponse.setResultTable(queryResults.getResultTable()); - long executionEndTimeNs = System.nanoTime(); updatePhaseTimingForTables(tableNames, BrokerQueryPhase.QUERY_EXECUTION, - executionEndTimeNs - executionStartTimeNs); + dispatchToReduceTimeNs); } brokerResponse.setClientRequestId(clientRequestId); @@ -885,6 +903,15 @@ private BrokerResponse query(QueryEnvironment.CompiledQuery query, long requestI } } + List stageCoverage = queryResults.getStageCoverage(); + if (processingException == null && brokerResponse.getExceptions().isEmpty() + && (stageCoverage == null || stageCoverage.stream().noneMatch(c -> c != null + && (c.getMissing() > 0 || c.getMergeFailed() > 0))) + && query.getWindowSortAutoPlan() instanceof WindowSortAutoTuner.Session) { + ((WindowSortAutoTuner.Session) query.getWindowSortAutoPlan()).observe(queryResults.getQueryStats(), + dispatchToReduceTimeNs); + } + fillOldBrokerResponseStats(brokerResponse, queryResults.getQueryStats(), dispatchableSubPlan, queryResults.getStageCoverage(), queryResults.getStageStatsTrees()); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java index e947045d3d4a..2f6a928bf6b3 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTuner.java @@ -50,7 +50,7 @@ public final class WindowSortAutoTuner { private static final int RESAMPLE_EVERY = 8; private static final int SLOWDOWN_PERCENT = 15; - private final Cache _observations = CacheBuilder.newBuilder() + private final Cache _observations = CacheBuilder.newBuilder() .maximumSize(1_024).expireAfterWrite(Duration.ofMinutes(30)).build(); private final AtomicLong _nextGeneration = new AtomicLong(); @@ -61,7 +61,8 @@ public Session newSession(long queryHash) { /// A single query's planner decisions and, on successful completion, its receiver-stage observations. public final class Session implements WindowSortAutoPlan { private final long _queryHash; - private final Map _decisions = new HashMap<>(); + private final Map _decisions = new HashMap<>(); + private final Map _receiverStages = new HashMap<>(); private final AtomicBoolean _observed = new AtomicBoolean(); private Session(long queryHash) { @@ -69,17 +70,24 @@ private Session(long queryHash) { } @Override - public boolean useSenderSort(int receiverStageId, int senderStageId, int inputHash, int collationHash) { - return decisionFor(receiverStageId, senderStageId, inputHash, collationHash).senderSort(); + public boolean useSenderSort(WindowSortAutoPlan.ExchangeKey key) { + return decisionFor(key).senderSort(); } @Override - public boolean shouldProfile(int receiverStageId, int senderStageId, int inputHash, int collationHash) { - return decisionFor(receiverStageId, senderStageId, inputHash, collationHash).profile(); + public boolean shouldProfile(WindowSortAutoPlan.ExchangeKey key) { + return decisionFor(key).profile(); } - private Decision decisionFor(int receiverStageId, int senderStageId, int inputHash, int collationHash) { - ExchangeKey key = new ExchangeKey(_queryHash, receiverStageId, senderStageId, inputHash, collationHash); + @Override + public void bind(WindowSortAutoPlan.ExchangeKey logicalKey, int receiverStageId, int senderStageId) { + ObservationKey key = new ObservationKey(_queryHash, logicalKey); + // Repeated planning may allocate a different stage. Ambiguous bindings never train a candidate. + _receiverStages.merge(key, receiverStageId, (previous, current) -> previous.equals(current) ? current : -1); + } + + private Decision decisionFor(WindowSortAutoPlan.ExchangeKey logicalKey) { + ObservationKey key = new ObservationKey(_queryHash, logicalKey); return _decisions.computeIfAbsent(key, ignored -> { Evidence evidence = _observations.getIfPresent(key); if (evidence == null) { @@ -147,12 +155,12 @@ public void observe(@Nullable List stage return; } Map candidateCountByStage = new HashMap<>(); - for (ExchangeKey key : _decisions.keySet()) { - candidateCountByStage.merge(key.receiverStageId(), 1, Integer::sum); + for (ObservationKey key : _decisions.keySet()) { + candidateCountByStage.merge(_receiverStages.getOrDefault(key, -1), 1, Integer::sum); } - for (Map.Entry entry : _decisions.entrySet()) { + for (Map.Entry entry : _decisions.entrySet()) { Decision decision = entry.getValue(); - ExchangeKey key = entry.getKey(); + ObservationKey key = entry.getKey(); if (decision.senderSort()) { // End-to-end latency cannot attribute a slowdown to one of several AUTO exchanges in a query. if (_decisions.size() == 1) { @@ -163,8 +171,9 @@ public void observe(@Nullable List stage if (!decision.profile()) { continue; } - ProfileResult result = candidateCountByStage.get(key.receiverStageId()) == 1 - ? profile(stageStats, key.receiverStageId()) + int receiverStageId = _receiverStages.getOrDefault(key, -1); + ProfileResult result = candidateCountByStage.get(receiverStageId) == 1 + ? profile(stageStats, receiverStageId) : ProfileResult.INCOMPLETE; Evidence snapshot = _observations.getIfPresent(key); if (snapshot == null || snapshot.generation() != decision.generation() @@ -207,7 +216,7 @@ public void observe(@Nullable List stage } } - private void observeSender(ExchangeKey key, Decision decision, long elapsedNanos) { + private void observeSender(ObservationKey key, Decision decision, long elapsedNanos) { if (elapsedNanos <= 0) { return; } @@ -295,8 +304,7 @@ private enum ProfileResult { CANDIDATE, NON_CANDIDATE, INCOMPLETE } - private record ExchangeKey(long queryHash, int receiverStageId, int senderStageId, int inputHash, - int collationHash) { + private record ObservationKey(long queryHash, WindowSortAutoPlan.ExchangeKey logicalKey) { } /// A negative observation uses -1 and counts down unprofiled receiver queries; -2 disables sender choice until TTL. diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java index 64a5288429bd..9b283519faad 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/WindowSortAutoTunerTest.java @@ -22,6 +22,7 @@ import java.util.Collections; import java.util.List; import org.apache.pinot.common.datatable.StatMap; +import org.apache.pinot.query.planner.logical.WindowSortAutoPlan.ExchangeKey; import org.apache.pinot.query.runtime.operator.BaseMailboxReceiveOperator; import org.apache.pinot.query.runtime.operator.MultiStageOperator; import org.apache.pinot.query.runtime.operator.OperatorTypeDescriptor; @@ -40,15 +41,15 @@ public void testRequiresTwoStrongObservationsForTheSameExchange() { assertFalse(observe(tuner, 17L, 1, 2, 31, 41, stats)); assertFalse(observe(tuner, 17L, 1, 2, 31, 41, stats)); - assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertTrue(choose(tuner.newSession(17L), 1, 2, 31, 41)); WindowSortAutoTuner.Session senderSession = tuner.newSession(17L); - assertTrue(senderSession.useSenderSort(1, 2, 31, 41)); + assertTrue(choose(senderSession, 1, 2, 31, 41)); senderSession.observe(null, 0); - assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertTrue(choose(tuner.newSession(17L), 1, 2, 31, 41)); - assertFalse(tuner.newSession(18L).useSenderSort(1, 2, 31, 41)); - assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 32, 41)); - assertFalse(tuner.newSession(17L).useSenderSort(1, 3, 31, 41)); + assertFalse(choose(tuner.newSession(18L), 1, 2, 31, 41)); + assertFalse(choose(tuner.newSession(17L), 1, 2, 32, 41)); + assertFalse(choose(tuner.newSession(17L), 1, 3, 31, 41)); } @Test @@ -59,22 +60,22 @@ public void testIncompleteOrChangedEvidenceResetsQualification() { assertFalse(observe(tuner, 17L, 1, 2, 31, 41, stats(1, 400_000, 4, 0, 0, 0, false))); assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); - assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertFalse(choose(tuner.newSession(17L), 1, 2, 31, 41)); assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); - assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertTrue(choose(tuner.newSession(17L), 1, 2, 31, 41)); WindowSortAutoTuner missingStats = new WindowSortAutoTuner(); assertFalse(observe(missingStats, 17L, 1, 2, 31, 41, strong)); assertFalse(observe(missingStats, 17L, 1, 2, 31, 41, null)); assertFalse(observe(missingStats, 17L, 1, 2, 31, 41, strong)); - assertFalse(missingStats.newSession(17L).useSenderSort(1, 2, 31, 41)); - assertTrue(missingStats.newSession(17L).shouldProfile(1, 2, 31, 41)); + assertFalse(choose(missingStats.newSession(17L), 1, 2, 31, 41)); + assertTrue(profile(missingStats.newSession(17L), 1, 2, 31, 41)); WindowSortAutoTuner cold = new WindowSortAutoTuner(); assertFalse(observe(cold, 17L, 1, 2, 31, 41, stats(1, 399_999, 4, 4, 4, 256, false))); assertFalse(observe(cold, 17L, 1, 2, 31, 41, stats(1, 400_000, 3, 3, 3, 256, false))); assertFalse(observe(cold, 17L, 1, 2, 31, 41, stats(1, 400_000, 4, 4, 3, 256, false))); - assertFalse(cold.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertFalse(choose(cold.newSession(17L), 1, 2, 31, 41)); } @Test @@ -83,43 +84,43 @@ public void testNegativeProfileCooldownAndReprobe() { List presorted = stats(1, 400_000, 4, 4, 0, 256, false); List candidate = stats(1, 400_000, 4, 4, 4, 256, false); WindowSortAutoTuner.Session first = tuner.newSession(17L); - assertFalse(first.useSenderSort(1, 2, 31, 41)); - assertTrue(first.shouldProfile(1, 2, 31, 41)); + assertFalse(choose(first, 1, 2, 31, 41)); + assertTrue(profile(first, 1, 2, 31, 41)); first.observe(presorted, 0); for (int i = 0; i < 7; i++) { WindowSortAutoTuner.Session skipped = tuner.newSession(17L); - assertFalse(skipped.useSenderSort(1, 2, 31, 41)); - assertFalse(skipped.shouldProfile(1, 2, 31, 41)); + assertFalse(choose(skipped, 1, 2, 31, 41)); + assertFalse(profile(skipped, 1, 2, 31, 41)); skipped.observe(candidate, 0); } WindowSortAutoTuner.Session probe = tuner.newSession(17L); - assertFalse(probe.useSenderSort(1, 2, 31, 41)); - assertTrue(probe.shouldProfile(1, 2, 31, 41)); + assertFalse(choose(probe, 1, 2, 31, 41)); + assertTrue(profile(probe, 1, 2, 31, 41)); WindowSortAutoTuner.Session follower = tuner.newSession(17L); - assertFalse(follower.useSenderSort(1, 2, 31, 41)); - assertFalse(follower.shouldProfile(1, 2, 31, 41)); + assertFalse(choose(follower, 1, 2, 31, 41)); + assertFalse(profile(follower, 1, 2, 31, 41)); follower.observe(candidate, 0); probe.observe(candidate, 0); WindowSortAutoTuner.Session second = tuner.newSession(17L); - assertFalse(second.useSenderSort(1, 2, 31, 41)); - assertTrue(second.shouldProfile(1, 2, 31, 41)); + assertFalse(choose(second, 1, 2, 31, 41)); + assertTrue(profile(second, 1, 2, 31, 41)); second.observe(candidate, 0); - assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertTrue(choose(tuner.newSession(17L), 1, 2, 31, 41)); WindowSortAutoTuner tiny = new WindowSortAutoTuner(); assertFalse(observe(tiny, 17L, 1, 2, 31, 41, stats(1, 1_000, 4, 0, 0, 40, false))); - assertFalse(tiny.newSession(17L).shouldProfile(1, 2, 31, 41)); + assertFalse(profile(tiny.newSession(17L), 1, 2, 31, 41)); WindowSortAutoTuner skewed = new WindowSortAutoTuner(); assertFalse(observe(skewed, 17L, 1, 2, 31, 41, stats(1, 400_000, 4, 1, 1, 68, false))); - assertFalse(skewed.newSession(17L).shouldProfile(1, 2, 31, 41)); + assertFalse(profile(skewed.newSession(17L), 1, 2, 31, 41)); WindowSortAutoTuner missingSample = new WindowSortAutoTuner(); assertFalse(observe(missingSample, 17L, 1, 2, 31, 41, stats(1, 1_000, 4, 0, 0, 0, false))); - assertFalse(missingSample.newSession(17L).shouldProfile(1, 2, 31, 41)); + assertFalse(profile(missingSample.newSession(17L), 1, 2, 31, 41)); assertFalse(observe(missingSample, 18L, 1, 2, 31, 41, stats(1, 400_000, 4, 0, 0, 0, false))); - assertTrue(missingSample.newSession(18L).shouldProfile(1, 2, 31, 41)); + assertTrue(profile(missingSample.newSession(18L), 1, 2, 31, 41)); } @Test @@ -127,33 +128,33 @@ public void testTwoConsecutiveSlowSenderRunsDisableChoice() { WindowSortAutoTuner tuner = new WindowSortAutoTuner(); List candidate = stats(1, 400_000, 4, 4, 4, 256, false); WindowSortAutoTuner.Session first = tuner.newSession(17L); - assertFalse(first.useSenderSort(1, 2, 31, 41)); + assertFalse(choose(first, 1, 2, 31, 41)); first.observe(candidate, 1_000); WindowSortAutoTuner.Session second = tuner.newSession(17L); - assertFalse(second.useSenderSort(1, 2, 31, 41)); + assertFalse(choose(second, 1, 2, 31, 41)); second.observe(candidate, 1_100); WindowSortAutoTuner.Session slow = tuner.newSession(17L); - assertTrue(slow.useSenderSort(1, 2, 31, 41)); + assertTrue(choose(slow, 1, 2, 31, 41)); slow.observe(candidate, 1_151); WindowSortAutoTuner.Session recovered = tuner.newSession(17L); - assertTrue(recovered.useSenderSort(1, 2, 31, 41)); + assertTrue(choose(recovered, 1, 2, 31, 41)); recovered.observe(candidate, 1_140); WindowSortAutoTuner.Session inFlight = tuner.newSession(17L); - assertTrue(inFlight.useSenderSort(1, 2, 31, 41)); + assertTrue(choose(inFlight, 1, 2, 31, 41)); WindowSortAutoTuner.Session slowAgain = tuner.newSession(17L); - assertTrue(slowAgain.useSenderSort(1, 2, 31, 41)); + assertTrue(choose(slowAgain, 1, 2, 31, 41)); slowAgain.observe(candidate, 1_160); WindowSortAutoTuner.Session slowTwice = tuner.newSession(17L); - assertTrue(slowTwice.useSenderSort(1, 2, 31, 41)); + assertTrue(choose(slowTwice, 1, 2, 31, 41)); slowTwice.observe(candidate, 1_170); inFlight.observe(candidate, 900); WindowSortAutoTuner.Session disabled = tuner.newSession(17L); - assertFalse(disabled.useSenderSort(1, 2, 31, 41)); - assertFalse(disabled.shouldProfile(1, 2, 31, 41)); + assertFalse(choose(disabled, 1, 2, 31, 41)); + assertFalse(profile(disabled, 1, 2, 31, 41)); disabled.observe(candidate, 500); - assertFalse(tuner.newSession(17L).shouldProfile(1, 2, 31, 41)); + assertFalse(profile(tuner.newSession(17L), 1, 2, 31, 41)); } @Test @@ -162,18 +163,18 @@ public void testAmbiguousReceiverStatsDoNotTrainAnyExchange() { List stats = stats(1, 400_000, 4, 4, 4, 256, false); for (int i = 0; i < 2; i++) { WindowSortAutoTuner.Session session = tuner.newSession(17L); - assertFalse(session.useSenderSort(1, 2, 31, 41)); - assertFalse(session.useSenderSort(1, 3, 32, 42)); + assertFalse(choose(session, 1, 2, 31, 41)); + assertFalse(choose(session, 1, 3, 32, 42)); session.observe(stats, 0); } - assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); - assertFalse(tuner.newSession(17L).useSenderSort(1, 3, 32, 42)); + assertFalse(choose(tuner.newSession(17L), 1, 2, 31, 41)); + assertFalse(choose(tuner.newSession(17L), 1, 3, 32, 42)); List twoReceives = stats(1, 400_000, 4, 4, 4, 256, true); assertFalse(observe(tuner, 17L, 1, 2, 31, 41, twoReceives)); assertFalse(observe(tuner, 17L, 1, 2, 31, 41, twoReceives)); - assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); - assertTrue(tuner.newSession(17L).shouldProfile(1, 2, 31, 41)); + assertFalse(choose(tuner.newSession(17L), 1, 2, 31, 41)); + assertTrue(profile(tuner.newSession(17L), 1, 2, 31, 41)); } @Test @@ -187,8 +188,8 @@ public void testPeriodicReceiverResampleCanRevokeSenderChoice() { assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); } WindowSortAutoTuner.Session inFlightProbe = tuner.newSession(17L); - assertFalse(inFlightProbe.useSenderSort(1, 2, 31, 41)); - assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41), + assertFalse(choose(inFlightProbe, 1, 2, 31, 41)); + assertFalse(choose(tuner.newSession(17L), 1, 2, 31, 41), "An in-flight receiver probe must not leave the sender plan ready"); inFlightProbe.observe(strong, 0); assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); @@ -197,8 +198,8 @@ public void testPeriodicReceiverResampleCanRevokeSenderChoice() { assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); } assertFalse(observe(tuner, 17L, 1, 2, 31, 41, stats(1, 400_000, 4, 4, 0, 256, false))); - assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); - assertFalse(tuner.newSession(17L).shouldProfile(1, 2, 31, 41)); + assertFalse(choose(tuner.newSession(17L), 1, 2, 31, 41)); + assertFalse(profile(tuner.newSession(17L), 1, 2, 31, 41)); } @Test @@ -210,9 +211,9 @@ public void testAbandonedReceiverProbeKeepsFallback() { for (int i = 0; i < 7; i++) { assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); } - assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertFalse(choose(tuner.newSession(17L), 1, 2, 31, 41)); assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); - assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41), + assertFalse(choose(tuner.newSession(17L), 1, 2, 31, 41), "A later receiver query must not rearm an abandoned probe"); } @@ -224,19 +225,19 @@ public void testTwoExchangesInOneQueryCanChooseDifferentPlans() { stats.set(3, stats(3, 400_000, 4, 4, 0, 256, false).get(3)); for (int i = 0; i < 2; i++) { WindowSortAutoTuner.Session session = tuner.newSession(17L); - assertFalse(session.useSenderSort(1, 2, 31, 41)); - assertFalse(session.useSenderSort(3, 4, 32, 42)); + assertFalse(choose(session, 1, 2, 31, 41)); + assertFalse(choose(session, 3, 4, 32, 42)); session.observe(stats, 1_000); } WindowSortAutoTuner.Session session = tuner.newSession(17L); - assertTrue(session.useSenderSort(1, 2, 31, 41)); - assertFalse(session.useSenderSort(3, 4, 32, 42)); + assertTrue(choose(session, 1, 2, 31, 41)); + assertFalse(choose(session, 3, 4, 32, 42)); session.observe(stats, 2_000); WindowSortAutoTuner.Session next = tuner.newSession(17L); - assertTrue(next.useSenderSort(1, 2, 31, 41)); - assertFalse(next.useSenderSort(3, 4, 32, 42)); + assertTrue(choose(next, 1, 2, 31, 41)); + assertFalse(choose(next, 3, 4, 32, 42)); next.observe(stats, 2_000); - assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41), + assertTrue(choose(tuner.newSession(17L), 1, 2, 31, 41), "Whole-query latency must not demote one exchange in a plan with multiple AUTO exchanges"); } @@ -245,25 +246,77 @@ public void testOnlyTheSelectedProbeCanRearmAfterConcurrentReceiverQueries() { WindowSortAutoTuner tuner = new WindowSortAutoTuner(); List strong = stats(1, 400_000, 4, 4, 4, 256, false); WindowSortAutoTuner.Session oldCold = tuner.newSession(17L); - assertFalse(oldCold.useSenderSort(1, 2, 31, 41)); + assertFalse(choose(oldCold, 1, 2, 31, 41)); assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); for (int i = 0; i < 7; i++) { assertTrue(observe(tuner, 17L, 1, 2, 31, 41, null)); } WindowSortAutoTuner.Session probe = tuner.newSession(17L); - assertFalse(probe.useSenderSort(1, 2, 31, 41)); + assertFalse(choose(probe, 1, 2, 31, 41)); oldCold.observe(strong, 0); assertFalse(observe(tuner, 17L, 1, 2, 31, 41, strong)); - assertFalse(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertFalse(choose(tuner.newSession(17L), 1, 2, 31, 41)); probe.observe(strong, 0); - assertTrue(tuner.newSession(17L).useSenderSort(1, 2, 31, 41)); + assertTrue(choose(tuner.newSession(17L), 1, 2, 31, 41)); + } + + @Test + public void testStageRenumberingReusesLogicalEvidence() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + ExchangeKey key = new ExchangeKey("root/0/1", 31, 41); + for (int receiverStage : List.of(1, 7)) { + WindowSortAutoTuner.Session session = tuner.newSession(17L); + assertFalse(session.useSenderSort(key)); + session.bind(key, receiverStage, receiverStage + 1); + session.observe(stats(receiverStage, 400_000, 4, 4, 4, 256, false), 1_000); + } + WindowSortAutoTuner.Session warm = tuner.newSession(17L); + assertTrue(warm.useSenderSort(key), "Allocated stage IDs must not enter the persistent evidence key"); + assertFalse(warm.useSenderSort(new ExchangeKey("root/1/1", 31, 41)), + "Identical input/order at another logical exchange remains separate"); + } + + @Test + public void testMissingOrAmbiguousBindingDoesNotTrain() { + WindowSortAutoTuner tuner = new WindowSortAutoTuner(); + ExchangeKey key = new ExchangeKey("root/0", 31, 41); + for (int i = 0; i < 2; i++) { + WindowSortAutoTuner.Session unbound = tuner.newSession(17L); + assertFalse(unbound.useSenderSort(key)); + unbound.observe(stats(1, 400_000, 4, 4, 4, 256, false), 1_000); + WindowSortAutoTuner.Session ambiguous = tuner.newSession(17L); + assertFalse(ambiguous.useSenderSort(key)); + ambiguous.bind(key, 1, 2); + ambiguous.bind(key, 3, 4); + ambiguous.observe(stats(1, 400_000, 4, 4, 4, 256, false), 1_000); + } + assertFalse(tuner.newSession(17L).useSenderSort(key)); + } + + private static ExchangeKey key(int ordinal, int inputHash, int collationHash) { + return new ExchangeKey("root/" + ordinal, inputHash, collationHash); + } + + private static boolean choose(WindowSortAutoTuner.Session session, int receiverStageId, int logicalOrdinal, + int inputHash, int collationHash) { + ExchangeKey key = key(logicalOrdinal, inputHash, collationHash); + boolean senderSort = session.useSenderSort(key); + session.bind(key, receiverStageId, receiverStageId + 1); + return senderSort; + } + + private static boolean profile(WindowSortAutoTuner.Session session, int receiverStageId, int logicalOrdinal, + int inputHash, int collationHash) { + ExchangeKey key = key(logicalOrdinal, inputHash, collationHash); + session.bind(key, receiverStageId, receiverStageId + 1); + return session.shouldProfile(key); } private static boolean observe(WindowSortAutoTuner tuner, long queryHash, int receiverStageId, int senderStageId, int inputHash, int collationHash, List stats) { WindowSortAutoTuner.Session session = tuner.newSession(queryHash); - boolean decision = session.useSenderSort(receiverStageId, senderStageId, inputHash, collationHash); + boolean decision = choose(session, receiverStageId, senderStageId, inputHash, collationHash); session.observe(stats, 0); return decision; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalSortExchange.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalSortExchange.java index d7881eff7802..488efab26862 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalSortExchange.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalSortExchange.java @@ -76,7 +76,7 @@ public static PinotLogicalSortExchange create(RelNode input, RelDistribution dis isSortOnReceiver, null); } - /// Defers the global-window ordering choice until stage ids are assigned by the fragmenter. + /// Marks an AUTO exchange for resolution on the finalized logical tree, before stage allocation. public static PinotLogicalSortExchange createWindowAuto(RelNode input, RelDistribution distribution, RelCollation collation) { return create(input, distribution, PinotRelExchangeType.getDefaultExchangeType(), collation, false, true, true, diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRule.java index c4a1e8a7ac5e..3e02ed7617ab 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRule.java @@ -23,8 +23,8 @@ import java.util.EnumSet; import java.util.HashSet; import java.util.List; +import java.util.Locale; import java.util.Set; -import org.apache.calcite.plan.Context; import org.apache.calcite.plan.RelOptCluster; import org.apache.calcite.plan.RelOptPlanner; import org.apache.calcite.plan.RelOptRule; @@ -52,7 +52,6 @@ import org.apache.pinot.calcite.rel.logical.PinotKWayMergeSortExchange; import org.apache.pinot.calcite.rel.logical.PinotLogicalExchange; import org.apache.pinot.calcite.rel.logical.PinotLogicalSortExchange; -import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.query.context.PlannerContext; import org.apache.pinot.spi.utils.CommonConstants; @@ -133,8 +132,11 @@ public void onMatch(RelOptRuleCall call) { // When enabled, sort each sender explicitly and merge the sorted mailbox streams at the receiver. Otherwise, // retain the legacy post-exchange full-sort path. This switch supports rolling upgrades and rapid rollback. // TODO: Revisit whether we should use hash distribution - boolean sortOnSender = isWindowSortOnSenderEnabled(call); - if (sortOnSender) { + String mode = getWindowSortMode(call); + if ("auto".equals(mode)) { + exchange = PinotLogicalSortExchange.createWindowAuto(input, RelDistributions.hash(List.of()), + windowGroup.orderKeys); + } else if ("true".equals(mode)) { // An unbounded sender sort retains every row, independent of the broker response limit. RelNode orderedInput = LogicalSort.create(input, windowGroup.orderKeys, null, window.getCluster().getRexBuilder().makeExactLiteral(BigDecimal.valueOf(Integer.MAX_VALUE))); @@ -170,9 +172,10 @@ public void onMatch(RelOptRuleCall call) { // old servers sort the same rows twice during a rolling upgrade. Keep an explicit Sort only for post-exchange // full-sort paths. // PinotSortExchangeNodeInsertRule does not re-fire on it: its matches() rejects a Sort whose input is an - // exchange. Window sorts must retain every input row, independently of the broker response limit. + // exchange. Window sorts retain every input row, independently of the response limit. RelNode windowInput = exchange instanceof PinotLogicalSortExchange && !((PinotLogicalSortExchange) exchange).isSortOnReceiver() + && !((PinotLogicalSortExchange) exchange).isAutoWindowSort() ? LogicalSort.create(exchange, ((PinotLogicalSortExchange) exchange).getCollation(), null, window.getCluster().getRexBuilder().makeExactLiteral(BigDecimal.valueOf(Integer.MAX_VALUE))) : exchange; // NOTE: Need to create a new LogicalWindow to use the modified window group. @@ -180,20 +183,26 @@ public void onMatch(RelOptRuleCall call) { List.of(windowGroup))); } - private static boolean isWindowSortOnSenderEnabled(RelOptRuleCall call) { + private static String getWindowSortMode(RelOptRuleCall call) { RelOptPlanner planner = call.getPlanner(); - if (planner != null) { - Context context = planner.getContext(); + if (planner != null && planner.getContext() != null) { + PlannerContext context = planner.getContext().unwrap(PlannerContext.class); if (context != null) { - PlannerContext plannerContext = context.unwrap(PlannerContext.class); - if (plannerContext != null) { - return plannerContext.getEnvConfig().isKWayMergeSupported() - && QueryOptionsUtils.isWindowSortOnSender(plannerContext.getOptions(), - plannerContext.getEnvConfig().defaultWindowSortOnSender()); - } + String option = context.getOptions().get(CommonConstants.Broker.Request.QueryOptionKey.WINDOW_SORT_ON_SENDER); + String mode = validateWindowSortMode(option != null ? option + : context.getEnvConfig().defaultWindowSortOnSenderMode()); + return context.getEnvConfig().isKWayMergeSupported() && !context.isUsePhysicalOptimizer() ? mode : "false"; } } - return CommonConstants.Broker.DEFAULT_WINDOW_SORT_ON_SENDER; + return CommonConstants.Broker.DEFAULT_WINDOW_SORT_ON_SENDER_MODE; + } + + private static String validateWindowSortMode(String mode) { + String normalized = mode.toLowerCase(Locale.ROOT); + if (!"auto".equals(normalized) && !"true".equals(normalized) && !"false".equals(normalized)) { + throw new IllegalArgumentException("windowSortOnSender must be auto, true or false: " + mode); + } + return normalized; } private boolean isPartitionByOnlyQuery(Window.Group windowGroup) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java index 08f6d8f3f3e2..f66d4db538d3 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java @@ -43,6 +43,7 @@ import org.apache.calcite.prepare.CalciteCatalogReader; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelRoot; +import org.apache.calcite.rel.core.Exchange; import org.apache.calcite.rex.RexBuilder; import org.apache.calcite.rex.RexExecutor; import org.apache.calcite.runtime.CalciteContextException; @@ -85,6 +86,8 @@ import org.apache.pinot.query.planner.logical.PinotLogicalQueryPlanner; import org.apache.pinot.query.planner.logical.RelToPlanNodeConverter; import org.apache.pinot.query.planner.logical.TransformationTracker; +import org.apache.pinot.query.planner.logical.WindowSortAutoPlan; +import org.apache.pinot.query.planner.logical.WindowSortAutoPlanner; import org.apache.pinot.query.planner.physical.DispatchableSubPlan; import org.apache.pinot.query.planner.physical.PinotDispatchPlanner; import org.apache.pinot.query.planner.physical.v2.PRelNode; @@ -319,8 +322,9 @@ public CompiledQuery compile(String sqlQuery, SqlNodeAndOptions sqlNodeAndOption } else { queryNode = sqlNode; } - RelRoot relRoot = compileQuery(queryNode, plannerContext); - return new CompiledQuery(_envConfig.getDatabase(), sqlQuery, relRoot, plannerContext, sqlNodeAndOptions); + CompiledRelations compiled = compileQuery(queryNode, plannerContext); + return new CompiledQuery(_envConfig.getDatabase(), sqlQuery, compiled.root(), plannerContext, sqlNodeAndOptions, + compiled.selections()); } catch (QueryException e) { throw e; } catch (Throwable t) { @@ -396,17 +400,29 @@ public int getLiteModeEffectiveSortLimit() { // steps // -------------------------------------------------------------------------- - private RelRoot compileQuery(SqlNode sqlNode, PlannerContext plannerContext) { + private record CompiledRelations(RelRoot root, Map selections) { + } + + private CompiledRelations compileQuery(SqlNode sqlNode, PlannerContext plannerContext) { SqlNode validated = validate(sqlNode, plannerContext); RelRoot relation = toRelation(validated, plannerContext); RelNode optimized = optimize(relation, plannerContext); + String windowSortMode = plannerContext.getOptions().getOrDefault( + CommonConstants.Broker.Request.QueryOptionKey.WINDOW_SORT_ON_SENDER, + _envConfig.defaultWindowSortOnSenderMode()); + WindowSortAutoPlanner.Result resolved = _envConfig.isKWayMergeSupported() + && !plannerContext.isUsePhysicalOptimizer() && "auto".equalsIgnoreCase(windowSortMode) + ? WindowSortAutoPlanner.resolve(optimized, _envConfig.getWindowSortAutoPlan()) + : new WindowSortAutoPlanner.Result(optimized, Map.of()); + optimized = resolved.root(); if (plannerContext.isUsePhysicalOptimizer()) { Preconditions.checkNotNull(plannerContext.getPhysicalPlannerContext(), "Physical planner context is null"); optimized = RelToPRelConverter.toPRelNode(optimized, plannerContext.getPhysicalPlannerContext(), _envConfig.getTableCache()).unwrap(); PRelNodeTreeValidator.validate((PRelNode) optimized, plannerContext.getPhysicalPlannerContext()); } - return relation.withRel(optimized); + return new CompiledRelations(relation.withRel(optimized), resolved.selections()); } /// Query validation is a transformation from SqlNode to SqlNode where each node is validated. @@ -531,12 +547,9 @@ private RelNode optimize(RelRoot relRoot, PlannerContext plannerContext) { } } - private DispatchableSubPlan toDispatchableSubPlan(RelRoot relRoot, PlannerContext plannerContext) { - return toDispatchableSubPlan(relRoot, plannerContext, null); - } - private DispatchableSubPlan toDispatchableSubPlan(RelRoot relRoot, PlannerContext plannerContext, - @Nullable TransformationTracker.Builder tracker) { + @Nullable TransformationTracker.Builder tracker, + Map autoExchanges) { long requestId = _envConfig.getRequestId(); if (plannerContext.isUsePhysicalOptimizer()) { Pair plan = PinotLogicalQueryPlanner.makePlanV2(relRoot, @@ -546,7 +559,8 @@ private DispatchableSubPlan toDispatchableSubPlan(RelRoot relRoot, PlannerContex return pinotDispatchPlanner.createDispatchableSubPlanV2(plan.getLeft(), plan.getRight()); } SubPlan plan = PinotLogicalQueryPlanner.makePlan(relRoot, tracker, useSpools(plannerContext.getOptions()), - _envConfig.defaultHashFunction(), pruneUnnestColumns(plannerContext.getOptions())); + _envConfig.defaultHashFunction(), pruneUnnestColumns(plannerContext.getOptions()), + _envConfig.getWindowSortAutoPlan(), autoExchanges); PinotDispatchPlanner pinotDispatchPlanner = new PinotDispatchPlanner(plannerContext, _envConfig.getWorkerManager(), _envConfig.getRequestId(), _envConfig.getTableCache()); @@ -970,6 +984,16 @@ default boolean isKWayMergeSupported() { default boolean defaultWindowSortOnSender() { return CommonConstants.Broker.DEFAULT_WINDOW_SORT_ON_SENDER; } + + /// Per-exchange adaptive mode is enabled explicitly; boolean defaults remain supported. + @Value.Default + default String defaultWindowSortOnSenderMode() { + return Boolean.toString(defaultWindowSortOnSender()); + } + + /// Optional broker-owned selector. Standalone planners conservatively use the receiver sort. + @Nullable + WindowSortAutoPlan getWindowSortAutoPlan(); } /// A query that have been parsed, validates, transformed into a [RelNode] and optimized with Calcite. @@ -989,9 +1013,12 @@ public class CompiledQuery implements Closeable { private final PlannerContext _plannerContext; private final SqlNodeAndOptions _sqlNodeAndOptions; private final Set _tableNames; + private final Map _autoExchanges; private CompiledQuery(String database, String textQuery, RelRoot relRoot, PlannerContext plannerContext, - SqlNodeAndOptions sqlNodeAndOptions) { + SqlNodeAndOptions sqlNodeAndOptions, + Map autoExchanges) { + _autoExchanges = autoExchanges; _database = database; _textQuery = textQuery; _relRoot = relRoot; @@ -1028,7 +1055,8 @@ public QueryEnvironment.QueryPlannerResult explain(long requestId, SqlExplainFormat format = _plannerContext.getSqlExplainFormat(); if (explain instanceof SqlPhysicalExplain) { // get the physical plan for query. - DispatchableSubPlan dispatchableSubPlan = toDispatchableSubPlan(_relRoot, _plannerContext); + DispatchableSubPlan dispatchableSubPlan = + toDispatchableSubPlan(_relRoot, _plannerContext, null, _autoExchanges); return getQueryPlannerResult(_plannerContext, dispatchableSubPlan, PhysicalExplainPlanVisitor.explain(dispatchableSubPlan), dispatchableSubPlan.getTableNames()); } else { @@ -1044,7 +1072,8 @@ public QueryEnvironment.QueryPlannerResult explain(long requestId, // Build the dispatchable subplan even though only the logical plan is rendered: some planning errors are // raised while converting the plan or assigning workers (e.g. a `tableOptions` partition hint that // disagrees with the table's actual partitioning), and EXPLAIN must fail wherever execution would. - DispatchableSubPlan dispatchableSubPlan = toDispatchableSubPlan(_relRoot, _plannerContext); + DispatchableSubPlan dispatchableSubPlan = + toDispatchableSubPlan(_relRoot, _plannerContext, null, _autoExchanges); return getQueryPlannerResult(_plannerContext, dispatchableSubPlan, PlannerUtils.explainPlan(_relRoot.rel, format, level), dispatchableSubPlan.getTableNames()); } else { @@ -1056,7 +1085,7 @@ public QueryEnvironment.QueryPlannerResult explain(long requestId, new TransformationTracker.ByIdentity.Builder<>(); // Transform RelNodes into DispatchableSubPlan DispatchableSubPlan dispatchableSubPlan = - toDispatchableSubPlan(_relRoot, _plannerContext, nodeTracker); + toDispatchableSubPlan(_relRoot, _plannerContext, nodeTracker, _autoExchanges); AskingServerStageExplainer serversExplainer = new AskingServerStageExplainer( onServerExplainer, explainPlanVerbose, RelBuilder.create(_config)); @@ -1081,7 +1110,8 @@ public QueryPlannerResult planQuery(long requestId) { // TODO: current code only assume one SubPlan per query, but we should support multiple SubPlans per query. // Each SubPlan should be able to run independently from Broker then set the results into the dependent // SubPlan for further processing. - DispatchableSubPlan dispatchableSubPlan = toDispatchableSubPlan(_relRoot, _plannerContext); + DispatchableSubPlan dispatchableSubPlan = + toDispatchableSubPlan(_relRoot, _plannerContext, null, _autoExchanges); return getQueryPlannerResult(_plannerContext, dispatchableSubPlan, null, dispatchableSubPlan.getTableNames()); } catch (QueryException e) { throw e; @@ -1111,6 +1141,11 @@ public Map getOptions() { return _sqlNodeAndOptions.getOptions(); } + @Nullable + public WindowSortAutoPlan getWindowSortAutoPlan() { + return _envConfig.getWindowSortAutoPlan(); + } + public RelRoot getRelRoot() { return _relRoot; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeMerger.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeMerger.java index 8128924909ba..f11ab0426ce9 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeMerger.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeMerger.java @@ -292,7 +292,8 @@ public PlanNode visitMailboxReceive(MailboxReceiveNode node, PlanNode context) { if (node.isSort() != otherNode.isSort()) { return null; } - if (node.isSortedOnSender() != otherNode.isSortedOnSender()) { + if (node.isSortedOnSender() != otherNode.isSortedOnSender() + || node.isAutoProfile() != otherNode.isAutoProfile()) { return null; } List children = mergeChildren(node, context); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinder.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinder.java index 448c87d6ec53..fd447f095cd1 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinder.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinder.java @@ -233,6 +233,7 @@ public Boolean visitMailboxReceive(MailboxReceiveNode node1, PlanNode node2) { // require some extra checks or transformation on the spooling logic. We are not doing that for now. && node1.isSort() == that.isSort() && node1.isSortedOnSender() == that.isSortedOnSender() + && node1.isAutoProfile() == that.isAutoProfile() && Objects.equals(node1.getCollations(), that.getCollations()) && node1.getExchangeType() == that.getExchangeType(); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java index e01301b9852e..e9a86706696e 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java @@ -32,6 +32,7 @@ import org.apache.calcite.rel.RelDistribution; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelRoot; +import org.apache.calcite.rel.core.Exchange; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType; import org.apache.pinot.query.context.PhysicalPlannerContext; @@ -58,10 +59,18 @@ private PinotLogicalQueryPlanner() { public static SubPlan makePlan(RelRoot relRoot, @Nullable TransformationTracker.Builder tracker, boolean useSpools, String hashFunction, boolean pruneUnnestColumns) { - PlanNode rootNode = new RelToPlanNodeConverter(tracker, hashFunction, - !CommonConstants.Helix.DEFAULT_ENABLE_CASE_INSENSITIVE, pruneUnnestColumns).toPlanNode(relRoot.rel); + return makePlan(relRoot, tracker, useSpools, hashFunction, pruneUnnestColumns, null, Map.of()); + } - PlanFragment rootFragment = planNodeToPlanFragment(rootNode, tracker, useSpools, hashFunction); + public static SubPlan makePlan(RelRoot relRoot, + @Nullable TransformationTracker.Builder tracker, boolean useSpools, + String hashFunction, boolean pruneUnnestColumns, @Nullable WindowSortAutoPlan selector, + Map autoExchanges) { + RelToPlanNodeConverter converter = new RelToPlanNodeConverter(tracker, hashFunction, + !CommonConstants.Helix.DEFAULT_ENABLE_CASE_INSENSITIVE, pruneUnnestColumns, autoExchanges); + PlanNode rootNode = converter.toPlanNode(relRoot.rel); + PlanFragment rootFragment = planNodeToPlanFragment(rootNode, tracker, useSpools, hashFunction, selector, + converter.getAutoBindings()); return new SubPlan(rootFragment, new SubPlanMetadata(RelToPlanNodeConverter.getTableNamesFromRelRoot(relRoot.rel), relRoot.fields), List.of()); @@ -107,8 +116,9 @@ public static Pair makePlanV2( private static PlanFragment planNodeToPlanFragment( PlanNode node, @Nullable TransformationTracker.Builder tracker, boolean useSpools, - String hashFunction) { - PlanFragmenter fragmenter = new PlanFragmenter(); + String hashFunction, @Nullable WindowSortAutoPlan selector, + Map autoBindings) { + PlanFragmenter fragmenter = new PlanFragmenter(selector, autoBindings); PlanFragmenter.Context fragmenterContext = fragmenter.createContext(); node = node.visit(fragmenter, fragmenterContext); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java index 9ffe129f4ada..561a614d7285 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java @@ -24,6 +24,8 @@ import java.util.ArrayList; import java.util.IdentityHashMap; import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; import org.apache.calcite.rel.RelDistribution; import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType; import org.apache.pinot.query.planner.PlanFragment; @@ -66,6 +68,20 @@ public class PlanFragmenter implements PlanNodeVisitor _mailboxReceiveToExchangeNodeMap = new IdentityHashMap<>(); + @Nullable + private final WindowSortAutoPlan _autoSelector; + private final Map _autoBindings; + + public PlanFragmenter() { + this(null, Map.of()); + } + + public PlanFragmenter(@Nullable WindowSortAutoPlan autoSelector, + Map autoBindings) { + _autoSelector = autoSelector; + _autoBindings = autoBindings; + } + // ROOT PlanFragment ID is 0, current PlanFragment ID starts with 1, next PlanFragment ID starts with 2. private int _nextPlanFragmentId = 2; @@ -185,6 +201,10 @@ public PlanNode visitExchange(ExchangeNode node, Context context) { int senderPlanFragmentId = _nextPlanFragmentId++; _childPlanFragmentIdsMap.computeIfAbsent(receiverPlanFragmentId, k -> new IntArrayList()).add(senderPlanFragmentId); + WindowSortAutoPlanner.Selection auto = _autoBindings.get(node); + if (auto != null && _autoSelector != null) { + _autoSelector.bind(auto.key(), receiverPlanFragmentId, senderPlanFragmentId); + } // Create a new context for the next PlanFragment with MailboxSendNode as the root node. PlanNode nextPlanFragmentRoot = node.getInputs().get(0).visit(this, new Context(senderPlanFragmentId)); PinotRelExchangeType exchangeType = node.getExchangeType(); @@ -204,8 +224,9 @@ public PlanNode visitExchange(ExchangeNode node, Context context) { senderPlanFragmentId, exchangeType, distributionType, keys, node.getCollations(), ((KWayMergeExchangeNode) node).getFetch(), ((KWayMergeExchangeNode) node).getOffset(), mailboxSendNode) : new MailboxReceiveNode(receiverPlanFragmentId, nextPlanFragmentRoot.getDataSchema(), - senderPlanFragmentId, exchangeType, distributionType, keys, node.getCollations(), node.isSortOnReceiver(), - node.isSortOnSender(), mailboxSendNode); + senderPlanFragmentId, exchangeType, distributionType, keys, + auto != null ? auto.collations() : node.getCollations(), node.isSortOnReceiver(), + node.isSortOnSender(), auto != null && auto.profile(), mailboxSendNode); _mailboxReceiveToExchangeNodeMap.put(mailboxReceiveNode, node); return mailboxReceiveNode; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java index e536ce103564..195baeab2774 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java @@ -22,7 +22,9 @@ import com.google.common.collect.Sets; import java.util.ArrayList; import java.util.Arrays; +import java.util.IdentityHashMap; import java.util.List; +import java.util.Map; import java.util.Set; import javax.annotation.Nullable; import org.apache.calcite.plan.RelOptTable; @@ -108,6 +110,8 @@ public final class RelToPlanNodeConverter { private final boolean _caseSensitive; // When true, UNNEST output is pruned to drop input (passthrough) columns not referenced downstream. Default off. private final boolean _pruneUnnestColumns; + private final Map _autoExchanges; + private final Map _autoBindings = new IdentityHashMap<>(); public RelToPlanNodeConverter(@Nullable TransformationTracker.Builder tracker, String hashFunction) { @@ -121,12 +125,23 @@ public RelToPlanNodeConverter(@Nullable TransformationTracker.Builder tracker, String hashFunction, boolean caseSensitive, boolean pruneUnnestColumns) { + this(tracker, hashFunction, caseSensitive, pruneUnnestColumns, Map.of()); + } + + public RelToPlanNodeConverter(@Nullable TransformationTracker.Builder tracker, + String hashFunction, boolean caseSensitive, boolean pruneUnnestColumns, + Map autoExchanges) { + _autoExchanges = autoExchanges; _tracker = tracker; _hashFunction = hashFunction; _caseSensitive = caseSensitive; _pruneUnnestColumns = pruneUnnestColumns; } + public Map getAutoBindings() { + return _autoBindings; + } + /// Converts a [RelNode] into its serializable counterpart. /// NOTE: Stage ID is not determined yet. public PlanNode toPlanNode(RelNode node) { @@ -147,6 +162,10 @@ public PlanNode toPlanNode(RelNode node) { result = convertLogicalSort((LogicalSort) node); } else if (node instanceof Exchange) { result = convertLogicalExchange((Exchange) node); + WindowSortAutoPlanner.Selection selection = _autoExchanges.get(node); + if (selection != null) { + _autoBindings.put((ExchangeNode) result, selection); + } } else if (node instanceof LogicalJoin) { _brokerMetrics.addMeteredGlobalValue(BrokerMeter.JOIN_COUNT, 1); if (!_joinFound) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlan.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlan.java index 5c61de5ecff6..8122265f3f0c 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlan.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlan.java @@ -19,16 +19,21 @@ package org.apache.pinot.query.planner.logical; -/// Selects one ordering strategy for every worker of a global window exchange. -/// -/// The receiver and sender stage identifiers are stable within one planned query. Implementations must return false -/// when evidence for this exchange is absent or ambiguous. A selector is shared only during planning and must not -/// mutate the physical plan after dispatch. +/// Selects ordering for global windows on the finalized logical tree, before stage allocation. +/// Implementations must fail closed when evidence is absent or ambiguous. The logical key remains stable across +/// repeated compilations; allocated stage identifiers are bound separately for one query's observations. public interface WindowSortAutoPlan { - boolean useSenderSort(int receiverStageId, int senderStageId, int inputHash, int collationHash); + boolean useSenderSort(ExchangeKey key); - /// Whether a receiver-sort execution should collect a fresh AUTO sample. The default preserves existing selectors. - default boolean shouldProfile(int receiverStageId, int senderStageId, int inputHash, int collationHash) { + default boolean shouldProfile(ExchangeKey key) { return true; } + + /// Binds an already selected logical exchange to this query's receiver stage. It cannot change the decision. + default void bind(ExchangeKey key, int receiverStageId, int senderStageId) { + } + + /// A deterministic path in the finalized logical tree plus ordering/input fingerprints, without stage identifiers. + record ExchangeKey(String path, int inputHash, int collationHash) { + } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlanner.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlanner.java new file mode 100644 index 000000000000..d5dc57fc74e5 --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/WindowSortAutoPlanner.java @@ -0,0 +1,99 @@ +/** + * 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.planner.logical; + +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.calcite.plan.RelOptUtil; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.Exchange; +import org.apache.calcite.rel.logical.LogicalSort; +import org.apache.pinot.calcite.rel.logical.PinotKWayMergeSortExchange; +import org.apache.pinot.calcite.rel.logical.PinotLogicalExchange; +import org.apache.pinot.calcite.rel.logical.PinotLogicalSortExchange; + + +/// Resolves AUTO on a finalized logical tree. This per-compilation pass is not shared between threads. +/// Deterministic tree paths separate exchanges before any stage allocation; the returned identity map binds only +/// this query's final logical exchanges to their observations during plain fragmentation. +public final class WindowSortAutoPlanner { + private WindowSortAutoPlanner() { + } + + public static Result resolve(RelNode root, @Nullable WindowSortAutoPlan selector) { + Map selections = new IdentityHashMap<>(); + return new Result(rewrite(root, "root", selector, selections), selections); + } + + private static RelNode rewrite(RelNode node, String path, @Nullable WindowSortAutoPlan selector, + Map selections) { + PinotLogicalSortExchange auto = node instanceof PinotLogicalSortExchange + && ((PinotLogicalSortExchange) node).isAutoWindowSort() ? (PinotLogicalSortExchange) node : null; + WindowSortAutoPlan.ExchangeKey key = auto != null + ? new WindowSortAutoPlan.ExchangeKey(path, RelOptUtil.toString(auto.getInput()).hashCode(), + auto.getCollation().hashCode()) : null; + boolean senderSort = key != null && selector != null && selector.useSenderSort(key); + boolean profile = key != null && selector != null && !senderSort && selector.shouldProfile(key); + List inputs = new ArrayList<>(); + boolean changed = false; + for (int i = 0; i < node.getInputs().size(); i++) { + RelNode original = node.getInput(i); + RelNode rewritten = rewrite(original, path + "/" + i, selector, selections); + inputs.add(rewritten); + changed |= original != rewritten; + } + if (auto == null) { + return changed ? node.copy(node.getTraitSet(), inputs) : node; + } + RelNode input = inputs.get(0); + Exchange exchange; + RelNode result; + if (senderSort) { + RelNode sorted = unboundedSort(input, auto); + exchange = PinotKWayMergeSortExchange.create(sorted, auto.getDistribution(), auto.getCollation(), + auto.getPrePartitioned()); + result = exchange; + } else { + exchange = PinotLogicalExchange.create(input, auto.getDistribution(), auto.getExchangeType(), + auto.getPrePartitioned()); + result = unboundedSort(exchange, auto); + } + if (selector != null) { + selections.put(exchange, new Selection(key, profile, auto.getCollation().getFieldCollations())); + } + return result; + } + + private static LogicalSort unboundedSort(RelNode input, PinotLogicalSortExchange auto) { + return LogicalSort.create(input, auto.getCollation(), null, + input.getCluster().getRexBuilder().makeExactLiteral(BigDecimal.valueOf(Integer.MAX_VALUE))); + } + + public record Selection(WindowSortAutoPlan.ExchangeKey key, boolean profile, + List collations) { + } + + public record Result(RelNode root, Map selections) { + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/ExchangeNode.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/ExchangeNode.java index 463f5131fd44..7a1eb15a5494 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/ExchangeNode.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/ExchangeNode.java @@ -41,7 +41,6 @@ public class ExchangeNode extends BasePlanNode { private final List _collations; private final boolean _sortOnSender; private final boolean _sortOnReceiver; - private final boolean _autoWindowSort; // Table names should be set for SUB_PLAN exchange type. private final Set _tableNames; @Nullable @@ -52,15 +51,6 @@ public ExchangeNode(int stageId, DataSchema dataSchema, List inputs, P RelDistribution.Type distributionType, @Nullable List keys, boolean prePartitioned, @Nullable List collations, boolean sortOnSender, boolean sortOnReceiver, @Nullable Set tableNames, ExchangeStrategy exchangeStrategy, String hashFunction) { - this(stageId, dataSchema, inputs, exchangeType, distributionType, keys, prePartitioned, collations, sortOnSender, - sortOnReceiver, false, tableNames, exchangeStrategy, hashFunction); - } - - public ExchangeNode(int stageId, DataSchema dataSchema, List inputs, PinotRelExchangeType exchangeType, - RelDistribution.Type distributionType, @Nullable List keys, boolean prePartitioned, - @Nullable List collations, boolean sortOnSender, boolean sortOnReceiver, - boolean autoWindowSort, @Nullable Set tableNames, ExchangeStrategy exchangeStrategy, - String hashFunction) { super(stageId, dataSchema, null, inputs); _exchangeType = exchangeType; _distributionType = distributionType; @@ -69,7 +59,6 @@ public ExchangeNode(int stageId, DataSchema dataSchema, List inputs, P _collations = collations; _sortOnSender = sortOnSender; _sortOnReceiver = sortOnReceiver; - _autoWindowSort = autoWindowSort; _tableNames = tableNames; _exchangeStrategy = exchangeStrategy; _hashFunction = hashFunction; @@ -105,10 +94,6 @@ public boolean isSortOnReceiver() { return _sortOnReceiver; } - public boolean isAutoWindowSort() { - return _autoWindowSort; - } - @Nullable public Set getTableNames() { return _tableNames; @@ -136,7 +121,7 @@ public T visit(PlanNodeVisitor visitor, C context) { @Override public PlanNode withInputs(List inputs) { return new ExchangeNode(_stageId, _dataSchema, inputs, _exchangeType, _distributionType, _keys, _prePartitioned, - _collations, _sortOnSender, _sortOnReceiver, _autoWindowSort, _tableNames, null, _hashFunction); + _collations, _sortOnSender, _sortOnReceiver, _tableNames, null, _hashFunction); } @Override @@ -152,7 +137,6 @@ public boolean equals(Object o) { } ExchangeNode that = (ExchangeNode) o; return _sortOnSender == that._sortOnSender && _sortOnReceiver == that._sortOnReceiver - && _autoWindowSort == that._autoWindowSort && _prePartitioned == that._prePartitioned && _exchangeType == that._exchangeType && _distributionType == that._distributionType && Objects.equals(_keys, that._keys) && Objects.equals(_collations, that._collations) && Objects.equals(_tableNames, that._tableNames) @@ -162,6 +146,6 @@ public boolean equals(Object o) { @Override public int hashCode() { return Objects.hash(super.hashCode(), _exchangeType, _distributionType, _keys, _sortOnSender, _sortOnReceiver, - _autoWindowSort, _prePartitioned, _collations, _tableNames, _hashFunction); + _prePartitioned, _collations, _tableNames, _hashFunction); } } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRuleTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRuleTest.java index 51e15187c58a..1380bc4e0427 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRuleTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRuleTest.java @@ -67,10 +67,10 @@ public class PinotWindowExchangeNodeInsertRuleTest { @BeforeMethod public void setUp() { _mocks = MockitoAnnotations.openMocks(this); - when(_cluster.getRexBuilder()).thenReturn(REX_BUILDER); RelTraitSet traits = RelTraitSet.createEmpty(); Mockito.when(_input.getTraitSet()).thenReturn(traits); Mockito.when(_input.getCluster()).thenReturn(_cluster); + when(_cluster.getRexBuilder()).thenReturn(REX_BUILDER); } @AfterMethod diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java index 202839721f3e..01dc72216527 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java @@ -27,14 +27,22 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; import java.util.stream.Collectors; import org.apache.calcite.rel.RelDistribution; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.Window; +import org.apache.calcite.rel.logical.LogicalSort; import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.pinot.calcite.rel.logical.PinotKWayMergeSortExchange; +import org.apache.pinot.calcite.rel.logical.PinotLogicalExchange; import org.apache.pinot.query.QueryEnvironment.CompiledQuery; import org.apache.pinot.query.planner.PlannerUtils; +import org.apache.pinot.query.planner.logical.WindowSortAutoPlan; import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; import org.apache.pinot.query.planner.physical.DispatchableSubPlan; import org.apache.pinot.query.planner.plannode.AggregateNode; @@ -1331,7 +1339,7 @@ public void testGlobalOrderedWindowSenderHasExplicitMatchingSortInput() { @Test public void testGlobalOrderedWindowSenderSortIsDisabledByDefault() { - DispatchableSubPlan plan = _queryEnvironment.planQuery( + DispatchableSubPlan plan = planWindowQuery(_queryEnvironment, "SELECT col1, SUM(col3) OVER (ORDER BY col3) FROM d"); MailboxSendNode sendNode = findWindowInputSendNode(plan); @@ -1339,8 +1347,6 @@ public void testGlobalOrderedWindowSenderSortIsDisabledByDefault() { assertFalse(sendNode.getInputs().get(0) instanceof SortNode); assertTrue(findWindowNode(plan).getInputs().get(0) instanceof SortNode, "The disabled path must retain the legacy post-exchange full sort"); - assertEquals(((SortNode) findWindowNode(plan).getInputs().get(0)).getFetch(), Integer.MAX_VALUE, - "The disabled path must retain the complete window input"); } @Test @@ -1349,12 +1355,174 @@ public void testQueryOptionCannotBypassMergeCapability() { try (CompiledQuery compiled = _queryEnvironment.compile(query)) { WindowNode window = findWindowNode(compiled.planQuery(0).getQueryPlan()); assertTrue(window.getInputs().get(0) instanceof SortNode); - assertEquals(((SortNode) window.getInputs().get(0)).getFetch(), Integer.MAX_VALUE, - "Compatibility fallback must retain the complete window input"); assertFalse(window.getInputs().get(0).getInputs().get(0) instanceof MailboxMergeReceiveNode); } } + @Test + public void testAutoColdPlanIsResolvedBeforeFragmentation() { + QueryEnvironment environment = windowAutoEnvironment(key -> false, true); + try (CompiledQuery compiled = environment.compile("SELECT col1, SUM(col3) OVER (ORDER BY col3) FROM d")) { + RelNode logicalInput = findLogicalWindow(compiled.getRelRoot().rel).getInput(); + assertTrue(logicalInput instanceof LogicalSort); + assertEquals(RexLiteral.intValue(((LogicalSort) logicalInput).fetch), Integer.MAX_VALUE); + assertTrue(logicalInput.getInput(0) instanceof PinotLogicalExchange); + DispatchableSubPlan plan = compiled.planQuery(0).getQueryPlan(); + SortNode sort = (SortNode) findWindowNode(plan).getInputs().get(0); + assertEquals(sort.getFetch(), Integer.MAX_VALUE, "Window input must outlive the broker response cap"); + MailboxReceiveNode receive = (MailboxReceiveNode) sort.getInputs().get(0); + assertFalse(receive instanceof MailboxMergeReceiveNode); + assertTrue(receive.isAutoProfile()); + assertFalse(findWindowInputSendNode(plan).getInputs().get(0) instanceof SortNode); + assertFalse(findWindowInputSendNode(plan).isSort()); + } + } + + @Test + public void testSupportedStandaloneAutoDoesNotProfile() { + QueryEnvironment environment = getQueryEnvironment(3, 1, 2, TABLE_SCHEMAS, SERVER1_SEGMENTS, + SERVER2_SEGMENTS, PARTITIONED_SEGMENTS_MAP, true); + DispatchableSubPlan plan = planWindowQuery(environment, + "SET windowSortOnSender='auto'; SELECT col1, SUM(col3) OVER (ORDER BY col3) FROM d"); + SortNode sort = (SortNode) findWindowNode(plan).getInputs().get(0); + assertEquals(sort.getFetch(), Integer.MAX_VALUE); + MailboxReceiveNode receive = (MailboxReceiveNode) sort.getInputs().get(0); + assertFalse(receive instanceof MailboxMergeReceiveNode); + assertFalse(receive.isAutoProfile(), "A standalone planner has no tuner to consume samples"); + } + + @Test + public void testAutoWarmPlanIsResolvedBeforeFragmentation() { + QueryEnvironment environment = windowAutoEnvironment(key -> true, true); + try (CompiledQuery compiled = environment.compile("SELECT col1, SUM(col3) OVER (ORDER BY col3) FROM d")) { + RelNode logicalInput = findLogicalWindow(compiled.getRelRoot().rel).getInput(); + assertTrue(logicalInput instanceof PinotKWayMergeSortExchange); + assertTrue(logicalInput.getInput(0) instanceof LogicalSort); + assertEquals(RexLiteral.intValue(((LogicalSort) logicalInput.getInput(0)).fetch), Integer.MAX_VALUE); + DispatchableSubPlan plan = compiled.planQuery(0).getQueryPlan(); + MailboxMergeReceiveNode receive = (MailboxMergeReceiveNode) findWindowNode(plan).getInputs().get(0); + assertFalse(receive.isAutoProfile()); + MailboxSendNode send = findWindowInputSendNode(plan); + assertFalse(send.isSort()); + assertTrue(send.getInputs().get(0) instanceof SortNode); + assertEquals(((SortNode) send.getInputs().get(0)).getFetch(), Integer.MAX_VALUE); + } + } + + @Test + public void testAutoCannotSelectOrProfileWithoutCapability() { + AtomicInteger calls = new AtomicInteger(); + QueryEnvironment environment = windowAutoEnvironment(key -> { + calls.incrementAndGet(); + return true; + }, false); + for (String prefix : List.of("", "SET windowSortOnSender='auto'; ", "SET windowSortOnSender=true; ")) { + DispatchableSubPlan plan = planWindowQuery(environment, + prefix + "SELECT col1, SUM(col3) OVER (ORDER BY col3) FROM d"); + SortNode sort = (SortNode) findWindowNode(plan).getInputs().get(0); + assertEquals(sort.getFetch(), Integer.MAX_VALUE); + MailboxReceiveNode receive = (MailboxReceiveNode) sort.getInputs().get(0); + assertFalse(receive instanceof MailboxMergeReceiveNode); + assertFalse(receive.isAutoProfile()); + } + assertEquals(calls.get(), 0); + } + + @Test + public void testExplicitWindowModesNeverTrainAuto() { + AtomicInteger calls = new AtomicInteger(); + QueryEnvironment environment = windowAutoEnvironment(key -> { + calls.incrementAndGet(); + return true; + }, true); + for (String mode : List.of("true", "false")) { + DispatchableSubPlan plan = planWindowQuery(environment, + "SET windowSortOnSender=" + mode + "; SELECT col1, SUM(col3) OVER (ORDER BY col3) FROM d"); + MailboxReceiveNode receive = findNodeOfType(findWindowNode(plan), MailboxReceiveNode.class); + assertFalse(receive.isAutoProfile()); + } + assertEquals(calls.get(), 0); + } + + @Test + public void testAutoLogicalKeysStableAndBoundToAllocatedStages() { + String query = "SELECT SUM(col3) OVER (ORDER BY col3), SUM(col3) OVER (ORDER BY col7) FROM d"; + List first = new ArrayList<>(); + List second = new ArrayList<>(); + Map bindings = new HashMap<>(); + WindowSortAutoPlan selector = new WindowSortAutoPlan() { + @Override + public boolean useSenderSort(ExchangeKey key) { + first.add(key); + return first.size() == 1; + } + + @Override + public void bind(ExchangeKey key, int receiverStageId, int senderStageId) { + bindings.put(key, receiverStageId); + } + }; + DispatchableSubPlan plan; + try (CompiledQuery compiled = windowAutoEnvironment(selector, true).compile(query)) { + assertEquals(first.size(), 2, "Selections must exist before physical stages are allocated"); + assertTrue(bindings.isEmpty()); + plan = compiled.planQuery(0).getQueryPlan(); + } + try (CompiledQuery compiled = windowAutoEnvironment(key -> { + second.add(key); + return false; + }, true).compile(query)) { + compiled.planQuery(1); + } + assertEquals(second, first, "Choosing warm/cold must not change any logical evidence key"); + assertNotEquals(first.get(0), first.get(1)); + assertEquals(bindings.size(), 2); + int merges = 0; + int profiled = 0; + for (DispatchablePlanFragment fragment : plan.getQueryStages()) { + MailboxReceiveNode receive = findNodeOfType(fragment.getPlanFragment().getFragmentRoot(), + MailboxReceiveNode.class); + if (receive instanceof MailboxMergeReceiveNode) { + merges++; + } else if (receive != null && receive.isAutoProfile()) { + profiled++; + assertTrue(bindings.containsValue(receive.getStageId())); + } + } + assertEquals(merges, 1); + assertEquals(profiled, 1); + } + + @Test + public void testInvalidAutoOptionRejectedEvenWithoutCapability() { + expectThrows(Exception.class, () -> _queryEnvironment.compile( + "SET windowSortOnSender='invalid'; SELECT SUM(col3) OVER (ORDER BY col3) FROM d")); + } + + private QueryEnvironment windowAutoEnvironment(WindowSortAutoPlan selector, boolean supported) { + return getQueryEnvironment(3, 1, 2, TABLE_SCHEMAS, SERVER1_SEGMENTS, SERVER2_SEGMENTS, + PARTITIONED_SEGMENTS_MAP, supported, selector); + } + + private static DispatchableSubPlan planWindowQuery(QueryEnvironment environment, String query) { + try (CompiledQuery compiled = environment.compile(query)) { + return compiled.planQuery(0).getQueryPlan(); + } + } + + private static Window findLogicalWindow(RelNode root) { + if (root instanceof Window) { + return (Window) root; + } + for (RelNode input : root.getInputs()) { + Window window = findLogicalWindow(input); + if (window != null) { + return window; + } + } + return null; + } + @DataProvider public Object[][] orderedLeafCases() { return new Object[][]{ diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java index c72c8fac069a..b4b982a9800b 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java @@ -34,6 +34,7 @@ import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo; import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo.PartitionInfo; +import org.apache.pinot.query.planner.logical.WindowSortAutoPlan; import org.apache.pinot.query.routing.WorkerManager; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.Schema; @@ -311,6 +312,14 @@ public static QueryEnvironment getQueryEnvironment(int reducerPort, int port1, i protected static QueryEnvironment getQueryEnvironment(int reducerPort, int port1, int port2, Map schemaMap, Map> segmentMap1, Map> segmentMap2, @Nullable Map>>> partitionedSegmentsMap, boolean mergeSupported) { + return getQueryEnvironment(reducerPort, port1, port2, schemaMap, segmentMap1, segmentMap2, partitionedSegmentsMap, + mergeSupported, null); + } + + protected static QueryEnvironment getQueryEnvironment(int reducerPort, int port1, int port2, + Map schemaMap, Map> segmentMap1, Map> segmentMap2, + @Nullable Map>>> partitionedSegmentsMap, boolean mergeSupported, + @Nullable WindowSortAutoPlan windowSortAutoPlan) { MockRoutingManagerFactory factory = new MockRoutingManagerFactory(port1, port2); for (Map.Entry entry : schemaMap.entrySet()) { factory.registerTable(entry.getValue(), entry.getKey()); @@ -351,7 +360,9 @@ protected static QueryEnvironment getQueryEnvironment(int reducerPort, int port1 return new QueryEnvironment(QueryEnvironment.configBuilder().requestId(-1L) .database(CommonConstants.DEFAULT_DATABASE).tableCache(tableCache).isNullHandlingEnabled(true) .workerManager(new WorkerManager("Broker_localhost", "localhost", reducerPort, routingManager)) - .isKWayMergeSupported(mergeSupported).build()); + .isKWayMergeSupported(mergeSupported) + .defaultWindowSortOnSenderMode(windowSortAutoPlan != null ? "auto" : "false") + .windowSortAutoPlan(windowSortAutoPlan).build()); } /// JSON test case definition for query planner test cases. Tables and schemas will come from those already defined diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java index 84297e8e0b32..48d389239610 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java @@ -37,10 +37,12 @@ import org.apache.pinot.query.planner.plannode.EnrichedJoinNode; import org.apache.pinot.query.planner.plannode.JoinNode; import org.apache.pinot.query.planner.plannode.MailboxMergeReceiveNode; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.planner.plannode.UnnestNode; import org.testng.annotations.Test; +import static org.apache.pinot.calcite.rel.logical.PinotRelExchangeType.STREAMING; import static org.testng.Assert.assertEquals; import static org.testng.Assert.expectThrows; @@ -72,6 +74,19 @@ public void testQueryStagePlanSerDe(String query) { } } + @Test + public void testAutoProfileReceiveNodeSerDe() { + DataSchema schema = new DataSchema(new String[]{"sort_key"}, new ColumnDataType[]{ColumnDataType.INT}); + MailboxReceiveNode receiveNode = new MailboxReceiveNode(1, schema, 2, STREAMING, + RelDistribution.Type.SINGLETON, List.of(), List.of(new RelFieldCollation(0)), false, false, true, null); + + MailboxReceiveNode deserialized = + (MailboxReceiveNode) PlanNodeDeserializer.process(PlanNodeSerializer.process(receiveNode)); + assertEquals(deserialized, receiveNode); + assertEquals(deserialized.isAutoProfile(), true, + "The AUTO profile flag must survive broker-to-server serialization"); + } + @Test public void testPrunedUnnestNodeSerDe() { // Round-trips the passthrough-pruning wire fields (passthroughInputIndexes, prunedPassthrough). A non-sequential diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java index 103315475ec7..685992bf298b 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java @@ -1188,10 +1188,10 @@ public void testSumWithUnboundedPrecedingLowerAndUnboundedFollowingUpper(WindowN assertTrue(operator.nextBlock().isSuccess(), "Second block is EOS (done processing)"); } - /// The global window receiver sort establishes window ordering without applying the broker response cap to its input. + /// The cold AUTO receiver sort establishes window ordering without applying the broker response cap to its input. /// Following frames and RANGE peers both need rows beyond that cap, even when the sort emits several blocks. @Test(dataProvider = "windowFrameTypes") - public void testReceiverSortRetainsCompleteWindowInput(WindowNode.WindowFrameType frameType) { + public void testColdReceiverSortRetainsCompleteWindowInput(WindowNode.WindowFrameType frameType) { DataSchema inputSchema = new DataSchema(new String[]{"key", "value"}, new ColumnDataType[]{INT, INT}); DataSchema resultSchema = new DataSchema(new String[]{"key", "value", "sum"}, new ColumnDataType[]{INT, INT, DOUBLE}); diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 5e6647d29663..b19af8f5906c 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -823,6 +823,7 @@ public static class Broker { public static final String CONFIG_OF_WINDOW_SORT_ON_SENDER = "pinot.broker.multistage.window.sort.on.sender"; public static final boolean DEFAULT_WINDOW_SORT_ON_SENDER = false; + public static final String DEFAULT_WINDOW_SORT_ON_SENDER_MODE = "false"; public static class Request { public static final String SQL = "sql"; @@ -1181,7 +1182,7 @@ public static class QueryOptionKey { /// whether the leaf sorts its rows in memory or streams physically sorted segments. No-op for the V2 planner. public static final String STREAMING_SORTED_MAILBOX_RECEIVE = "streamingSortedMailboxReceive"; - /// Per-query override of [Broker#CONFIG_OF_WINDOW_SORT_ON_SENDER]. + /// Per-query auto, true, or false override of [Broker#CONFIG_OF_WINDOW_SORT_ON_SENDER]. public static final String WINDOW_SORT_ON_SENDER = "windowSortOnSender"; // Vector search query options From b43d2c8b43422fb9eff7564d5b12bb5f834dca56 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 4 Oct 2026 03:26:45 +0800 Subject: [PATCH 8/9] Unify shared window ordering regression coverage Align the shared full-input window regression name and comments, and the RexBuilder mock setup order, with the corrected ordering dependency. Preserve the identical test bodies and all AUTO coverage. --- .../rel/rules/PinotWindowExchangeNodeInsertRuleTest.java | 2 +- .../query/runtime/operator/WindowAggregateOperatorTest.java | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRuleTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRuleTest.java index 1380bc4e0427..51e15187c58a 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRuleTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRuleTest.java @@ -67,10 +67,10 @@ public class PinotWindowExchangeNodeInsertRuleTest { @BeforeMethod public void setUp() { _mocks = MockitoAnnotations.openMocks(this); + when(_cluster.getRexBuilder()).thenReturn(REX_BUILDER); RelTraitSet traits = RelTraitSet.createEmpty(); Mockito.when(_input.getTraitSet()).thenReturn(traits); Mockito.when(_input.getCluster()).thenReturn(_cluster); - when(_cluster.getRexBuilder()).thenReturn(REX_BUILDER); } @AfterMethod diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java index 685992bf298b..103315475ec7 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java @@ -1188,10 +1188,10 @@ public void testSumWithUnboundedPrecedingLowerAndUnboundedFollowingUpper(WindowN assertTrue(operator.nextBlock().isSuccess(), "Second block is EOS (done processing)"); } - /// The cold AUTO receiver sort establishes window ordering without applying the broker response cap to its input. + /// The global window receiver sort establishes window ordering without applying the broker response cap to its input. /// Following frames and RANGE peers both need rows beyond that cap, even when the sort emits several blocks. @Test(dataProvider = "windowFrameTypes") - public void testColdReceiverSortRetainsCompleteWindowInput(WindowNode.WindowFrameType frameType) { + public void testReceiverSortRetainsCompleteWindowInput(WindowNode.WindowFrameType frameType) { DataSchema inputSchema = new DataSchema(new String[]{"key", "value"}, new ColumnDataType[]{INT, INT}); DataSchema resultSchema = new DataSchema(new String[]{"key", "value", "sum"}, new ColumnDataType[]{INT, INT, DOUBLE}); From f8a15779fa5a52934275307f2ed23f6de6727a02 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 6 Oct 2026 03:19:31 -0700 Subject: [PATCH 9/9] [ci] Enable existing checks for stacked PR #19714 --- .github/workflows/pinot_compatibility_checks.yml | 1 + .github/workflows/pinot_integration_tests.yml | 1 + .github/workflows/pinot_java11_client_compatibility.yml | 1 + .github/workflows/pinot_quickstart_tests.yml | 1 + .github/workflows/pinot_unit_tests.yml | 1 + .github/workflows/pinot_vuln_check.yml | 1 + 6 files changed, 6 insertions(+) diff --git a/.github/workflows/pinot_compatibility_checks.yml b/.github/workflows/pinot_compatibility_checks.yml index 7e5da173d63b..18923aca4d85 100644 --- a/.github/workflows/pinot_compatibility_checks.yml +++ b/.github/workflows/pinot_compatibility_checks.yml @@ -34,6 +34,7 @@ on: pull_request: branches: - master + - 'xiangfu0/codex1/sorted-exchange-ordering' paths-ignore: - "contrib/**" - "docs/**" diff --git a/.github/workflows/pinot_integration_tests.yml b/.github/workflows/pinot_integration_tests.yml index 8d64b2fae233..747d57026b45 100644 --- a/.github/workflows/pinot_integration_tests.yml +++ b/.github/workflows/pinot_integration_tests.yml @@ -34,6 +34,7 @@ on: pull_request: branches: - master + - 'xiangfu0/codex1/sorted-exchange-ordering' paths-ignore: - "contrib/**" - "docs/**" diff --git a/.github/workflows/pinot_java11_client_compatibility.yml b/.github/workflows/pinot_java11_client_compatibility.yml index 0f7429a18bd7..6c9295d62679 100644 --- a/.github/workflows/pinot_java11_client_compatibility.yml +++ b/.github/workflows/pinot_java11_client_compatibility.yml @@ -34,6 +34,7 @@ on: pull_request: branches: - master + - 'xiangfu0/codex1/sorted-exchange-ordering' paths-ignore: - "contrib/**" - "docs/**" diff --git a/.github/workflows/pinot_quickstart_tests.yml b/.github/workflows/pinot_quickstart_tests.yml index 910f0ca2b6a9..56d2722ba3fb 100644 --- a/.github/workflows/pinot_quickstart_tests.yml +++ b/.github/workflows/pinot_quickstart_tests.yml @@ -34,6 +34,7 @@ on: pull_request: branches: - master + - 'xiangfu0/codex1/sorted-exchange-ordering' paths-ignore: - "contrib/**" - "docs/**" diff --git a/.github/workflows/pinot_unit_tests.yml b/.github/workflows/pinot_unit_tests.yml index e6976ee16407..fe9eeee1757a 100644 --- a/.github/workflows/pinot_unit_tests.yml +++ b/.github/workflows/pinot_unit_tests.yml @@ -34,6 +34,7 @@ on: pull_request: branches: - master + - 'xiangfu0/codex1/sorted-exchange-ordering' paths-ignore: - "contrib/**" - "docs/**" diff --git a/.github/workflows/pinot_vuln_check.yml b/.github/workflows/pinot_vuln_check.yml index 0f0f569e6ec4..89e716dfb4e1 100644 --- a/.github/workflows/pinot_vuln_check.yml +++ b/.github/workflows/pinot_vuln_check.yml @@ -25,6 +25,7 @@ on: pull_request: branches: - master + - 'xiangfu0/codex1/sorted-exchange-ordering' paths: - "**/pom.xml" - "**/package.json"