gortiz commented on code in PR #19353: URL: https://github.com/apache/pinot/pull/19353#discussion_r3864460423
########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/SpoolBroadcastExchange.java: ########## @@ -0,0 +1,92 @@ +/** + * 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.exchange; + +import java.util.List; +import java.util.function.Function; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.query.mailbox.SendingMailbox; +import org.apache.pinot.query.runtime.blocks.BlockSplitter; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; + + +/// Broadcasts blocks to the per-receiver-stage exchanges of a multi-send (spool) node. +/// +/// Unlike [BroadcastExchange], which routes the very same block instance to every destination, this exchange gives +/// every destination except the first its own [copy][RowHeapDataBlock#copy()] of blocks that carry mutable cells, +/// i.e. aggregation intermediate results in [OBJECT][ColumnDataType#OBJECT] columns of on-heap blocks. Local +/// (same-JVM) mailboxes deliver on-heap blocks by reference, so without the copies multiple receiver stages would +/// observe the same mutable intermediate result objects and corrupt them by mutating them when merging them or +/// extracting final results. +/// +/// Within a single receiver stage the rows of one copy are still shared: the inner per-stage exchange either routes +/// each row to exactly one worker (hash/singleton distribution), or broadcasts rows that downstream operators never +/// mutate. Aggregation intermediate results only travel on the hash/singleton edge between the partial and the final +/// aggregation, never on broadcast edges, so only the fan-out across receiver stages needs the copies. +class SpoolBroadcastExchange extends BroadcastExchange { + + SpoolBroadcastExchange(List<SendingMailbox> sendingMailboxes, BlockSplitter splitter, + Function<List<SendingMailbox>, Integer> statsIndexChooser) { + super(sendingMailboxes, splitter, statsIndexChooser); + } + + @Override + protected void route(List<SendingMailbox> destinations, MseBlock.Data block) { + int numDestinations = destinations.size(); + if (numDestinations == 1 || !mayContainMutableCells(block)) { Review Comment: Spooling isn't really the defect here — in-place mutation of a payload that an exchange handed to more than one consumer is. A plain `BROADCAST` to a stage with two workers on the same server reproduces this with no spool involved. Looking at the other exchanges: `SingletonExchange` asserts a single mailbox, `RandomExchange` picks one destination per block, and `HashExchange` partitions rows disjointly into fresh blocks. `BroadcastExchange.route` is the **only** other `route()` that hands the same `MseBlock.Data` instance to N mailboxes — and `MailboxSendOperator` creates one `SendingMailbox` per receiver *worker*, so any two workers of that stage co-located on one server get in-memory mailboxes sharing the block by reference. Same instance, same two consumers mutating it, same corruption. So the invariant this PR relies on is a two-way one: *no edge that duplicates a block ever carries a type that a downstream operator mutates in place*. Today that holds only because the exchange above a partial aggregate is always hash/singleton (a broadcast there would double-count), which is a planner property now asserted in a runtime class's javadoc. Two ways to make it structural instead: 1. **Move the copy into `BroadcastExchange`** (or into `BlockExchange#sendBlock`, for local mailboxes after the first). `mayContainMutableCells()` is a scan of a cached array (`DataSchema` caches `getStoredColumnDataTypes()`), so this is near-free, it covers the multi-worker broadcast case, and it makes `SpoolBroadcastExchange` unnecessary along with the unenforced invariant documented in `BroadcastExchange`. 2. **Fix the mutation contract instead** — `AggregationFunction#merge` mutating its left argument and `extractFinalResult` draining the accumulator are what make a shared block unsafe. That's the deeper fix, but it's a deliberate perf choice, so (1) is the realistic one. Either way, could you state in the PR description why the multi-worker broadcast case is out of scope, rather than leaving it implied? ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/blocks/RowHeapDataBlock.java: ########## @@ -102,6 +108,40 @@ public RowHeapDataBlock asRowHeap() { return this; } + /// Returns a copy of this block that does not share any mutable cell values with this block. Review Comment: `copy()` promises more than it delivers: it copies `OBJECT` cells and shares every other mutable cell. "a copy of this block that does not share any mutable cell values" isn't quite what the method does. `OBJECT` is one of roughly ten stored types backed by a mutable Java object: `MAP` holds a live `Map`, `BYTES` a `ByteArray` over a `byte[]`, and every `*_ARRAY` type an `int[]`/`long[]`/`String[]`/`Object[]`. `ColumnDataType.UUID`'s own javadoc in `DataSchema` already flags that its placeholder "wraps a mutable 16-byte array". The javadoc's justification — "cells of all other column types are effectively immutable" — is a statement about current operator behaviour, not about the types, and it's exactly the assumption that will rot. The day an operator sorts an array cell in place or merges into a `MAP` cell, this method keeps silently sharing it and the bug comes back in a form nobody will connect to this code. Suggestion: have the method take an `EnumSet<ColumnDataType>` of the types to copy, so the caller — which knows what operators are downstream — makes that decision explicitly, and the assumption becomes an argument someone has to look at rather than a hidden invariant. It also gives the `BroadcastExchange` case above somewhere to express a wider policy if it ever needs one. At minimum, renaming to `copyAggregationIntermediates()` / `copyObjectColumns()` would stop the name overpromising. -- 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]
