This is an automated email from the ASF dual-hosted git repository.
gortiz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new e83b59f0d8e Release MSE operator buffers on error and cancel, not only
on end of stream (#19462)
e83b59f0d8e is described below
commit e83b59f0d8e8c39be3138c5629c531471b63f73f
Author: Gonzalo Ortiz Jaureguizar <[email protected]>
AuthorDate: Mon Sep 7 12:45:31 2026 +0200
Release MSE operator buffers on error and cancel, not only on end of stream
(#19462)
---
.../query/runtime/operator/AggregateOperator.java | 33 +-
.../query/runtime/operator/AsofJoinOperator.java | 7 +-
.../query/runtime/operator/BaseJoinOperator.java | 6 +-
.../query/runtime/operator/HashJoinOperator.java | 7 +-
.../pinot/query/runtime/operator/LeafOperator.java | 9 +-
.../query/runtime/operator/MultiStageOperator.java | 60 ++
.../runtime/operator/NonEquiJoinOperator.java | 13 +-
.../query/runtime/operator/RepeatOperator.java | 26 +
.../pinot/query/runtime/operator/SortOperator.java | 37 +-
.../operator/SortedMailboxReceiveOperator.java | 27 +-
.../runtime/operator/set/BinarySetOperator.java | 19 +-
.../query/runtime/operator/set/UnionOperator.java | 18 +-
.../plan/pipeline/PipelineBreakerOperator.java | 10 +
.../operator/OperatorBufferReleaseTest.java | 640 +++++++++++++++++++++
.../operator/SortedMailboxReceiveOperatorTest.java | 34 ++
.../operator/set/SetOperatorBufferReleaseTest.java | 228 ++++++++
16 files changed, 1141 insertions(+), 33 deletions(-)
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
index 0e7501c635a..4717d5a2b96 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
@@ -50,6 +50,7 @@ import org.apache.pinot.query.planner.plannode.AggregateNode;
import org.apache.pinot.query.planner.plannode.PlanNode;
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.SortUtils;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.apache.pinot.spi.exception.QueryErrorCode;
@@ -71,6 +72,9 @@ public class AggregateOperator extends MultiStageOperator {
private final MultiStageOperator _input;
private final DataSchema _resultSchema;
private final AggregationFunction<?, ?>[] _aggFunctions;
+ /// Whether this operator groups. Decides which of the two executors below
is in use; kept separate from them so
+ /// that releasing an executor cannot change the mode.
+ private final boolean _isGroupBy;
@Nullable
private MultistageAggregationExecutor _aggregationExecutor;
@Nullable
@@ -136,7 +140,8 @@ public class AggregateOperator extends MultiStageOperator {
AggregateNode.AggType aggType = node.getAggType();
// TODO: Allow leaf return final result for non-group-by queries
boolean leafReturnFinalResult = node.isLeafReturnFinalResult();
- if (groupKeys.isEmpty()) {
+ _isGroupBy = !groupKeys.isEmpty();
+ if (!_isGroupBy) {
_aggregationExecutor =
new MultistageAggregationExecutor(_aggFunctions, filterArgIds,
maxFilterArgId, aggType, _resultSchema);
_groupByExecutor = null;
@@ -201,20 +206,40 @@ public class AggregateOperator extends MultiStageOperator
{
if (_eosBlock != null) {
return _eosBlock;
}
- MseBlock.Eos finalBlock = _aggregationExecutor != null ?
consumeAggregation() : consumeGroupBy();
+ MseBlock.Eos finalBlock = _isGroupBy ? consumeGroupBy() :
consumeAggregation();
_eosBlock = finalBlock;
if (finalBlock.isError()) {
+ // The upstream failed, so no result will ever be produced from what we
accumulated: drop it right away instead
+ // of waiting for close()/cancel().
+ releaseBuffers();
return finalBlock;
}
MseBlock mseBlock = produceAggregatedBlock();
+ releaseBuffers();
+ return mseBlock;
+ }
+
+ /// Drops the executors, and with them the group-by hash maps and the
aggregate result holders they own. Marks the
+ /// operator finished at the same time, so that a later [#getNextBlock()]
returns the cached end of stream instead
+ /// of trying to consume the input again with a dropped executor.
+ @Override
+ protected void releaseBuffers() {
_aggregationExecutor = null;
_groupByExecutor = null;
- return mseBlock;
+ if (_eosBlock == null) {
+ _eosBlock = SuccessMseBlock.INSTANCE;
+ }
+ }
+
+ @Override
+ protected boolean hasBufferedState() {
+ return _aggregationExecutor != null || _groupByExecutor != null;
}
private MseBlock produceAggregatedBlock() {
- if (_aggregationExecutor != null) {
+ if (!_isGroupBy) {
+ assert _aggregationExecutor != null;
return new RowHeapDataBlock(_aggregationExecutor.getResult(),
_resultSchema, _aggFunctions);
} else {
assert _groupByExecutor != null;
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AsofJoinOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AsofJoinOperator.java
index b072f272c1c..f577782607e 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AsofJoinOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AsofJoinOperator.java
@@ -91,10 +91,15 @@ public class AsofJoinOperator extends BaseJoinOperator {
}
@Override
- protected void onEosProduced() {
+ protected void releaseBuffers() {
_rightTable = null; // Release memory in case we keep the operator around
for a while
}
+ @Override
+ protected boolean hasBufferedState() {
+ return _rightTable != null;
+ }
+
@Override
protected List<Object[]> buildJoinedRows(MseBlock.Data leftBlock) {
assert _rightTable != null : "Right table should not be null when building
joined rows";
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java
index cf481837b12..c6fcdbd390e 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java
@@ -191,13 +191,13 @@ public abstract class BaseJoinOperator extends
MultiStageOperator {
LOGGER.trace("Returning {} for join operator", mseBlock);
if (mseBlock.isEos()) {
_eos = (MseBlock.Eos) mseBlock;
- onEosProduced();
+ // The right table is dead weight from here on, whether this is a
successful end of stream or an error block
+ // propagated from the left input. Release it now rather than waiting
for close()/cancel() to do it.
+ releaseBuffers();
}
return mseBlock;
}
- protected abstract void onEosProduced();
-
protected void buildRightTable() {
LOGGER.trace("Building right table for join operator");
long startTime = System.currentTimeMillis();
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java
index 1a6ea42c2f7..9a2dc958a78 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java
@@ -157,12 +157,17 @@ public class HashJoinOperator extends BaseJoinOperator {
}
@Override
- protected void onEosProduced() {
+ protected void releaseBuffers() {
_rightTable = null;
_matchedRightRows = null;
_nullKeyRightRows = null;
}
+ @Override
+ protected boolean hasBufferedState() {
+ return _rightTable != null || _matchedRightRows != null ||
_nullKeyRightRows != null;
+ }
+
@Override
protected List<Object[]> buildJoinedRows(MseBlock.Data leftBlock) {
assert _rightTable != null : "Right table should not be null when building
joined rows";
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java
index d80adc418c8..efa846b409e 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java
@@ -320,14 +320,17 @@ public class LeafOperator extends MultiStageOperator {
cancelSseTasks();
}
+ /// Stops the single-stage tasks feeding this operator and drains the blocks
they have already queued. Routed
+ /// through the release hook rather than through separate
`close()`/`cancel()` overrides so that it runs on every
+ /// termination path by construction, and so a failure here cannot abort the
rest of the teardown.
@Override
- public void cancel(Throwable e) {
+ protected void releaseBuffers() {
cancelSseTasks();
}
@Override
- public void close() {
- cancelSseTasks();
+ protected boolean hasBufferedState() {
+ return !_blockingQueue.isEmpty();
}
@VisibleForTesting
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
index 0bead828f0f..88dafd2e012 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.query.runtime.operator;
+import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.base.Stopwatch;
@@ -238,6 +239,52 @@ public abstract class MultiStageOperator implements
Operator<MseBlock>, AutoClos
/// the stat map the operator keeps.
public abstract StatMap<?> copyStatMaps();
+ /// Drops the per-query row and hash state this operator is holding —
whatever it accumulated while running, plus
+ /// any input it is still pointing at — so that it becomes collectable even
while the operator itself stays
+ /// reachable.
+ ///
+ /// Both [#close()] and [#cancel(Throwable)] call this, so the state is
released on *every* termination path: end
+ /// of stream, error and cancellation alike. Operators that can release
earlier (when they produce their
+ /// end-of-stream block, say) should keep doing that as well — this is the
backstop, not the prompt path.
+ ///
+ /// **Threading.** Like the rest of teardown, this runs on the thread that
executes the op chain (see [OpChain]) —
+ /// either the worker thread itself, via the scheduler's direct-executor
callback, or a thread holding a chain that
+ /// never started. It is therefore safe to touch operator state without
synchronization, and callers must not
+ /// invoke [#close()] or [#cancel(Throwable)] from anywhere else: nulling a
field that a concurrently running
+ /// `getNextBlock()` is dereferencing would be a use-after-free, not a
missed release.
+ ///
+ /// Rules for implementations:
+ ///
+ /// 1. **Be idempotent.** This can run more than once, and it runs after
[#cancel(Throwable)] on the error path.
+ /// 2. **Leave the stats alone.** [#calculateStats()] and [#copyStatMaps()]
are called after termination.
+ /// 3. **Never mutate something whose identity left the operator.** A block
sitting in a local mailbox still
+ /// points at the list it was built from, and emptying that list
silently drops rows rather than failing.
+ /// Drop or replace the reference; do not `clear()` in place. The same
goes for state an external consumer
+ /// reads after termination — see
[org.apache.pinot.query.runtime.plan.pipeline.PipelineBreakerOperator], whose
+ /// buffer is its output and which therefore releases nothing.
+ /// 4. **Prefer replacing to clearing.** `clear()` drops the elements but
keeps the backing table at whatever
+ /// capacity it grew to — `HashMap`, `ObjectOpenHashSet`, `ArrayList`
and `PriorityQueue` all behave this way,
+ /// which for a large buffer leaves tens of megabytes of empty slots
reachable. Assign a fresh, empty instance
+ /// (or `null`) instead.
+ /// 5. **Do not let a released field double as control flow.** After
release the operator is done, so a field that
+ /// also serves as a mode discriminator or a "have I read the input yet"
marker must not be the one being
+ /// dropped. Keep the discriminator in a separate final field.
+ ///
+ /// An operator that overrides [#close()] or [#cancel(Throwable)] must chain
to `super`, or its state is never
+ /// released — that is not a compile error, so it is on the implementer.
+ protected void releaseBuffers() {
+ }
+
+ /// Whether this operator is currently holding any of the state that
[#releaseBuffers()] drops.
+ ///
+ /// The two are a pair: an operator that overrides one must override the
other, and
+ /// `releaseBuffers(); assert !hasBufferedState();` must hold. It exists so
the release invariant can be asserted
+ /// without reflecting into private fields; nothing in production reads it.
+ @VisibleForTesting
+ protected boolean hasBufferedState() {
+ return false;
+ }
+
// TODO: Ideally close() call should finish within request deadline.
// TODO: Consider passing deadline as part of the API.
@Override
@@ -250,6 +297,7 @@ public abstract class MultiStageOperator implements
Operator<MseBlock>, AutoClos
// Continue processing because even one operator failed to be close,
we should still close the rest.
}
}
+ releaseBuffersSafely();
}
public void cancel(Throwable e) {
@@ -261,6 +309,18 @@ public abstract class MultiStageOperator implements
Operator<MseBlock>, AutoClos
// Continue processing because even one operator failed to be
cancelled, we should still cancel the rest.
}
}
+ releaseBuffersSafely();
+ }
+
+ private void releaseBuffersSafely() {
+ try {
+ releaseBuffers();
+ } catch (Throwable t) {
+ // Releasing buffers is best-effort cleanup; never let it break the rest
of the teardown. Throwable rather than
+ // Exception so that an AssertionError from an implementation (tests run
with -ea) cannot abort a parent's
+ // close loop and leave its siblings unclosed.
+ logger().error("Failed to release the buffers of operator: {}", this, t);
+ }
}
@Override
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/NonEquiJoinOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/NonEquiJoinOperator.java
index bb1840dfa97..315c32d2272 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/NonEquiJoinOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/NonEquiJoinOperator.java
@@ -37,7 +37,7 @@ public class NonEquiJoinOperator extends BaseJoinOperator {
private static final String BUILD_JOINED_ROWS_SCOPE =
"NonEquiJoinOperator#buildJoinedRows";
private static final String BUILD_NON_MATCH_RIGHT_ROWS_SCOPE =
"NonEquiJoinOperator#buildNonMatchRightRows";
- private final List<Object[]> _rightTable;
+ private List<Object[]> _rightTable;
// Track matched right rows for right join and full join to output
non-matched right rows.
// TODO: Revisit whether we should use IntList or RoaringBitmap for smaller
memory footprint.
@Nullable
@@ -69,11 +69,20 @@ public class NonEquiJoinOperator extends BaseJoinOperator {
}
}
+ /// Replaces `_rightTable` with a fresh, empty list rather than clearing it:
`clear()` would keep the element array
+ /// it grew to. A new `ArrayList` rather than `List.of()` because the field
is read — and, if the operator were ever
+ /// re-entered, written — unconditionally, so it must stay both non-null and
mutable.
@Override
- protected void onEosProduced() {
+ protected void releaseBuffers() {
+ _rightTable = new ArrayList<>();
_matchedRightRows = null;
}
+ @Override
+ protected boolean hasBufferedState() {
+ return !_rightTable.isEmpty() || _matchedRightRows != null;
+ }
+
@Override
protected List<Object[]> buildJoinedRows(MseBlock.Data leftBlock) {
ArrayList<Object[]> rows = new ArrayList<>();
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java
index 5045b3b8d07..53013c0ee7a 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java
@@ -25,6 +25,7 @@ import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
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.plan.OpChainExecutionContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -70,6 +71,10 @@ public class RepeatOperator extends MultiStageOperator {
/// Rows of the input block currently being expanded, or null when the next
input block must be pulled.
@Nullable
private List<Object[]> _currentRows;
+ /// Set once the input is exhausted (or the operator is terminated), and
returned from every later call. Without
+ /// it, a null [#_currentRows] would send us back to a spent — after
[#releaseBuffers()], closed — input.
+ @Nullable
+ private MseBlock.Eos _eosBlock;
/// The grouping set (ordinal) the next getNextBlock() call will expand the
current input block for.
private int _currentSet;
@@ -124,9 +129,13 @@ public class RepeatOperator extends MultiStageOperator {
@Override
protected MseBlock getNextBlock() {
+ if (_eosBlock != null) {
+ return _eosBlock;
+ }
if (_currentRows == null) {
MseBlock block = _input.nextBlock();
if (block.isEos()) {
+ _eosBlock = (MseBlock.Eos) block;
return block;
}
_currentRows = ((MseBlock.Data) block).asRowHeap().getRows();
@@ -166,6 +175,23 @@ public class RepeatOperator extends MultiStageOperator {
return new StatMap<>(_statMap);
}
+ /// Drops the input block being expanded. The rows belong to that block, not
to this operator, and the emitted rows
+ /// are freshly allocated, so dropping the reference is enough — clearing
would corrupt a block another operator may
+ /// still be holding. Marks the operator finished at the same time, because
a null [#_currentRows] otherwise means
+ /// "pull the next input block".
+ @Override
+ protected void releaseBuffers() {
+ _currentRows = null;
+ if (_eosBlock == null) {
+ _eosBlock = SuccessMseBlock.INSTANCE;
+ }
+ }
+
+ @Override
+ protected boolean hasBufferedState() {
+ return _currentRows != null;
+ }
+
public enum StatKey implements StatMap.Key {
EXECUTION_TIME_MS(StatMap.Type.LONG) {
@Override
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java
index 82db0410b46..a06e0a483be 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java
@@ -24,6 +24,7 @@ import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;
+import javax.annotation.Nullable;
import org.apache.calcite.rel.RelFieldCollation;
import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
@@ -46,8 +47,18 @@ public class SortOperator extends MultiStageOperator {
private final DataSchema _dataSchema;
private final int _offset;
private final int _numRowsToKeep;
- private final PriorityQueue<Object[]> _priorityQueue;
- private final ArrayList<Object[]> _rows;
+ /// Whether this operator has to sort, or is a plain limit/offset over an
already-sorted input. Decides which of
+ /// the two buffers below is in use. Kept separate from them so that
releasing a buffer cannot change the mode.
+ private final boolean _requiresSort;
+ /// The rows-to-keep heap, used when [#_requiresSort]. Never handed
downstream — [#produceSortedBlock()] drains it
+ /// into a fresh array — so releasing replaces it outright rather than
clearing it, which would keep the backing
+ /// array alive.
+ @Nullable
+ private PriorityQueue<Object[]> _priorityQueue;
+ /// The buffered rows, used when not [#_requiresSort]. Handed downstream as
a sublist view of itself, so releasing
+ /// must drop the reference rather than empty the list.
+ @Nullable
+ private ArrayList<Object[]> _rows;
private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
private boolean _hasConstructedSortedBlock;
@@ -73,7 +84,8 @@ public class SortOperator extends MultiStageOperator {
// - There is no collation
// - Input is already sorted
List<RelFieldCollation> collations = node.getCollations();
- if (collations.isEmpty() || input instanceof SortedMailboxReceiveOperator)
{
+ _requiresSort = !(collations.isEmpty() || input instanceof
SortedMailboxReceiveOperator);
+ if (!_requiresSort) {
_priorityQueue = null;
_rows = new ArrayList<>(Math.min(defaultHolderCapacity, _numRowsToKeep));
} else {
@@ -109,7 +121,14 @@ public class SortOperator extends MultiStageOperator {
}
@Override
- public void cancel(Throwable e) {
+ protected void releaseBuffers() {
+ _priorityQueue = null;
+ _rows = null;
+ }
+
+ @Override
+ protected boolean hasBufferedState() {
+ return _priorityQueue != null || _rows != null;
}
@Override
@@ -125,7 +144,7 @@ public class SortOperator extends MultiStageOperator {
}
_eosBlock = consumeInputBlocks();
// returning upstream error block if finalBlock contains error.
- _statMap.merge(StatKey.REQUIRE_SORT, _priorityQueue != null);
+ _statMap.merge(StatKey.REQUIRE_SORT, _requiresSort);
if (_eosBlock.isError()) {
return _eosBlock;
}
@@ -139,7 +158,8 @@ public class SortOperator extends MultiStageOperator {
private MseBlock produceSortedBlock() {
_hasConstructedSortedBlock = true;
- if (_priorityQueue == null) {
+ if (!_requiresSort) {
+ assert _rows != null : "Rows should not be null when producing the
sorted block";
if (_rows.size() > _offset) {
List<Object[]> row = _rows.subList(_offset, _rows.size());
return new RowHeapDataBlock(row, _dataSchema);
@@ -147,6 +167,7 @@ public class SortOperator extends MultiStageOperator {
return _eosBlock;
}
} else {
+ assert _priorityQueue != null : "Priority queue should not be null when
producing the sorted block";
int resultSize = _priorityQueue.size() - _offset;
if (resultSize <= 0) {
return _eosBlock;
@@ -164,7 +185,8 @@ public class SortOperator extends MultiStageOperator {
MseBlock block = _input.nextBlock();
while (block.isData()) {
List<Object[]> container = ((MseBlock.Data) block).asRowHeap().getRows();
- if (_priorityQueue == null) {
+ if (!_requiresSort) {
+ assert _rows != null : "Rows should not be null when consuming input
blocks";
// TODO: when push-down properly, we shouldn't get more than
_numRowsToKeep
int numRows = _rows.size();
if (numRows < _numRowsToKeep) {
@@ -184,6 +206,7 @@ public class SortOperator extends MultiStageOperator {
}
}
} else {
+ assert _priorityQueue != null : "Priority queue should not be null
when consuming input blocks";
for (Object[] row : container) {
SelectionOperatorUtils.addToPriorityQueue(row, _priorityQueue,
_numRowsToKeep);
}
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java
index bede521ad63..42ad6d05bdf 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java
@@ -21,6 +21,7 @@ package org.apache.pinot.query.runtime.operator;
import com.google.common.base.Preconditions;
import java.util.ArrayList;
import java.util.List;
+import javax.annotation.Nullable;
import org.apache.calcite.rel.RelFieldCollation;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.pinot.common.utils.DataSchema;
@@ -46,7 +47,11 @@ public class SortedMailboxReceiveOperator extends
BaseMailboxReceiveOperator {
private final DataSchema _dataSchema;
private final List<RelFieldCollation> _collations;
- private final List<Object[]> _rows = new ArrayList<>();
+ /// The rows collected from the mailboxes. Sorted in place and handed
downstream as-is, so [#releaseBuffers()]
+ /// drops the reference instead of clearing it: a local mailbox may still be
holding the block that wraps this very
+ /// list, and emptying it would silently drop those rows.
+ @Nullable
+ private List<Object[]> _rows = new ArrayList<>();
private MseBlock _eosBlock;
@@ -77,20 +82,24 @@ public class SortedMailboxReceiveOperator extends
BaseMailboxReceiveOperator {
while (true) {
MseBlock block = _multiConsumer.readMseBlockBlocking();
if (block.isData()) {
+ assert _rows != null : "Rows should not be null while collecting
mailbox blocks";
_rows.addAll(((MseBlock.Data) block).asRowHeap().getRows());
continue;
}
MseBlock.Eos eosBlock = (MseBlock.Eos) block;
onEos();
_eosBlock = eosBlock;
+ List<Object[]> rows = _rows;
+ assert rows != null : "Rows should not be null when the end of stream is
reached";
+ releaseBuffers();
if (eosBlock.isError()) {
return eosBlock;
} else {
- if (!_rows.isEmpty()) {
+ if (!rows.isEmpty()) {
// TODO: This might not be efficient because we are sorting all the
received rows. We should use a k-way merge
// when sender side is sorted.
- _rows.sort(new SortUtils.SortComparator(_collations, false));
- return new RowHeapDataBlock(_rows, _dataSchema);
+ rows.sort(new SortUtils.SortComparator(_collations, false));
+ return new RowHeapDataBlock(rows, _dataSchema);
} else {
return block;
}
@@ -99,14 +108,12 @@ public class SortedMailboxReceiveOperator extends
BaseMailboxReceiveOperator {
}
@Override
- public void close() {
- super.close();
- _rows.clear();
+ protected void releaseBuffers() {
+ _rows = null;
}
@Override
- public void cancel(Throwable t) {
- super.cancel(t);
- _rows.clear();
+ protected boolean hasBufferedState() {
+ return _rows != null;
}
}
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java
index 76f94b12688..3ef64f8e339 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java
@@ -36,7 +36,7 @@ public abstract class BinarySetOperator extends SetOperator {
protected final MultiStageOperator _leftChildOperator;
protected final MultiStageOperator _rightChildOperator;
- protected final Multiset<Record> _rightRowSet;
+ protected Multiset<Record> _rightRowSet;
private MseBlock.Eos _eos;
private boolean _isRightChildOperatorProcessed;
@@ -105,6 +105,7 @@ public abstract class BinarySetOperator extends SetOperator
{
return mseBlock;
} else if (mseBlock.isError()) {
_eos = (MseBlock.Eos) mseBlock;
+ releaseBuffers();
return _eos;
} else if (mseBlock.isSuccess()) {
// If it's a regular EOS block, we continue to process the left child
operator.
@@ -115,12 +116,28 @@ public abstract class BinarySetOperator extends
SetOperator {
MseBlock mseBlock = processLeftOperator();
if (mseBlock.isEos()) {
_eos = (MseBlock.Eos) mseBlock;
+ releaseBuffers();
return _eos;
} else {
return mseBlock;
}
}
+ /// Replaces `_rightRowSet` with a fresh, empty multiset rather than
clearing it: `HashMultiset.clear()` walks every
+ /// entry and still leaves the backing table at the capacity it grew to,
which for a wide INTERSECT / EXCEPT is the
+ /// bulk of what we are trying to release. Safe to swap because
[#handleRowMatched(Object[])] is only reached from
+ /// [#processLeftOperator()], which [#getNextBlock()] stops calling once
`_eos` is set — and `_eos` is set before
+ /// every release.
+ @Override
+ protected void releaseBuffers() {
+ _rightRowSet = HashMultiset.create();
+ }
+
+ @Override
+ protected boolean hasBufferedState() {
+ return !_rightRowSet.isEmpty();
+ }
+
/// Returns true if the row matches the criteria defined by the set
operation.
///
/// Also updates the right row set based on the operator.
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/UnionOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/UnionOperator.java
index 3d258db671a..515f754eb48 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/UnionOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/UnionOperator.java
@@ -41,7 +41,7 @@ public class UnionOperator extends SetOperator {
private MseBlock _eosBlock = null;
private int _currentOperatorIndex = 0;
- private final Set<Record> _seenRecords = new ObjectOpenHashSet<>();
+ private Set<Record> _seenRecords = new ObjectOpenHashSet<>();
public UnionOperator(OpChainExecutionContext opChainExecutionContext,
List<MultiStageOperator> inputOperators, DataSchema dataSchema) {
@@ -60,11 +60,13 @@ public class UnionOperator extends SetOperator {
MseBlock block = currentOperator.nextBlock();
if (block.isError()) {
_eosBlock = block;
+ releaseBuffers();
return block;
} else if (block.isSuccess()) {
_currentOperatorIndex++;
if (_currentOperatorIndex == _inputOperators.size()) {
_eosBlock = block;
+ releaseBuffers();
return block;
}
} else if (block.isData()) {
@@ -103,4 +105,18 @@ public class UnionOperator extends SetOperator {
public String toExplainString() {
return EXPLAIN_NAME;
}
+
+ /// Replaces `_seenRecords` with a fresh, empty set rather than clearing it:
`ObjectOpenHashSet.clear()` nulls the
+ /// entries but deliberately keeps the key table at the capacity it grew to,
and for a de-duplicating UNION that
+ /// table is sized by the full distinct cardinality of every input. Safe to
swap because [#getNextBlock()] returns
+ /// the cached `_eosBlock` without touching the set once it is set, and it
is set before every release.
+ @Override
+ protected void releaseBuffers() {
+ _seenRecords = new ObjectOpenHashSet<>();
+ }
+
+ @Override
+ protected boolean hasBufferedState() {
+ return !_seenRecords.isEmpty();
+ }
}
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
index c2a0aab6056..a9f10d7a0d9 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
@@ -71,6 +71,16 @@ public class PipelineBreakerOperator extends
MultiStageOperator {
_statMap.merge(StatKey.EMITTED_ROWS, numRows);
}
+ /// Deliberately releases nothing, unlike every other operator.
+ ///
+ /// `_resultMap` is this operator's *output*, not scratch space:
[PipelineBreakerExecutor] reads it through
+ /// [#getResultMap()] after the op chain has finished, and `OpChain#close()`
fires the callback that unblocks that
+ /// read only after `close()` has already run. Dropping the map here would
hand the main op chain empty
+ /// pipeline-breaker results, which `nextBlock()` would then surface as an
unexplained error block.
+ @Override
+ protected void releaseBuffers() {
+ }
+
@Override
public List<MultiStageOperator> getChildOperators() {
return _childOperators;
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorBufferReleaseTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorBufferReleaseTest.java
new file mode 100644
index 00000000000..88851a436c3
--- /dev/null
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorBufferReleaseTest.java
@@ -0,0 +1,640 @@
+/**
+ * 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 java.util.List;
+import java.util.function.Supplier;
+import org.apache.calcite.rel.RelFieldCollation;
+import org.apache.calcite.rel.RelFieldCollation.Direction;
+import org.apache.calcite.rel.RelFieldCollation.NullDirection;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.pinot.common.datatable.StatMap;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.query.planner.logical.RexExpression;
+import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.AggregateNode.AggType;
+import org.apache.pinot.query.planner.plannode.JoinNode;
+import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.query.planner.plannode.SortNode;
+import org.apache.pinot.query.runtime.blocks.ErrorMseBlock;
+import org.apache.pinot.query.runtime.blocks.MseBlock;
+import org.apache.pinot.query.runtime.blocks.SuccessMseBlock;
+import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.mockito.Mock;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.BOOLEAN;
+import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.INT;
+import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.LONG;
+import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.STRING;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+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.assertTrue;
+
+
+/// Asserts the invariant that a [MultiStageOperator] releases the row and
hash state it is holding on *every*
+/// termination path — not only when it produces its end-of-stream block, but
also when the query errors out and when
+/// it is cancelled.
+///
+/// While an operator stays reachable its un-released buffers are live
references, so no GC can reclaim them. That is
+/// why the release has to happen in the operator itself rather than depend on
nothing else holding the operator tree.
+///
+/// This is deliberately one cross-operator suite rather than a case bolted
onto each operator's own test class: the
+/// invariant is a property of [MultiStageOperator], the interesting failure
is one operator quietly not having it,
+/// and [#shouldReleaseByEitherPathForEveryOperator()] can only be written in
one place. Each case drives an operator
+/// to the point where it holds state, asserts
[MultiStageOperator#hasBufferedState] so the scenario cannot go
+/// vacuous, terminates it, and asserts the state is gone.
+///
+/// Set operators live in their own package and are covered by
+/// [org.apache.pinot.query.runtime.operator.set.SetOperatorBufferReleaseTest].
+public class OperatorBufferReleaseTest {
+ private static final DataSchema SCHEMA =
+ new DataSchema(new String[]{"int_col", "string_col"}, new
ColumnDataType[]{INT, STRING});
+ private static final DataSchema JOIN_RESULT_SCHEMA =
+ new DataSchema(new String[]{"int_col1", "string_col1", "int_col2",
"string_col2"},
+ new ColumnDataType[]{INT, STRING, INT, STRING});
+ private static final RuntimeException ERROR = new RuntimeException("boom");
+
+ private AutoCloseable _mocks;
+ @Mock
+ private MultiStageOperator _input;
+ @Mock
+ private MultiStageOperator _leftInput;
+ @Mock
+ private MultiStageOperator _rightInput;
+
+ @BeforeMethod
+ public void setUp() {
+ _mocks = openMocks(this);
+ }
+
+ @AfterMethod
+ public void tearDown()
+ throws Exception {
+ _mocks.close();
+ }
+
+ //
---------------------------------------------------------------------------------------------------------------
+ // The base-class contract
+ //
---------------------------------------------------------------------------------------------------------------
+
+ @Test
+ public void shouldReleaseFromBothCloseAndCancel() {
+ RecordingOperator closed = new RecordingOperator();
+ closed.close();
+ assertEquals(closed._releaseCount, 1, "close() must release");
+
+ RecordingOperator cancelled = new RecordingOperator();
+ cancelled.cancel(ERROR);
+ assertEquals(cancelled._releaseCount, 1, "cancel() must release");
+ }
+
+ /// close() may run more than once, and it runs after cancel() on the error
path, so releasing has to be idempotent.
+ @Test
+ public void shouldTolerateRepeatedTermination() {
+ RecordingOperator operator = new RecordingOperator();
+ operator.cancel(ERROR);
+ operator.close();
+ operator.close();
+ assertEquals(operator._releaseCount, 3);
+ }
+
+ /// Children are torn down first, so an operator may still reach into its
inputs while releasing.
+ @Test
+ public void shouldReleaseAfterClosingChildren() {
+ RecordingOperator operator = new RecordingOperator(_input);
+ operator.close();
+ verify(_input, times(1)).close();
+ assertTrue(operator._childrenWereClosedFirst);
+ }
+
+ /// Releasing is best-effort cleanup: a throwing implementation must not
abort the rest of the teardown. An
+ /// `AssertionError` is the interesting case, since tests run with
assertions enabled.
+ @Test
+ public void shouldSurviveAThrowingRelease() {
+ MultiStageOperator throwsException = new RecordingOperator(_input) {
+ @Override
+ protected void releaseBuffers() {
+ throw new RuntimeException("release failed");
+ }
+ };
+ throwsException.close();
+ throwsException.cancel(ERROR);
+
+ MultiStageOperator throwsAssertionError = new
RecordingOperator(_leftInput) {
+ @Override
+ protected void releaseBuffers() {
+ throw new AssertionError("release failed");
+ }
+ };
+ throwsAssertionError.close();
+ throwsAssertionError.cancel(ERROR);
+
+ verify(_input, times(1)).close();
+ verify(_input, times(1)).cancel(any());
+ verify(_leftInput, times(1)).close();
+ verify(_leftInput, times(1)).cancel(any());
+ }
+
+ //
---------------------------------------------------------------------------------------------------------------
+ // SortOperator
+ //
---------------------------------------------------------------------------------------------------------------
+
+ @Test
+ public void shouldReleaseSortPriorityQueueOnError() {
+ givenInputProducesThenFails();
+ SortOperator operator = sortOperator(_input);
+
+ assertTrue(operator.nextBlock().isError());
+ assertTrue(operator.hasBufferedState(), "the sort is still holding its
heap when the error propagates");
+
+ operator.cancel(ERROR);
+
+ assertFalse(operator.hasBufferedState());
+ }
+
+ @Test
+ public void shouldReleaseSortRowsOnError() {
+ // No collations => the operator buffers into _rows instead of the
priority queue.
+ givenInputProducesThenFails();
+ SortOperator operator = sortOperator(_input, List.of());
+
+ assertTrue(operator.nextBlock().isError());
+ assertTrue(operator.hasBufferedState(), "the sort is still holding its
rows when the error propagates");
+
+ operator.close();
+
+ assertFalse(operator.hasBufferedState());
+ }
+
+ /// [SortOperator] used to override `cancel()` with an empty body, so a
cancel never reached the operators feeding
+ /// it. Releasing state must not come at the cost of that recursion.
+ @Test
+ public void shouldPropagateSortCancelToChildren() {
+ when(_input.nextBlock()).thenReturn(SuccessMseBlock.INSTANCE);
+ SortOperator operator = sortOperator(_input);
+
+ operator.cancel(ERROR);
+
+ verify(_input, times(1)).cancel(ERROR);
+ }
+
+ /// The block [SortOperator] emits is a sublist view of `_rows`, and a block
handed to a local mailbox can still be
+ /// read after this op chain has been closed. Releasing must therefore drop
the reference, never empty the list.
+ @Test
+ public void shouldNotEmptyTheBlockEmittedBySort() {
+ when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{2, "b"}, new Object[]{1, "a"}))
+ .thenReturn(SuccessMseBlock.INSTANCE);
+ SortOperator operator = sortOperator(_input, List.of());
+ List<Object[]> rows = ((MseBlock.Data)
operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 2);
+
+ operator.close();
+
+ assertEquals(rows.size(), 2, "the emitted block must survive the operator
being closed");
+ assertEquals(rows.get(0), new Object[]{2, "b"});
+ }
+
+ //
---------------------------------------------------------------------------------------------------------------
+ // Joins
+ //
---------------------------------------------------------------------------------------------------------------
+
+ /// The right side is materialized, the join starts emitting, and then the
op chain is cancelled mid-flight — the
+ /// case where nothing on the success path ever gets to run.
+ @Test
+ public void shouldReleaseHashJoinRightTableOnCancel() {
+ givenRightBuiltThenLeftKeepsFlowing();
+ HashJoinOperator operator = hashJoinOperator();
+
+ assertTrue(operator.nextBlock().isData());
+ assertTrue(operator.hasBufferedState(), "the right table is still held
mid-join");
+
+ operator.cancel(ERROR);
+
+ assertFalse(operator.hasBufferedState());
+ }
+
+ /// When the right side itself fails, the partially built right table never
reaches the success path.
+ @Test
+ public void shouldReleaseHashJoinPartialRightTableOnError() {
+ givenRightSideFails();
+ HashJoinOperator operator = hashJoinOperator();
+
+ assertTrue(operator.nextBlock().isError());
+ assertTrue(operator.hasBufferedState(), "the partial right table is still
held when the error propagates");
+
+ operator.close();
+
+ assertFalse(operator.hasBufferedState());
+ }
+
+ @Test
+ public void shouldReleaseNonEquiJoinRightTableOnCancel() {
+ givenRightBuiltThenLeftKeepsFlowing();
+ NonEquiJoinOperator operator = nonEquiJoinOperator();
+
+ assertTrue(operator.nextBlock().isData());
+ assertTrue(operator.hasBufferedState(), "the right table is still held
mid-join");
+
+ operator.cancel(ERROR);
+
+ assertFalse(operator.hasBufferedState());
+ }
+
+ @Test
+ public void shouldReleaseNonEquiJoinPartialRightTableOnError() {
+ givenRightSideFails();
+ NonEquiJoinOperator operator = nonEquiJoinOperator();
+
+ assertTrue(operator.nextBlock().isError());
+ assertTrue(operator.hasBufferedState(), "the partial right table is still
held when the error propagates");
+
+ operator.close();
+
+ assertFalse(operator.hasBufferedState());
+ }
+
+ /// The emitted rows are copies, so releasing the right table must not
disturb a block already sent downstream.
+ @Test
+ public void shouldNotEmptyTheBlockEmittedByNonEquiJoin() {
+ when(_rightInput.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA,
new Object[]{5, "r"}))
+ .thenReturn(SuccessMseBlock.INSTANCE);
+ when(_leftInput.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{1, "l"}))
+ .thenReturn(SuccessMseBlock.INSTANCE);
+ NonEquiJoinOperator operator = nonEquiJoinOperator();
+ List<Object[]> rows = ((MseBlock.Data)
operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 1);
+
+ operator.close();
+
+ assertEquals(rows.size(), 1, "the emitted block must survive the operator
being closed");
+ assertEquals(rows.get(0), new Object[]{1, "l", 5, "r"});
+ }
+
+ @Test
+ public void shouldReleaseAsofJoinRightTableOnCancel() {
+ givenRightBuiltThenLeftKeepsFlowing();
+ AsofJoinOperator operator = asofJoinOperator();
+
+ assertTrue(operator.nextBlock().isData());
+ assertTrue(operator.hasBufferedState(), "the right table is still held
mid-join");
+
+ operator.cancel(ERROR);
+
+ assertFalse(operator.hasBufferedState());
+ }
+
+ @Test
+ public void shouldReleaseAsofJoinPartialRightTableOnError() {
+ givenRightSideFails();
+ AsofJoinOperator operator = asofJoinOperator();
+
+ assertTrue(operator.nextBlock().isError());
+ assertTrue(operator.hasBufferedState(), "the partial right table is still
held when the error propagates");
+
+ operator.close();
+
+ assertFalse(operator.hasBufferedState());
+ }
+
+ //
---------------------------------------------------------------------------------------------------------------
+ // AggregateOperator and RepeatOperator
+ //
---------------------------------------------------------------------------------------------------------------
+
+ @Test
+ public void shouldReleaseGroupByExecutorOnError() {
+ givenInputProducesThenFails();
+ AggregateOperator operator = aggregateOperator(List.of(0));
+ assertTrue(operator.hasBufferedState());
+
+ assertTrue(operator.nextBlock().isError());
+
+ // The upstream failed, so the group-by hash maps are dead weight from the
moment the error block is produced.
+ assertFalse(operator.hasBufferedState(), "the group-by executor should be
released on the error path");
+ operator.close();
+ assertFalse(operator.hasBufferedState());
+ }
+
+ @Test
+ public void shouldReleaseAggregationExecutorOnError() {
+ givenInputProducesThenFails();
+ AggregateOperator operator = aggregateOperator(List.of());
+ assertTrue(operator.hasBufferedState());
+
+ assertTrue(operator.nextBlock().isError());
+
+ assertFalse(operator.hasBufferedState(), "the aggregation executor should
be released on the error path");
+ operator.cancel(ERROR);
+ assertFalse(operator.hasBufferedState());
+ }
+
+ /// Releasing the executors must not change which of the two the operator
thinks it is: a group-by that has been
+ /// closed still has to behave like a group-by, not silently turn into a
plain aggregation.
+ @Test
+ public void shouldKeepAggregateModeAfterRelease() {
+ givenInputProducesThenFails();
+ AggregateOperator groupBy = aggregateOperator(List.of(0));
+ groupBy.close();
+ assertTrue(groupBy.nextBlock().isEos(), "a released group-by must not be
re-dispatched as an aggregation");
+
+ givenInputProducesThenFails();
+ AggregateOperator aggregation = aggregateOperator(List.of());
+ aggregation.close();
+ assertTrue(aggregation.nextBlock().isEos());
+ }
+
+ @Test
+ public void shouldReleaseRepeatCurrentRowsOnCancel() {
+ when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{1, "a"}, new Object[]{2, "b"}));
+ RepeatOperator operator = repeatOperator();
+
+ assertTrue(operator.nextBlock().isData());
+ assertTrue(operator.hasBufferedState(), "the input block is still held
between grouping sets");
+
+ operator.cancel(ERROR);
+
+ assertFalse(operator.hasBufferedState());
+ }
+
+ /// `_currentRows == null` means "pull the next input block", so releasing
it must also mark the operator finished —
+ /// otherwise a released operator would go back to an input that has already
been closed.
+ @Test
+ public void shouldNotPullFromTheInputAfterRepeatIsReleased() {
+ when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{1, "a"}));
+ RepeatOperator operator = repeatOperator();
+ assertTrue(operator.nextBlock().isData());
+ reset(_input);
+
+ operator.close();
+
+ assertTrue(operator.nextBlock().isEos(), "a released operator must report
end of stream");
+ verify(_input, times(0)).nextBlock();
+ assertFalse(operator.hasBufferedState());
+ }
+
+ //
---------------------------------------------------------------------------------------------------------------
+ // Cross-operator sweep
+ //
---------------------------------------------------------------------------------------------------------------
+
+ /// close() and cancel() must be interchangeable and repeatable for every
operator that buffers: whichever one runs,
+ /// and however often, the operator ends up holding nothing. Guards against
an operator that releases on only one
+ /// of the two paths, which is the exact defect this change set fixes.
+ @Test
+ public void shouldReleaseByEitherPathForEveryOperator() {
+ for (boolean cancelFirst : new boolean[]{true, false}) {
+ assertReleasedByEitherPath(cancelFirst, this::sortHoldingRows);
+ assertReleasedByEitherPath(cancelFirst, this::sortHoldingPriorityQueue);
+ assertReleasedByEitherPath(cancelFirst, this::hashJoinHoldingRightTable);
+ assertReleasedByEitherPath(cancelFirst,
this::nonEquiJoinHoldingRightTable);
+ assertReleasedByEitherPath(cancelFirst, this::asofJoinHoldingRightTable);
+ assertReleasedByEitherPath(cancelFirst,
this::aggregateHoldingGroupByExecutor);
+ assertReleasedByEitherPath(cancelFirst, this::repeatHoldingCurrentRows);
+ }
+ }
+
+ private void assertReleasedByEitherPath(boolean cancelFirst,
Supplier<MultiStageOperator> factory) {
+ MultiStageOperator operator = factory.get();
+ String name = operator.getClass().getSimpleName();
+ assertTrue(operator.hasBufferedState(), name + " should be holding state
before termination");
+
+ terminate(operator, cancelFirst);
+ assertFalse(operator.hasBufferedState(),
+ name + " should have released after " + (cancelFirst ? "cancel" :
"close"));
+
+ // Terminating again, by the other path, must be safe and must leave it
released.
+ terminate(operator, !cancelFirst);
+ assertFalse(operator.hasBufferedState(), name + " should stay released
after a second termination");
+ }
+
+ private static void terminate(MultiStageOperator operator, boolean cancel) {
+ if (cancel) {
+ operator.cancel(ERROR);
+ } else {
+ operator.close();
+ }
+ }
+
+ private MultiStageOperator sortHoldingRows() {
+ resetMocks();
+ givenInputProducesThenFails();
+ SortOperator operator = sortOperator(_input, List.of());
+ operator.nextBlock();
+ return operator;
+ }
+
+ private MultiStageOperator sortHoldingPriorityQueue() {
+ resetMocks();
+ givenInputProducesThenFails();
+ SortOperator operator = sortOperator(_input);
+ operator.nextBlock();
+ return operator;
+ }
+
+ private MultiStageOperator hashJoinHoldingRightTable() {
+ resetMocks();
+ givenRightBuiltThenLeftKeepsFlowing();
+ HashJoinOperator operator = hashJoinOperator();
+ operator.nextBlock();
+ return operator;
+ }
+
+ private MultiStageOperator nonEquiJoinHoldingRightTable() {
+ resetMocks();
+ givenRightBuiltThenLeftKeepsFlowing();
+ NonEquiJoinOperator operator = nonEquiJoinOperator();
+ operator.nextBlock();
+ return operator;
+ }
+
+ private MultiStageOperator asofJoinHoldingRightTable() {
+ resetMocks();
+ givenRightBuiltThenLeftKeepsFlowing();
+ AsofJoinOperator operator = asofJoinOperator();
+ operator.nextBlock();
+ return operator;
+ }
+
+ private MultiStageOperator aggregateHoldingGroupByExecutor() {
+ resetMocks();
+ when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{1, "a"}));
+ return aggregateOperator(List.of(0));
+ }
+
+ private MultiStageOperator repeatHoldingCurrentRows() {
+ resetMocks();
+ when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{1, "a"}));
+ RepeatOperator operator = repeatOperator();
+ operator.nextBlock();
+ return operator;
+ }
+
+ private void resetMocks() {
+ reset(_input, _leftInput, _rightInput);
+ }
+
+ //
---------------------------------------------------------------------------------------------------------------
+ // Fixtures
+ //
---------------------------------------------------------------------------------------------------------------
+
+ private void givenInputProducesThenFails() {
+ when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{2, "b"}, new Object[]{1, "a"}))
+ .thenReturn(ErrorMseBlock.fromException(ERROR));
+ }
+
+ /// Right side completes, left side keeps producing, so the join is
mid-flight when we terminate it.
+ private void givenRightBuiltThenLeftKeepsFlowing() {
+ when(_rightInput.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA,
new Object[]{1, "a"}))
+ .thenReturn(SuccessMseBlock.INSTANCE);
+ when(_leftInput.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{0, "z"}));
+ }
+
+ /// Right side produces one block and then fails, so the right table is only
partially built.
+ private void givenRightSideFails() {
+ when(_rightInput.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA,
new Object[]{1, "a"}))
+ .thenReturn(ErrorMseBlock.fromException(ERROR));
+ }
+
+ private HashJoinOperator hashJoinOperator() {
+ return new HashJoinOperator(OperatorTestUtil.getTracingContext(),
_leftInput, SCHEMA, _rightInput,
+ new JoinNode(-1, JOIN_RESULT_SCHEMA, PlanNode.NodeHint.EMPTY,
List.of(), JoinRelType.FULL, List.of(0),
+ List.of(0), List.of(), JoinNode.JoinStrategy.HASH));
+ }
+
+ private NonEquiJoinOperator nonEquiJoinOperator() {
+ // Condition: left.int_col < right.int_col
+ List<RexExpression> nonEquiConditions = List.of(
+ new RexExpression.FunctionCall(BOOLEAN, SqlKind.LESS_THAN.name(),
+ List.of(new RexExpression.InputRef(0), new
RexExpression.InputRef(2))));
+ return new NonEquiJoinOperator(OperatorTestUtil.getTracingContext(),
_leftInput, SCHEMA, _rightInput,
+ new JoinNode(-1, JOIN_RESULT_SCHEMA, PlanNode.NodeHint.EMPTY,
List.of(), JoinRelType.FULL, List.of(),
+ List.of(), nonEquiConditions, JoinNode.JoinStrategy.HASH));
+ }
+
+ private AsofJoinOperator asofJoinOperator() {
+ // Joined on int_col, MATCH_CONDITION on string_col.
+ RexExpression matchCondition = new RexExpression.FunctionCall(BOOLEAN,
"LESS_THAN_OR_EQUAL",
+ List.of(new RexExpression.InputRef(1), new RexExpression.InputRef(3)));
+ return new AsofJoinOperator(OperatorTestUtil.getTracingContext(),
_leftInput, SCHEMA, _rightInput,
+ new JoinNode(-1, JOIN_RESULT_SCHEMA, PlanNode.NodeHint.EMPTY,
List.of(), JoinRelType.LEFT_ASOF, List.of(0),
+ List.of(0), List.of(), JoinNode.JoinStrategy.ASOF,
matchCondition));
+ }
+
+ private SortOperator sortOperator(MultiStageOperator input) {
+ return sortOperator(input, List.of(new RelFieldCollation(0,
Direction.ASCENDING, NullDirection.LAST)));
+ }
+
+ private SortOperator sortOperator(MultiStageOperator input,
List<RelFieldCollation> collations) {
+ return new SortOperator(OperatorTestUtil.getTracingContext(), input,
+ new SortNode(-1, SCHEMA, PlanNode.NodeHint.EMPTY, List.of(),
collations, 10, 0));
+ }
+
+ private AggregateOperator aggregateOperator(List<Integer> groupKeys) {
+ RexExpression.FunctionCall countStar =
+ new RexExpression.FunctionCall(LONG, SqlKind.COUNT.name(), List.of());
+ DataSchema resultSchema = groupKeys.isEmpty()
+ ? new DataSchema(new String[]{"count"}, new ColumnDataType[]{LONG})
+ : new DataSchema(new String[]{"int_col", "count"}, new
ColumnDataType[]{INT, LONG});
+ return new AggregateOperator(OperatorTestUtil.getTracingContext(), _input,
+ new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY,
List.of(), List.of(countStar), List.of(-1),
+ groupKeys, AggType.DIRECT, false, null, 0));
+ }
+
+ /// Two grouping sets, so the operator still holds the input block after
emitting the first expansion.
+ private RepeatOperator repeatOperator() {
+ DataSchema resultSchema = new DataSchema(new String[]{"int_col",
"string_col", "int_col_key", "$groupingId"},
+ new ColumnDataType[]{INT, STRING, INT, INT});
+ return new RepeatOperator(OperatorTestUtil.getTracingContext(), _input,
new int[]{0},
+ List.of(List.of(0), List.of()), resultSchema);
+ }
+
+ /// Minimal operator that records how the base class drove it.
+ private static class RecordingOperator extends MultiStageOperator {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(RecordingOperator.class);
+
+ private final List<MultiStageOperator> _children;
+ private final StatMap<SortOperator.StatKey> _statMap = new
StatMap<>(SortOperator.StatKey.class);
+ private int _releaseCount;
+ private boolean _childrenWereClosedFirst;
+
+ RecordingOperator(MultiStageOperator... children) {
+ this(OperatorTestUtil.getTracingContext(), children);
+ }
+
+ private RecordingOperator(OpChainExecutionContext context,
MultiStageOperator... children) {
+ super(context);
+ _children = List.of(children);
+ }
+
+ @Override
+ protected void releaseBuffers() {
+ _releaseCount++;
+ for (MultiStageOperator child : _children) {
+ verify(child, atLeastOnce()).close();
+ }
+ _childrenWereClosedFirst = true;
+ }
+
+ @Override
+ public List<MultiStageOperator> getChildOperators() {
+ return _children;
+ }
+
+ @Override
+ protected MseBlock getNextBlock() {
+ return SuccessMseBlock.INSTANCE;
+ }
+
+ @Override
+ public void registerExecution(long time, int numRows, long
memoryUsedBytes, long gcTimeMs) {
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.SORT_OR_LIMIT;
+ }
+
+ @Override
+ public StatMap<SortOperator.StatKey> copyStatMaps() {
+ return new StatMap<>(_statMap);
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
+ @Override
+ public String toExplainString() {
+ return "RECORDING";
+ }
+ }
+}
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java
index 59aeccd6d5d..d5ae3a2abc5 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java
@@ -54,6 +54,7 @@ import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertTrue;
@@ -252,6 +253,39 @@ public class SortedMailboxReceiveOperatorTest {
}
}
+ /// The block this operator emits wraps the very list it buffered the rows
into, and a block handed to a local
+ /// mailbox can still be read after this op chain has been closed. Releasing
the buffer must therefore drop the
+ /// reference, never empty the list.
+ @Test
+ public void shouldNotEmptyTheEmittedBlockOnClose() {
+
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1);
+ Object[] row = new Object[]{1, 1};
+
when(_mailbox1.poll()).thenReturn(OperatorTestUtil.blockWithStats(DATA_SCHEMA,
row),
+ OperatorTestUtil.eosWithEmptyStats());
+ SortedMailboxReceiveOperator operator = getOperator(_stageMetadata1,
RelDistribution.Type.SINGLETON);
+ List<Object[]> resultRows = ((MseBlock.Data)
operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(resultRows.size(), 1);
+
+ operator.close();
+
+ assertEquals(resultRows.size(), 1, "the emitted block must survive the
operator being closed");
+ assertEquals(resultRows.get(0), row);
+ }
+
+ /// The buffered rows must be gone on the error path too, not only once the
sorted block has been produced.
+ @Test
+ public void shouldReleaseBufferedRowsOnError() {
+
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1);
+
when(_mailbox1.poll()).thenReturn(OperatorTestUtil.blockWithStats(DATA_SCHEMA,
new Object[]{1, 1}),
+ OperatorTestUtil.errorWithEmptyStats(new RuntimeException("TEST
ERROR")));
+ SortedMailboxReceiveOperator operator = getOperator(_stageMetadata1,
RelDistribution.Type.SINGLETON);
+ assertTrue(operator.nextBlock().isError());
+
+ operator.cancel(new RuntimeException("TEST ERROR"));
+
+ assertFalse(operator.hasBufferedState());
+ }
+
private SortedMailboxReceiveOperator getOperator(StageMetadata
stageMetadata, RelDistribution.Type distributionType,
DataSchema resultSchema, List<RelFieldCollation> collations, long
deadlineMs) {
OpChainExecutionContext context =
OperatorTestUtil.getOpChainContext(_mailboxService, deadlineMs, stageMetadata);
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/SetOperatorBufferReleaseTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/SetOperatorBufferReleaseTest.java
new file mode 100644
index 00000000000..01ece5ae152
--- /dev/null
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/SetOperatorBufferReleaseTest.java
@@ -0,0 +1,228 @@
+/**
+ * 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.set;
+
+import java.util.List;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.query.runtime.blocks.ErrorMseBlock;
+import org.apache.pinot.query.runtime.blocks.MseBlock;
+import org.apache.pinot.query.runtime.operator.BlockListMultiStageOperator;
+import org.apache.pinot.query.runtime.operator.MultiStageOperator;
+import org.apache.pinot.query.runtime.operator.OperatorTestUtil;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+
+/// The set-operator half of the release invariant asserted by
+/// [org.apache.pinot.query.runtime.operator.OperatorBufferReleaseTest]: every
operator must drop the row state it is
+/// holding on end of stream, on error and on cancellation alike.
+///
+/// These live in their own class because `hasBufferedState()` is `protected`
and the overrides that matter here are
+/// declared in this package.
+///
+/// The [BinarySetOperator] subclasses all buffer the whole right child in
`_rightRowSet`, and [UnionOperator] buffers
+/// every distinct row it has seen — neither was released on any path before.
+public class SetOperatorBufferReleaseTest {
+ private static final DataSchema SCHEMA =
+ new DataSchema(new String[]{"int_col", "string_col"}, new
ColumnDataType[]{
+ ColumnDataType.INT, ColumnDataType.STRING
+ });
+ private static final RuntimeException ERROR = new RuntimeException("boom");
+
+ @Test
+ public void shouldReleaseRightRowSetOnCancelForEveryBinarySetOperator() {
+ for (BinarySetOperatorFactory factory : binarySetOperators()) {
+ // The right child completes and the left child produces a block, so the
operator is mid-flight.
+ MultiStageOperator right = twoRowOperator();
+ BinarySetOperator operator = factory.create(mixedLeftOperator(), right);
+ String name = operator.getClass().getSimpleName();
+
+ assertTrue(operator.nextBlock().isData(), name + " should emit a data
block");
+ assertTrue(operator.hasBufferedState(), name + " should still hold the
right row set mid-flight");
+
+ operator.cancel(ERROR);
+
+ assertFalse(operator.hasBufferedState(), name + " should release the
right row set on cancel");
+ }
+ }
+
+ @Test
+ public void shouldReleaseRightRowSetOnCloseForEveryBinarySetOperator() {
+ for (BinarySetOperatorFactory factory : binarySetOperators()) {
+ MultiStageOperator right = twoRowOperator();
+ BinarySetOperator operator = factory.create(mixedLeftOperator(), right);
+ String name = operator.getClass().getSimpleName();
+
+ assertTrue(operator.nextBlock().isData());
+ assertTrue(operator.hasBufferedState());
+
+ operator.close();
+
+ assertFalse(operator.hasBufferedState(), name + " should release the
right row set on close");
+ // Terminating twice, by the other path, must be safe.
+ operator.cancel(ERROR);
+ assertFalse(operator.hasBufferedState(), name + " should stay released
after a second termination");
+ }
+ }
+
+ /// When the left child fails, the right row set is released in band with
the error block rather than waiting for
+ /// teardown.
+ @Test
+ public void shouldReleaseRightRowSetWhenTheLeftChildFails() {
+ for (BinarySetOperatorFactory factory : binarySetOperators()) {
+ MultiStageOperator right = twoRowOperator();
+ MultiStageOperator left = mock(MultiStageOperator.class);
+ when(left.nextBlock()).thenReturn(ErrorMseBlock.fromException(ERROR));
+ BinarySetOperator operator = factory.create(left, right);
+ String name = operator.getClass().getSimpleName();
+
+ assertTrue(operator.nextBlock().isError(), name + " should propagate the
error");
+
+ assertFalse(operator.hasBufferedState(), name + " should release the
right row set on the error path");
+ operator.close();
+ assertFalse(operator.hasBufferedState());
+ }
+ }
+
+ /// The right child failing leaves the right row set partially built and no
result will ever be produced from it.
+ @Test
+ public void shouldReleasePartialRightRowSetWhenTheRightChildFails() {
+ for (BinarySetOperatorFactory factory : binarySetOperators()) {
+ MultiStageOperator right = mock(MultiStageOperator.class);
+ when(right.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{1, "AA"}))
+ .thenReturn(ErrorMseBlock.fromException(ERROR));
+ BinarySetOperator operator = factory.create(mixedLeftOperator(), right);
+ String name = operator.getClass().getSimpleName();
+
+ assertTrue(operator.nextBlock().isError(), name + " should propagate the
error");
+
+ assertFalse(operator.hasBufferedState(), name + " should release the
partial right row set");
+ }
+ }
+
+ /// The emitted rows come from the left child, so releasing the right row
set must not disturb a block already
+ /// handed downstream.
+ @Test
+ public void shouldNotEmptyTheEmittedBlockOnClose() {
+ MultiStageOperator right = twoRowOperator();
+ MultiStageOperator left = new
BlockListMultiStageOperator.Builder(SCHEMA).addRow(1, "AA").buildWithEos();
+ IntersectOperator operator =
+ new IntersectOperator(OperatorTestUtil.getTracingContext(),
List.of(left, right), SCHEMA);
+ List<Object[]> rows = ((MseBlock.Data)
operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 1);
+
+ operator.close();
+
+ assertEquals(rows.size(), 1, "the emitted block must survive the operator
being closed");
+ assertEquals(rows.get(0), new Object[]{1, "AA"});
+ }
+
+ @Test
+ public void shouldReleaseUnionSeenRecordsOnCancel() {
+ MultiStageOperator first = mock(MultiStageOperator.class);
+ when(first.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{1, "AA"}));
+ MultiStageOperator second = mock(MultiStageOperator.class);
+ UnionOperator operator = new
UnionOperator(OperatorTestUtil.getTracingContext(), List.of(first, second),
SCHEMA);
+
+ assertTrue(operator.nextBlock().isData());
+ assertTrue(operator.hasBufferedState(), "the seen records should be
tracked before termination");
+
+ operator.cancel(ERROR);
+
+ assertFalse(operator.hasBufferedState());
+ }
+
+ @Test
+ public void shouldReleaseUnionSeenRecordsOnError() {
+ MultiStageOperator first = mock(MultiStageOperator.class);
+ when(first.nextBlock()).thenReturn(OperatorTestUtil.block(SCHEMA, new
Object[]{1, "AA"}))
+ .thenReturn(ErrorMseBlock.fromException(ERROR));
+ MultiStageOperator second = mock(MultiStageOperator.class);
+ UnionOperator operator = new
UnionOperator(OperatorTestUtil.getTracingContext(), List.of(first, second),
SCHEMA);
+
+ assertTrue(operator.nextBlock().isData());
+ assertTrue(operator.hasBufferedState());
+ assertTrue(operator.nextBlock().isError());
+
+ assertFalse(operator.hasBufferedState(), "the seen records should be
released on the error path");
+ operator.close();
+ assertFalse(operator.hasBufferedState());
+ }
+
+ /// Draining every input successfully must release too — the success path is
the common one.
+ @Test
+ public void shouldReleaseUnionSeenRecordsOnEndOfStream() {
+ MultiStageOperator first = new
BlockListMultiStageOperator.Builder(SCHEMA).addRow(1, "AA").buildWithEos();
+ MultiStageOperator second = new
BlockListMultiStageOperator.Builder(SCHEMA).addRow(2, "BB").buildWithEos();
+ UnionOperator operator = new
UnionOperator(OperatorTestUtil.getTracingContext(), List.of(first, second),
SCHEMA);
+
+ MseBlock block = operator.nextBlock();
+ while (block.isData()) {
+ block = operator.nextBlock();
+ }
+
+ assertTrue(block.isSuccess());
+ assertFalse(operator.hasBufferedState(), "the seen records should be
released once every input is drained");
+ }
+
+ /// The emitted blocks reference the left rows directly, so releasing must
not touch them.
+ @Test
+ public void shouldNotEmptyTheBlockEmittedByUnion() {
+ MultiStageOperator first = new
BlockListMultiStageOperator.Builder(SCHEMA).addRow(1, "AA").buildWithEos();
+ MultiStageOperator second = new
BlockListMultiStageOperator.Builder(SCHEMA).addRow(2, "BB").buildWithEos();
+ UnionOperator operator = new
UnionOperator(OperatorTestUtil.getTracingContext(), List.of(first, second),
SCHEMA);
+ List<Object[]> rows = ((MseBlock.Data)
operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 1);
+
+ operator.close();
+
+ assertEquals(rows.size(), 1, "the emitted block must survive the operator
being closed");
+ assertEquals(rows.get(0), new Object[]{1, "AA"});
+ }
+
+ private static MultiStageOperator twoRowOperator() {
+ // Two rows so that the set is still non-empty after the left child has
matched one of them.
+ return new BlockListMultiStageOperator.Builder(SCHEMA).addRow(1,
"AA").addRow(2, "BB").buildWithEos();
+ }
+
+ /// One row the right child also has and one it does not, so that every
flavour emits a data block: INTERSECT and
+ /// INTERSECT ALL keep the matching row, MINUS and MINUS ALL keep the other
one.
+ private static MultiStageOperator mixedLeftOperator() {
+ return new BlockListMultiStageOperator.Builder(SCHEMA).addRow(1,
"AA").addRow(3, "CC").buildWithEos();
+ }
+
+ private static List<BinarySetOperatorFactory> binarySetOperators() {
+ return List.of(
+ (left, right) -> new
IntersectOperator(OperatorTestUtil.getTracingContext(), List.of(left, right),
SCHEMA),
+ (left, right) -> new
IntersectAllOperator(OperatorTestUtil.getTracingContext(), List.of(left,
right), SCHEMA),
+ (left, right) -> new
MinusOperator(OperatorTestUtil.getTracingContext(), List.of(left, right),
SCHEMA),
+ (left, right) -> new
MinusAllOperator(OperatorTestUtil.getTracingContext(), List.of(left, right),
SCHEMA));
+ }
+
+ @FunctionalInterface
+ private interface BinarySetOperatorFactory {
+ BinarySetOperator create(MultiStageOperator left, MultiStageOperator
right);
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
