gortiz commented on code in PR #19396: URL: https://github.com/apache/pinot/pull/19396#discussion_r4207838063
########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/KWayMergeSupportPredicate.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.broker.requesthandler; + +import java.util.List; +import java.util.function.BooleanSupplier; +import org.apache.helix.HelixAdmin; +import org.apache.helix.HelixManager; +import org.apache.helix.NotificationContext; +import org.apache.helix.api.listeners.BatchMode; +import org.apache.helix.api.listeners.InstanceConfigChangeListener; +import org.apache.helix.api.listeners.PreFetch; +import org.apache.helix.model.InstanceConfig; +import org.apache.pinot.common.version.PinotVersion; +import org.apache.pinot.spi.config.instance.InstanceType; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.InstanceTypeUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Enables new merge plan nodes only after a complete homogeneous broker/server release-version snapshot. +/// Uses the SAFE stats policy's version comparison, but fails closed on initialization and read failures. +/// Snapshot builds are excluded because equal snapshot versions can contain different protocol implementations. +/// Helix drives updates serially; query threads read the volatile result without locking. +@BatchMode(enabled = false) +@PreFetch(enabled = false) +public class KWayMergeSupportPredicate implements InstanceConfigChangeListener, BooleanSupplier { + private static final Logger LOGGER = LoggerFactory.getLogger(KWayMergeSupportPredicate.class); + private final HelixManager _manager; + private final String _version; + private volatile boolean _supported; + + public KWayMergeSupportPredicate(HelixManager manager) { + this(manager, PinotVersion.VERSION); + } + + KWayMergeSupportPredicate(HelixManager manager, String version) { + _manager = manager; + _version = version; + } + + @Override + public boolean getAsBoolean() { + return _supported && _manager.isConnected(); + } + + @Override + public synchronized void onInstanceConfigChange(List<InstanceConfig> configs, NotificationContext context) { + // Clear first: no partial snapshot or failed refresh may leave the new wire node enabled. + _supported = false; Review Comment: Clearing the flag first makes the feature flap off during every rescan, and if the scan throws it stays off until the next instance config event, which may never come in a stable cluster. I'd compute into a local and publish once. Also, unlike `SendStatsPredicate`, nothing logs which instances keep the gate closed, so an operator who enables the option can't tell why it does nothing. ########## pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/MailboxMergeReceiveNode.java: ########## @@ -0,0 +1,87 @@ +/** + * 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.plannode; + +import java.util.List; +import java.util.Objects; +import javax.annotation.Nullable; +import org.apache.calcite.rel.RelDistribution; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType; +import org.apache.pinot.common.utils.DataSchema; + + +/// Receives streams ordered by the logical plan and preserves that order with a k-way merge. +/// This distinct wire node is emitted only when the broker verifies homogeneous cluster versions. +/// Instances belong to a single query plan and are not thread-safe. +public class MailboxMergeReceiveNode extends MailboxReceiveNode { + private final int _fetch; + private final int _offset; + + public MailboxMergeReceiveNode(int stageId, DataSchema schema, int senderStageId, + PinotRelExchangeType exchangeType, RelDistribution.Type distribution, @Nullable List<Integer> keys, + List<RelFieldCollation> collations, @Nullable MailboxSendNode sender) { + this(stageId, schema, senderStageId, exchangeType, distribution, keys, collations, -1, -1, sender); + } + + public MailboxMergeReceiveNode(int stageId, DataSchema schema, int senderStageId, + PinotRelExchangeType exchangeType, RelDistribution.Type distribution, @Nullable List<Integer> keys, + List<RelFieldCollation> collations, int fetch, int offset, @Nullable MailboxSendNode sender) { + super(stageId, schema, senderStageId, exchangeType, distribution, keys, collations, false, false, sender); Review Comment: The subclass passes `sort=false, sortedOnSender=false` to a parent whose flags mean the opposite of what this node does. Any code that reads `isSort()` gets the wrong answer for it, and every visitor needs an `instanceof` (serializer, `PlanNodeToRelConverter`, `PlanNodeMerger`, `EquivalentStagesFinder`, `PlanFragmenter`, `PlanNodeToOpChain`). I'd make it a sibling `PlanNode` with its own `visitMailboxMergeReceive`, and deprecate the old `sort`/`sortedOnSender` flags in this PR so we end up with a single model. ########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/KWayMergeSupportPredicate.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.broker.requesthandler; + +import java.util.List; +import java.util.function.BooleanSupplier; +import org.apache.helix.HelixAdmin; +import org.apache.helix.HelixManager; +import org.apache.helix.NotificationContext; +import org.apache.helix.api.listeners.BatchMode; +import org.apache.helix.api.listeners.InstanceConfigChangeListener; +import org.apache.helix.api.listeners.PreFetch; +import org.apache.helix.model.InstanceConfig; +import org.apache.pinot.common.version.PinotVersion; +import org.apache.pinot.spi.config.instance.InstanceType; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.InstanceTypeUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Enables new merge plan nodes only after a complete homogeneous broker/server release-version snapshot. +/// Uses the SAFE stats policy's version comparison, but fails closed on initialization and read failures. +/// Snapshot builds are excluded because equal snapshot versions can contain different protocol implementations. +/// Helix drives updates serially; query threads read the volatile result without locking. +@BatchMode(enabled = false) +@PreFetch(enabled = false) +public class KWayMergeSupportPredicate implements InstanceConfigChangeListener, BooleanSupplier { + private static final Logger LOGGER = LoggerFactory.getLogger(KWayMergeSupportPredicate.class); + private final HelixManager _manager; + private final String _version; + private volatile boolean _supported; + + public KWayMergeSupportPredicate(HelixManager manager) { + this(manager, PinotVersion.VERSION); + } + + KWayMergeSupportPredicate(HelixManager manager, String version) { + _manager = manager; + _version = version; + } + + @Override + public boolean getAsBoolean() { + return _supported && _manager.isConnected(); Review Comment: The `isConnected()` check doesn't add anything here: a session expiry already triggers FINALIZE and a new INIT. Returning `_supported` alone avoids taking the ZK event lock on the query path. ########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/KWayMergeSupportPredicate.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.broker.requesthandler; + +import java.util.List; +import java.util.function.BooleanSupplier; +import org.apache.helix.HelixAdmin; +import org.apache.helix.HelixManager; +import org.apache.helix.NotificationContext; +import org.apache.helix.api.listeners.BatchMode; +import org.apache.helix.api.listeners.InstanceConfigChangeListener; +import org.apache.helix.api.listeners.PreFetch; +import org.apache.helix.model.InstanceConfig; +import org.apache.pinot.common.version.PinotVersion; +import org.apache.pinot.spi.config.instance.InstanceType; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.InstanceTypeUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Enables new merge plan nodes only after a complete homogeneous broker/server release-version snapshot. +/// Uses the SAFE stats policy's version comparison, but fails closed on initialization and read failures. +/// Snapshot builds are excluded because equal snapshot versions can contain different protocol implementations. +/// Helix drives updates serially; query threads read the volatile result without locking. +@BatchMode(enabled = false) +@PreFetch(enabled = false) +public class KWayMergeSupportPredicate implements InstanceConfigChangeListener, BooleanSupplier { + private static final Logger LOGGER = LoggerFactory.getLogger(KWayMergeSupportPredicate.class); + private final HelixManager _manager; + private final String _version; + private volatile boolean _supported; + + public KWayMergeSupportPredicate(HelixManager manager) { + this(manager, PinotVersion.VERSION); + } + + KWayMergeSupportPredicate(HelixManager manager, String version) { + _manager = manager; + _version = version; + } + + @Override + public boolean getAsBoolean() { + return _supported && _manager.isConnected(); + } + + @Override + public synchronized void onInstanceConfigChange(List<InstanceConfig> configs, NotificationContext context) { + // Clear first: no partial snapshot or failed refresh may leave the new wire node enabled. + _supported = false; + if (context.getType() != NotificationContext.Type.INIT + && context.getType() != NotificationContext.Type.CALLBACK) { + return; + } + if (_version == null || _version.equals(PinotVersion.UNKNOWN) || _version.contains("SNAPSHOT")) { Review Comment: Excluding SNAPSHOT means no CI run, quickstart or locally built cluster can take the merge path, so it will run for the first time in a release cluster. I'd keep a test-only way to force the gate open and add an integration test that compares the results with the receiver sort. ########## pinot-segment-local/pom.xml: ########## @@ -88,6 +88,13 @@ <artifactId>h3</artifactId> </dependency> + <!-- Zstd's API annotations are required when javac resolves annotated return types with warnings enabled. --> Review Comment: This looks unrelated to the PR. Nothing in `pinot-segment-local` uses these annotations. If it's needed for warning-enabled builds, I'd move it to its own PR. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxMergeReceiveOperator.java: ########## @@ -0,0 +1,351 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.operator; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.Deque; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.PriorityQueue; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.commons.collections4.CollectionUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.mailbox.ReceivingMailbox; +import org.apache.pinot.query.planner.plannode.MailboxMergeReceiveNode; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; +import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; +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.apache.pinot.spi.query.QueryThreadContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// K-way merges streams whose ordering is guaranteed by the logical plan. +/// +/// The broker only emits this operator's plan node in a homogeneous cluster. Ordering is established upstream; +/// mailbox senders preserve it without transport metadata or runtime structural checks. +/// +/// Reads whichever mailbox is ready to avoid cross-receiver backpressure cycles. Multi-sender output blocks contain +/// at most 10,000 rows; a single sender preserves its input blocks. Read-ahead retained while another sender is starved +/// can approach the legacy full receiver sort's memory footprint. +/// +/// Driven by a single consumer thread; not thread-safe. +public class SortedMailboxMergeReceiveOperator extends BaseMailboxReceiveOperator { + private static final Logger LOGGER = LoggerFactory.getLogger(SortedMailboxMergeReceiveOperator.class); + + private static final String EXPLAIN_NAME = "SORTED_MAILBOX_MERGE_RECEIVE"; + private static final String MERGE_SCOPE = "SortedMailboxMergeReceiveOperator"; + private final DataSchema _dataSchema; + private final Comparator<Object[]> _comparator; + private final PriorityQueue<SenderCursor> _readyCursors; + private final boolean _singleSortedSender; + private int _rowsToSkip; + private long _rowsToEmit; + /// Senders that have not finished but do not currently have a row ready. Nothing can be emitted while this is + /// non-empty because any one of these senders may hold the next row. + private final Set<SenderCursor> _starvedCursors = Collections.newSetFromMap(new IdentityHashMap<>()); + private final Map<AsyncStream<ReceivingMailbox.MseBlockWithStats>, SenderCursor> _cursorsByStream = + new IdentityHashMap<>(); + + @Nullable + private MseBlock _eosBlock; + + public SortedMailboxMergeReceiveOperator(OpChainExecutionContext context, MailboxMergeReceiveNode node) { + super(context, node); + Preconditions.checkState(!CollectionUtils.isEmpty(node.getCollations()), "Field collations must be set"); + _dataSchema = node.getDataSchema(); + _rowsToSkip = Math.max(node.getOffset(), 0); + _rowsToEmit = node.getFetch() < 0 ? Long.MAX_VALUE : node.getFetch(); + _comparator = new SortUtils.SortComparator(List.copyOf(node.getCollations()), false); + List<AsyncStream<ReceivingMailbox.MseBlockWithStats>> streams = _multiConsumer.getLiveStreamsSnapshot(); + _readyCursors = new PriorityQueue<>(Math.max(streams.size(), 1), + (left, right) -> _comparator.compare(left.peek(), right.peek())); + _singleSortedSender = streams.size() == 1; + if (!_singleSortedSender) { + for (AsyncStream<ReceivingMailbox.MseBlockWithStats> stream : streams) { + SenderCursor cursor = new SenderCursor(); + _cursorsByStream.put(stream, cursor); + _starvedCursors.add(cursor); + } + } + } + + @Override + protected Logger logger() { + return LOGGER; + } + + @Override + public String toExplainString() { + return EXPLAIN_NAME; + } + + @Override + protected MseBlock getNextBlock() { + if (_eosBlock != null) { + return _eosBlock; + } + if (_isEarlyTerminated) { + return readUntilEos(); + } + if (_rowsToEmit == 0) { + earlyTerminate(); + return readUntilEos(); + } + while (true) { + MseBlock block = _singleSortedSender ? readSingleSortedSender() : mergeNextBlock(); + if (block.isEos()) { + return block; + } + if (_rowsToSkip == 0 && _rowsToEmit == Long.MAX_VALUE) { + return block; + } + List<Object[]> rows = ((MseBlock.Data) block).asRowHeap().getRows(); + int from = Math.min(_rowsToSkip, rows.size()); + _rowsToSkip -= from; + int count = (int) Math.min(rows.size() - from, _rowsToEmit); + _rowsToEmit -= count; + if (_rowsToEmit == 0) { + earlyTerminate(); + } + if (count > 0) { + return from == 0 && count == rows.size() ? block + : new RowHeapDataBlock(rows.subList(from, from + count), _dataSchema); + } + } + } + + /// Passes through one sorted sender without copying its rows through the merge heap. + private MseBlock readSingleSortedSender() { + while (true) { + MseBlock block = _multiConsumer.readMseBlockBlocking(); + if (block.isEos()) { + return terminate(block); + } + MseBlock.Data dataBlock = (MseBlock.Data) block; + checkActiveTerminationAndSampleUsage(); + + if (dataBlock.getNumRows() > 0) { + return dataBlock; + } + } + } + + /// Merges the sorted senders, emitting at most [SortOperator#DEFAULT_MAX_ROWS_PER_BLOCK] rows per call. + private MseBlock mergeNextBlock() { + List<Object[]> rows = new ArrayList<>(); + while (rows.size() < SortOperator.DEFAULT_MAX_ROWS_PER_BLOCK) { + if (!_starvedCursors.isEmpty()) { + MseBlock.Eos error; + if (rows.isEmpty()) { + error = readOneBlock(); + } else { + // Rows already removed from the heap are a globally ordered prefix. Consume any immediately available + // cursor progress so blocks can still be coalesced, but return that safe prefix instead of waiting merely + // to fill the output block. + MseBlock block = _multiConsumer.pollMseBlockOrStreamCompletion(); + if (block == null && _multiConsumer.getFinishedStreamsLastRead().isEmpty()) { + break; + } + error = processReadBlock(block); + } + if (error != null) { + return terminate(error); + } + + continue; + } + if (_readyCursors.isEmpty()) { + break; + } + SenderCursor cursor = _readyCursors.remove(); + rows.add(cursor.next()); + if (cursor.hasRow()) { + _readyCursors.add(cursor); + } else if (!cursor._finished) { + _starvedCursors.add(cursor); + } + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(rows.size(), MERGE_SCOPE, + _context.getActiveDeadlineMs()); Review Comment: Receive operators use the passive deadline on purpose (`BaseMailboxReceiveOperator.getDeadlineMs()`), so the upstream operator that is actually working times out first and reports a useful error. Using the active deadline here makes this stage time out first with a generic error. I'd use the inherited `checkTerminationAndSampleUsage*` methods. Same at line 257. ########## pinot-common/src/main/proto/plan.proto: ########## @@ -201,6 +202,18 @@ message MailboxReceiveNode { bool sortedOnSender = 7; } +// Each input stream is ordered by the logical plan. Older servers reject the unknown PlanNode oneof tag. +// Brokers only emit this node after verifying homogeneous supported versions. +message MailboxMergeReceiveNode { + int32 senderStageId = 1; + ExchangeType exchangeType = 2; + DistributionType distributionType = 3; + repeated int32 keys = 4; + repeated Collation collations = 5; + int32 fetch = 6; Review Comment: These fields are never set by the planner, and once released they are permanent. Their semantics also differ from `SortNode`: here `fetch=0` means "emit nothing", while `SortOperator` treats `fetch<=0` as unbounded. Since proto3 encodes an unset int32 as 0, any producer that forgets `setFetch(-1)` silently returns empty windows. I'd drop them for now. ########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java: ########## @@ -623,6 +632,9 @@ private ImmutableQueryEnvironment.Config getQueryEnvConf(HttpHeaders httpHeaders .defaultDisabledPlannerRules(_defaultDisabledPlannerRules) .defaultSortExchangeCopyLimit(sortExchangeCopyThreshold) .defaultSealedInListThreshold(sealedInListThreshold) + .defaultWindowSortOnSender(windowSortOnSender) + .isKWayMergeSupported(!QueryOptionsUtils.isMultiClusterRoutingEnabled(queryOptions, false) + && _kWayMergeSupported.getAsBoolean()) Review Comment: This runs for every MSE query, with or without windows and with the option on or off. `getAsBoolean()` calls `isConnected()`, which goes through `ZkClient.waitForKeeperState` and acquires the ZK event lock (I checked the Helix 2.0.1 bytecode). During a session expiry and reconnect, every MSE query blocks here. Could the predicate return only the volatile flag, and could we evaluate it only when the window option is enabled? ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java: ########## @@ -282,6 +283,12 @@ public Void visitProject(ProjectNode node, ServerPlanRequestContext context) { @Override public Void visitSort(SortNode node, ServerPlanRequestContext context) { if (visit(node.getInputs().get(0), context)) { + if (node.getFetch() == Integer.MAX_VALUE) { Review Comment: `fetch == Integer.MAX_VALUE` becomes an implicit contract between the planner and the server. A user query with `ORDER BY x LIMIT 2147483647` also hits this branch and loses the V1 sort pushdown. The requirement belongs to the exchange (the merge needs a single run per sender), so I'd use an explicit marker for it. ########## pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/KWayMergeExchangeNode.java: ########## @@ -0,0 +1,44 @@ +/** + * 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.plannode; + +import java.util.List; +import javax.annotation.Nullable; +import org.apache.calcite.rel.RelDistribution; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType; +import org.apache.pinot.common.utils.DataSchema; + + +/// Planning-only exchange requiring ordered inputs and producing a k-way merge receive and a plain send. +/// Instances are local to a query's planning thread and are never serialized. +public class KWayMergeExchangeNode extends ExchangeNode { + public KWayMergeExchangeNode(int stageId, DataSchema schema, List<PlanNode> inputs, + RelDistribution.Type distribution, @Nullable List<Integer> keys, boolean prePartitioned, + List<RelFieldCollation> collations, String hashFunction) { + super(stageId, schema, inputs, PinotRelExchangeType.STREAMING, distribution, keys, prePartitioned, collations, Review Comment: Same here: the parent's `sortOnSender`/`sortOnReceiver` are hard-coded to false. Since `ExchangeNode` is planner-internal, I'd either keep a single exchange node, or make the merge exchange a separate type, but not a subtype that contradicts the parent. ########## pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java: ########## @@ -583,6 +584,13 @@ public void start() _accessControlFactory, _queryQuotaManager, _tableCache, _multiStageQueryThrottler, _failureDetector, _threadAccountant, multiClusterRoutingContext, workerManager, multiClusterWorkerManager, _serverRoutingStatsManager); + KWayMergeSupportPredicate mergeSupport = new KWayMergeSupportPredicate(_spectatorHelixManager); + try { + _spectatorHelixManager.addInstanceConfigChangeListener(mergeSupport); Review Comment: This is registered on every broker, even when the broker config is off and no query will ever enable the option. I'd register it only when the feature can be enabled, or make it cheap enough that it doesn't matter. ########## pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java: ########## @@ -197,8 +199,10 @@ public PlanNode visitExchange(ExchangeNode node, Context context) { _mailboxSendToExchangeNodeMap.put(mailboxSendNode, node); // Return the MailboxReceiveNode as the leave node of the current PlanFragment. - MailboxReceiveNode mailboxReceiveNode = - new MailboxReceiveNode(receiverPlanFragmentId, nextPlanFragmentRoot.getDataSchema(), + MailboxReceiveNode mailboxReceiveNode = node instanceof KWayMergeExchangeNode + ? new MailboxMergeReceiveNode(receiverPlanFragmentId, nextPlanFragmentRoot.getDataSchema(), Review Comment: This is the only production constructor call and it always passes fetch/offset -1, so the fetch/offset support in the node, the proto, the operator, the explain and the equivalence check is only reachable from tests. I'd remove it until a rule pushes a limit into the merge receive, or add that rule and its plan tests in this PR. ########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/KWayMergeSupportPredicate.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.broker.requesthandler; + +import java.util.List; +import java.util.function.BooleanSupplier; +import org.apache.helix.HelixAdmin; +import org.apache.helix.HelixManager; +import org.apache.helix.NotificationContext; +import org.apache.helix.api.listeners.BatchMode; +import org.apache.helix.api.listeners.InstanceConfigChangeListener; +import org.apache.helix.api.listeners.PreFetch; +import org.apache.helix.model.InstanceConfig; +import org.apache.pinot.common.version.PinotVersion; +import org.apache.pinot.spi.config.instance.InstanceType; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.InstanceTypeUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Enables new merge plan nodes only after a complete homogeneous broker/server release-version snapshot. +/// Uses the SAFE stats policy's version comparison, but fails closed on initialization and read failures. +/// Snapshot builds are excluded because equal snapshot versions can contain different protocol implementations. +/// Helix drives updates serially; query threads read the volatile result without locking. +@BatchMode(enabled = false) +@PreFetch(enabled = false) +public class KWayMergeSupportPredicate implements InstanceConfigChangeListener, BooleanSupplier { + private static final Logger LOGGER = LoggerFactory.getLogger(KWayMergeSupportPredicate.class); + private final HelixManager _manager; + private final String _version; + private volatile boolean _supported; + + public KWayMergeSupportPredicate(HelixManager manager) { + this(manager, PinotVersion.VERSION); + } + + KWayMergeSupportPredicate(HelixManager manager, String version) { + _manager = manager; + _version = version; + } + + @Override + public boolean getAsBoolean() { + return _supported && _manager.isConnected(); + } + + @Override + public synchronized void onInstanceConfigChange(List<InstanceConfig> configs, NotificationContext context) { + // Clear first: no partial snapshot or failed refresh may leave the new wire node enabled. + _supported = false; + if (context.getType() != NotificationContext.Type.INIT + && context.getType() != NotificationContext.Type.CALLBACK) { + return; + } + if (_version == null || _version.equals(PinotVersion.UNKNOWN) || _version.contains("SNAPSHOT")) { + return; + } + try { + HelixAdmin admin = _manager.getClusterManagmentTool(); + String cluster = _manager.getClusterName(); + List<String> instances = admin.getInstancesInCluster(cluster); Review Comment: This rescans every broker and server on every callback: 1 + N ZK reads on the spectator event thread, for data changes too. Servers write their instance config on start and shutdown, so a rolling restart costs O(N²) reads per broker. `SendStatsPredicate.Safe` only rescans on INIT or child changes and reads the single changed path otherwise. `getInstancesInCluster` also returns non-live instances, so a stale znode from a decommissioned host keeps the feature off. Could we extract the version tracking from `SendStatsPredicate` into something shared instead of duplicating it? ########## pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotWindowExchangeNodeInsertRule.java: ########## @@ -141,27 +158,44 @@ public void onMatch(RelOptRuleCall call) { exchange = PinotLogicalExchange.create(input, RelDistributions.hash(windowGroup.keys.toList()), prePartitioned); } else { // PARTITION BY and ORDER BY on different key(s) - // Add a LogicalSortExchange hashed on the partition by keys and collation based on order by keys. - // The ordering itself is established by the Sort placed over the exchange below, not by the receive - // operator - see the comment at the transformTo call. + // Keep the receiver full-sort path for a partitioned exchange. Sorting before the hash exchange compares rows + // routed to different receivers and blocks streaming; the explicit Sort retained above the exchange establishes + // the required ordering after partitioning. exchange = PinotLogicalSortExchange.create(input, RelDistributions.hash(windowGroup.keys.toList()), windowGroup.orderKeys, false, false, prePartitioned); } } - // WindowAggregateOperator requires its input ordered on the ORDER BY keys and does no ordering of its own, so - // where the exchange carries a collation the ordering has to be established above it. Place an explicit Sort - // rather than asking the receive operator to sort: SortOperator is the operator that knows fetch/offset, and - // SortedMailboxReceiveOperator is deprecated. The Sort carries no fetch, so it keeps every row - the same - // semantics as the unbounded list the receive operator used. + // WindowAggregateOperator requires its input ordered on the ORDER BY keys and does no ordering of its own. A sort + // exchange that sorts on the receiver establishes that contract itself; adding another Sort above it would make + // 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. PinotSortExchangeCopyRule does not either: it declines when there is no fetch. + // exchange. Window sorts must retain every input row, independently of the broker response limit. RelNode windowInput = exchange instanceof PinotLogicalSortExchange - ? LogicalSort.create(exchange, ((PinotLogicalSortExchange) exchange).getCollation(), null, null) : exchange; + && !((PinotLogicalSortExchange) exchange).isSortOnReceiver() Review Comment: Every `PinotLogicalSortExchange.create` in this rule passes `sortOnReceiver=false`, so this condition is always true, and the comment describes a rolling-upgrade case this rule can't produce. I'd remove both. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
