diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java index 4a39ab5d8e70..41c0046fd921 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java @@ -391,6 +391,26 @@ public static Integer getStreamingGroupByFlushThreshold(Map quer return checkedParseIntNonNegative(QueryOptionKey.STREAMING_GROUP_BY_FLUSH_THRESHOLD, value); } + public static boolean isSortedSelectionMergeEnabled(Map queryOptions) { + return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SORTED_SELECTION_MERGE_ENABLED)); + } + + @Nullable + public static Integer getSortedSelectionMergeBlockSize(Map queryOptions) { + String value = queryOptions.get(QueryOptionKey.SORTED_SELECTION_MERGE_BLOCK_SIZE); + return checkedParseIntPositive(QueryOptionKey.SORTED_SELECTION_MERGE_BLOCK_SIZE, value); + } + + public static boolean isStreamingSortedMailboxReceiveEnabled(Map queryOptions) { + return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE)); + } + + @Nullable + public static Integer getStreamingSortedMailboxReceiveBlockSize(Map queryOptions) { + String value = queryOptions.get(QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE); + return checkedParseIntPositive(QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE, value); + } + public static boolean isNullHandlingEnabled(Map queryOptions) { return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.ENABLE_NULL_HANDLING)); } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java index 5ece39b66cb2..c65d251343a5 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java @@ -38,7 +38,8 @@ public class QueryOptionsUtilsTest { private static final List POSITIVE_INT_KEYS = List.of(NUM_REPLICA_GROUPS_TO_QUERY, MAX_EXECUTION_THREADS, NUM_GROUPS_LIMIT, MAX_INITIAL_RESULT_HOLDER_CAPACITY, - MAX_STREAMING_PENDING_BLOCKS, MAX_ROWS_IN_JOIN, MAX_ROWS_IN_WINDOW); + MAX_STREAMING_PENDING_BLOCKS, MAX_ROWS_IN_JOIN, MAX_ROWS_IN_WINDOW, SORTED_SELECTION_MERGE_BLOCK_SIZE, + STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE); private static final List NON_NEGATIVE_INT_KEYS = List.of(MULTI_STAGE_LEAF_LIMIT); private static final List UNBOUNDED_INT_KEYS = List.of(MIN_SEGMENT_GROUP_TRIM_SIZE, MIN_SERVER_GROUP_TRIM_SIZE, MIN_BROKER_GROUP_TRIM_SIZE, @@ -303,6 +304,21 @@ public void testInvertedIndexDistinctCostRatioRejectsNonFiniteValues() { } } + @Test + public void testStreamingSortedMailboxReceiveEnabled() { + // Unset and any non-"true" value are equivalent: the k-way merge stays off. + assertFalse(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled(Map.of())); + assertFalse(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled(new HashMap<>())); + assertTrue(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled( + Map.of(STREAMING_SORTED_MAILBOX_RECEIVE, "true"))); + assertTrue(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled( + Map.of(STREAMING_SORTED_MAILBOX_RECEIVE, "TRUE"))); + assertFalse(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled( + Map.of(STREAMING_SORTED_MAILBOX_RECEIVE, "false"))); + assertFalse(QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled( + Map.of(STREAMING_SORTED_MAILBOX_RECEIVE, "1"))); + } + private static Object getValue(Map map, String key) { switch (key) { // Positive ints @@ -320,6 +336,10 @@ private static Object getValue(Map map, String key) { return QueryOptionsUtils.getMaxRowsInJoin(map); case MAX_ROWS_IN_WINDOW: return QueryOptionsUtils.getMaxRowsInWindow(map); + case SORTED_SELECTION_MERGE_BLOCK_SIZE: + return QueryOptionsUtils.getSortedSelectionMergeBlockSize(map); + case STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE: + return QueryOptionsUtils.getStreamingSortedMailboxReceiveBlockSize(map); // Non-negative ints case MULTI_STAGE_LEAF_LIMIT: return QueryOptionsUtils.getMultiStageLeafLimit(map); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java new file mode 100644 index 000000000000..b99d6f0edc70 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java @@ -0,0 +1,543 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.operator.combine; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.PriorityQueue; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.OrderByExpressionContext; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator; +import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock; +import org.apache.pinot.core.operator.blocks.results.MetadataResultsBlock; +import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock; +import org.apache.pinot.core.operator.query.StreamingSelectionOrderByOperator; +import org.apache.pinot.core.operator.streaming.BaseStreamingCombineOperator; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunctionFactory; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.selection.SelectionOperatorUtils; +import org.apache.pinot.core.query.utils.OrderByComparatorFactory; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryErrorMessage; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/// Streaming, lazy combine operator for selection ORDER BY queries whose first order-by expression is an identifier. +/// +/// It performs an incremental k-way heap merge across the per-segment operators in {@code _operators}, returning +/// globally sorted rows in bounded blocks. Each segment is exposed through a {@link SegmentCursor} that yields that +/// segment's locally-sorted rows in order: +/// +/// - Segments physically sorted on the first order-by column are backed by +/// {@link StreamingSelectionOrderByOperator}, which is pulled lazily one run/block at a time. +/// - Other (e.g. consuming/unsorted) segments are backed by a single materialized top-K block (any +/// {@link SelectionResultsBlock}-producing operator such as {@code SelectionOrderByOperator}); the cursor reads that +/// one block and iterates its rows. +/// +/// A {@link PriorityQueue} of {@link SegmentCursor} ordered by the {@link OrderByComparatorFactory} comparator on +/// each +/// cursor's current head row drives the merge with an at-most-one-head-per-active-segment invariant (the heap holds the +/// cursors themselves, never all rows, which would degenerate into a full heap-sort that materializes everything). Each +/// cycle pops the global-min cursor, appends its head to the current output block, advances that one cursor by a single +/// row, and re-offers it if it still has a head. +/// +/// **Min/max lazy segment activation (pruning).** Cursors are sorted by the first order-by column's min value +/// (ASC) / max value (DESC) reusing the {@code MinMaxValueContext} idea from +/// {@link MinMaxValueBasedSelectionOrderByCombineOperator}. A cursor is only activated (its segment acquired and first +/// block read) when the merge frontier reaches its min/max, so once {@code limit + offset} rows are emitted the +/// remaining segments are never acquired or read. See {@link #activateEligibleCursors()} for the correctness argument. +/// Pruning is disabled when null handling is enabled (an unsorted segment's first order-by column may then contain +/// nulls whose ordering position the raw min/max cannot capture), in which case every segment is activated. +/// +/// **Segment acquire/release lifecycle.** A cursor acquires its +/// {@link AcquireReleaseColumnsSegmentOperator} on activation and releases it only when its child operator is fully +/// drained (acquire-on-activate / release-on-exhaust), rather than per run. This is intentional: the backing +/// {@link StreamingSelectionOrderByOperator} retains a buffer-backed {@code ValueBlock} across {@code nextBlock()} +/// calls in its tail-to-sort mode, so releasing between interleaved runs could read segment buffers after a release +/// under prefetch. Holding the acquire for the cursor's lifetime guarantees no release happens between a cursor's +/// own reads; min/max pruning bounds the number of simultaneously-active (acquired) segments to the merge frontier. +/// The rows handed out by the child operators are already deep-copied to heap {@code Object[]} (via +/// {@code RowBasedBlockValueFetcher}), so they remain valid after the segment is released. Any cursors still +/// acquired when the merge ends early (LIMIT reached) or errors out are released via {@link #releaseAllCursors()}. +/// +/// **Streaming vs single-block.** When {@code _streaming} is {@code true} (MSE leaf path, driven by +/// {@link org.apache.pinot.core.operator.streaming.StreamingInstanceResponseOperator}) the merge emits many bounded +/// {@link SelectionResultsBlock}s from successive {@link #getNextBlock()} calls followed by a final +/// {@link MetadataResultsBlock}. When {@code false} (classic single-stage path) the merge runs to completion and the +/// first {@link #getNextBlock()} call returns a single block with execution stats attached. +/// +/// **Threading.** This operator overrides {@link #start()}/{@link #stop()} to no-ops (other than releasing +/// segments) and runs the merge single-threaded and lazily in {@link #getNextBlock()} on the consumer thread; it does +/// not use the base worker-queue model, and {@link #processSegments()} is overridden to fail loud. The base +/// {@code Phaser} (which exists only to fence worker threads against segment release) is intentionally bypassed because +/// all child/segment access is synchronous on the single consumer thread that holds the segment references; no async +/// work may be introduced here without restoring that fence. The instance is single-use (driven once to completion) and +/// is not thread-safe. +@SuppressWarnings({"rawtypes", "unchecked"}) +public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombineOperator { + private static final Logger LOGGER = LoggerFactory.getLogger(StreamingSelectionOrderByCombineOperator.class); + private static final String EXPLAIN_NAME = "COMBINE_SELECT_ORDERBY_STREAMING"; + + private final boolean _streaming; + private final boolean _asc; + private final boolean _pruningEnabled; + private final int _numRowsToKeep; + private final int _blockSize; + private final Comparator _comparator; + private final SegmentCursor[] _sortedCursors; + private final PriorityQueue _priorityQueue; + + // Merge progress (single-threaded; mutated only by the consumer thread driving getNextBlock()) + private int _nextToActivate; + private int _numRowsEmitted; + private List _outputRows; + private boolean _done; + /// Captured from the first child block seen; all child blocks share the same schema + private DataSchema _dataSchema; + /// Deduplicated MERGE_RESPONSE errors for segment blocks dropped on schema mismatch; null until the first mismatch + @Nullable + private Set _dataSchemaMismatchErrors; + /// Subset of the above not yet attached to an emitted block; drained on each attach + @Nullable + private List _unreportedDataSchemaMismatchErrors; + + public StreamingSelectionOrderByCombineOperator(List operators, QueryContext queryContext, + ExecutorService executorService, boolean streaming) { + // Pass a null merger: we override the consumption path entirely and never touch the base merger / worker queue. + super(null, operators, queryContext, executorService); + _streaming = streaming; + _numRowsToKeep = queryContext.getLimit() + queryContext.getOffset(); + // Streaming mode flushes bounded blocks; single-stage mode flushes once at the end as a single block. + _blockSize = streaming ? queryContext.getSortedSelectionMergeBlockSize() : Integer.MAX_VALUE; + _pruningEnabled = !queryContext.isNullHandlingEnabled(); + + List orderByExpressions = queryContext.getOrderByExpressions(); + assert orderByExpressions != null && !orderByExpressions.isEmpty(); + OrderByExpressionContext firstOrderByExpression = orderByExpressions.get(0); + assert firstOrderByExpression.getExpression().getType() == ExpressionContext.Type.IDENTIFIER; + _asc = firstOrderByExpression.isAsc(); + String firstOrderByColumn = firstOrderByExpression.getExpression().getIdentifier(); + _comparator = OrderByComparatorFactory.getComparator(orderByExpressions, queryContext.isNullHandlingEnabled()); + + // Build one cursor per segment operator and read its first order-by column min/max for lazy activation ordering. + // Reading DataSourceMetadata does not touch column buffers, so no segment acquire is needed here (mirrors + // MinMaxValueBasedSelectionOrderByCombineOperator). + _sortedCursors = new SegmentCursor[_numOperators]; + for (int i = 0; i < _numOperators; i++) { + Operator operator = _operators.get(i); + DataSourceMetadata metadata = + operator.getIndexSegment().getDataSource(firstOrderByColumn, queryContext.getSchema()) + .getDataSourceMetadata(); + _sortedCursors[i] = new SegmentCursor(operator, metadata.getMinValue(), metadata.getMaxValue()); + } + sortCursorsByMinMax(); + + _priorityQueue = new PriorityQueue<>(Math.max(1, _numOperators), + (o1, o2) -> _comparator.compare(o1.currentHead(), o2.currentHead())); + _outputRows = newOutputList(); + } + + /// Sorts the cursors so the merge can activate them lazily in frontier order: ascending by the column min value for + /// ASC, descending by the column max value for DESC. Cursors without a min/max are placed first because they must + /// always be processed (mirrors {@link MinMaxValueBasedSelectionOrderByCombineOperator}). + private void sortCursorsByMinMax() { + if (_asc) { + Arrays.sort(_sortedCursors, (o1, o2) -> { + if (o1._minValue == null) { + return o2._minValue == null ? 0 : -1; + } + if (o2._minValue == null) { + return 1; + } + return o1._minValue.compareTo(o2._minValue); + }); + } else { + Arrays.sort(_sortedCursors, (o1, o2) -> { + if (o1._maxValue == null) { + return o2._maxValue == null ? 0 : -1; + } + if (o2._maxValue == null) { + return 1; + } + return o2._maxValue.compareTo(o1._maxValue); + }); + } + } + + @Override + public String toExplainString() { + return EXPLAIN_NAME; + } + + /// Override to a no-op: the merge is single-threaded and lazy in {@link #getNextBlock()}, so we do not spin up the + /// base worker threads / blocking-queue model. + @Override + public void start() { + } + + /// Override the base worker-queue stop: no worker threads / phaser tasks were started. Release any segments still + /// acquired (idempotent) so an early stop by the driver cannot leak acquires. + @Override + public void stop() { + _done = true; + releaseAllCursors(); + } + + /// The base worker-thread entry point must never run here ({@link #start()} is a no-op). Fail loud if it ever does. + @Override + protected void processSegments() { + throw new IllegalStateException( + "StreamingSelectionOrderByCombineOperator runs single-threaded; processSegments() must not be called"); + } + + @Override + protected BaseResultsBlock getNextBlock() { + if (_done) { + // Streaming mode: terminal metadata block after the last data block. Idempotent if called again. + return attachExecutionStats(new MetadataResultsBlock()); + } + try { + long endTimeMs = _queryContext.getEndTimeMs(); + while (_numRowsEmitted < _numRowsToKeep) { + // The merge drains the heap on this thread and, for single-block cursors, may never re-enter a child operator. + // Without this check nothing would observe cancellation, pause, or the query deadline for up to + // (limit + offset) iterations -- a regression against MinMaxValueBasedSelectionOrderByCombineOperator, which + // this operator replaces on the same query shape. + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(_numRowsEmitted, EXPLAIN_NAME, endTimeMs); + activateEligibleCursors(); + SegmentCursor cursor = _priorityQueue.poll(); + if (cursor == null) { + // All active cursors exhausted (and pruning guarantees the rest cannot contribute). + break; + } + _outputRows.add(cursor.currentHead()); + _numRowsEmitted++; + cursor.advance(); + if (cursor.currentHead() != null) { + _priorityQueue.offer(cursor); + } + if (_streaming && _outputRows.size() >= _blockSize) { + return flushDataBlock(); + } + } + // Merge complete. + finish(); + if (!_outputRows.isEmpty()) { + // Streaming: the final partial data block (next call returns the terminal metadata block). + // Single-stage: the single complete block. + return flushDataBlock(); + } + if (_streaming) { + return attachDataSchemaMismatchErrors(attachExecutionStats(new MetadataResultsBlock())); + } + // Single-stage with no rows: still return a (single) block carrying the schema and execution stats. + return attachDataSchemaMismatchErrors(attachExecutionStats( + new SelectionResultsBlock(resolveDataSchema(), List.of(), _comparator, _queryContext))); + } catch (Exception e) { + _done = true; + releaseAllCursors(); + return createExceptionResultsBlockAndAttachExecutionStats(e, "merging sorted selection results"); + } catch (Throwable t) { + // An Error (e.g. OutOfMemoryError while accumulating output rows) must not leave segments acquired for the + // lifetime of the server. In single-stage mode there is no start()/stop() backstop around this operator. + _done = true; + releaseAllCursors(); + throw t; + } + } + + /// Activates not-yet-active cursors whose min/max value can still contribute before the current merge frontier. + /// + /// Cursors are visited in min/max-sorted order and {@code _nextToActivate} is advanced only when a cursor is + /// actually activated; a {@code break} merely defers the current cursor, which is re-evaluated against the (rising) + /// frontier on every subsequent call. Correctness: a not-yet-activated cursor whose first order-by value range starts + /// strictly beyond the current heap head cannot contain a row that sorts before that head (the first order-by column + /// is the primary sort key), so the head is the true global minimum and is safe to emit; the deferred cursor is + /// activated later, exactly when the frontier reaches its min/max. When the heap is empty the frontier is unknown, so + /// activation is forced (never pruned), which also drains any segments that sort entirely after the ones seen so far. + private void activateEligibleCursors() { + while (_nextToActivate < _sortedCursors.length) { + SegmentCursor cursor = _sortedCursors[_nextToActivate]; + if (_pruningEnabled) { + Comparable bound = _asc ? cursor._minValue : cursor._maxValue; + // A null bound means the segment must always be processed. Otherwise, only prune against a non-null frontier; + // if the head's first order-by value is null we cannot compare, so fall through and activate. + if (bound != null && !_priorityQueue.isEmpty()) { + Object headValue = _priorityQueue.peek().currentHead()[0]; + if (headValue != null) { + // Both come from the same first order-by column: the metadata min/max and the materialized row[0] share + // the column's stored type, so this comparison is type-safe (same assumption as MinMaxValueBased...). + int cmp = bound.compareTo(headValue); + if (_asc ? cmp > 0 : cmp < 0) { + break; + } + } + } + } + cursor.activate(); + _nextToActivate++; + if (cursor.currentHead() != null) { + _priorityQueue.offer(cursor); + } + } + } + + /// Marks the merge done and releases any segments still held by un-drained cursors (e.g. when LIMIT is reached). + private void finish() { + _done = true; + releaseAllCursors(); + } + + private void releaseAllCursors() { + for (SegmentCursor cursor : _sortedCursors) { + cursor.release(); + } + } + + /// Returns the accumulated output rows as a sorted {@link SelectionResultsBlock} and resets the output buffer. The + /// block carries the comparator so the broker-side n-way reduce stays correct. In single-stage mode it is the only + /// block, so execution stats are attached; in streaming mode stats are attached to the terminal metadata block. + private BaseResultsBlock flushDataBlock() { + List rows = _outputRows; + _outputRows = newOutputList(); + SelectionResultsBlock block = new SelectionResultsBlock(resolveDataSchema(), rows, _comparator, _queryContext); + attachDataSchemaMismatchErrors(block); + return _streaming ? block : attachExecutionStats(block); + } + + /// Records that a segment's block was dropped because its schema disagreed with the merge schema. Deduplicated by + /// message so a mid-reload table with many divergent segments cannot flood the response. + private void recordDataSchemaMismatch(@Nullable DataSchema mismatched) { + String errorMessage = + String.format("Data schema mismatch between merged block: %s and block to merge: %s, drop block to merge", + _dataSchema, mismatched); + // NOTE: This is segment level log, so log at debug level to prevent flooding the log. + LOGGER.debug(errorMessage); + if (_dataSchemaMismatchErrors == null) { + _dataSchemaMismatchErrors = new HashSet<>(); + _unreportedDataSchemaMismatchErrors = new ArrayList<>(); + } + if (_dataSchemaMismatchErrors.add(errorMessage)) { + _unreportedDataSchemaMismatchErrors.add(errorMessage); + } + } + + /// Attaches mismatch errors recorded since the last emitted block. In streaming mode blocks are emitted many times, + /// so errors are drained rather than re-attached, and the terminal block picks up any recorded after the last flush. + private T attachDataSchemaMismatchErrors(T block) { + if (_unreportedDataSchemaMismatchErrors != null && !_unreportedDataSchemaMismatchErrors.isEmpty()) { + for (String errorMessage : _unreportedDataSchemaMismatchErrors) { + block.addErrorMessage(QueryErrorMessage.safeMsg(QueryErrorCode.MERGE_RESPONSE, errorMessage)); + } + _unreportedDataSchemaMismatchErrors.clear(); + } + return block; + } + + private List newOutputList() { + int capacity = + Math.min(_blockSize, Math.min(_numRowsToKeep, SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY)); + return new ArrayList<>(Math.max(1, capacity)); + } + + /// Returns the data schema captured from the first child block. If no segment produced a block (every segment is a + /// streaming operator that matched zero rows), reconstructs the schema from the first segment so that an empty result + /// still carries a valid, correctly-ordered schema (order-by expressions first, matching the child blocks' layout). + private DataSchema resolveDataSchema() { + if (_dataSchema != null) { + return _dataSchema; + } + IndexSegment indexSegment = _operators.get(0).getIndexSegment(); + List expressions = SelectionOperatorUtils.extractExpressions(_queryContext, indexSegment); + Set columns = new HashSet<>(); + for (ExpressionContext expression : expressions) { + expression.getColumns(columns); + } + Map dataSourceMap = new HashMap<>(); + for (String column : columns) { + dataSourceMap.put(column, indexSegment.getDataSource(column, _queryContext.getSchema())); + } + int numExpressions = expressions.size(); + String[] columnNames = new String[numExpressions]; + ColumnDataType[] columnDataTypes = new ColumnDataType[numExpressions]; + for (int i = 0; i < numExpressions; i++) { + ExpressionContext expression = expressions.get(i); + columnNames[i] = expression.toString(); + TransformFunction transformFunction = TransformFunctionFactory.get(expression, dataSourceMap); + columnDataTypes[i] = ColumnDataType.fromDataType(transformFunction.getResultMetadata().getDataType(), + transformFunction.getResultMetadata().isSingleValue()); + } + _dataSchema = new DataSchema(columnNames, columnDataTypes); + return _dataSchema; + } + + /// Iterates a single segment's locally-sorted rows, pulling blocks lazily from its operator. Streaming-backed cursors + /// loop until the operator returns {@code null}; single-block-backed cursors read exactly one block. The segment is + /// acquired on activation and released once exhausted or when the combine finishes (see + /// {@link StreamingSelectionOrderByCombineOperator}). + private class SegmentCursor { + private final Operator _operator; + @Nullable + private final Comparable _minValue; + @Nullable + private final Comparable _maxValue; + + private boolean _streamingChild; + private List _rows; + private int _pos; + private boolean _acquired; + /// Cached current head (rows.get(pos)) so the hot heap comparator does not re-index per comparison + @Nullable + private Object[] _head; + + SegmentCursor(Operator operator, @Nullable Comparable minValue, @Nullable Comparable maxValue) { + _operator = operator; + _minValue = minValue; + _maxValue = maxValue; + } + + /// Returns the current head row to be merged next, or {@code null} if not activated or exhausted. + @Nullable + Object[] currentHead() { + return _head; + } + + /// Acquires the segment, resolves whether the child is the lazy streaming operator, and reads the first block. + /// After this call {@link #currentHead()} returns the first row, or {@code null} if the segment contributes + /// nothing. + void activate() { + acquireSegment(); + _streamingChild = isStreamingChild(); + if (!pullBlock()) { + exhaust(); + } + } + + /// Advances past the current head, pulling the next block lazily for streaming cursors. Releases the segment + /// when the cursor is exhausted; afterwards {@link #currentHead()} returns {@code null}. + void advance() { + _pos++; + if (_pos < _rows.size()) { + _head = _rows.get(_pos); + return; + } + // Current block drained: streaming cursors pull the next run/block; single-block cursors are done. + if (_streamingChild && pullBlock()) { + return; + } + exhaust(); + } + + /// Loads the next non-empty block of rows from the operator, capturing the combine-level data schema on first + /// sight. Returns {@code false} when the operator is exhausted (no more rows). Streaming-backed operators emit + /// one run/block per call and {@code null} when done; single-block operators emit a single block and must not be + /// called again afterwards, so an empty/null block from a single-block child is treated as exhausted. + /// + /// A block whose schema differs from the one already captured is dropped and reported, mirroring + /// {@link org.apache.pinot.core.operator.combine.merger.SelectionOrderByResultsBlockMerger}. Segments on a server + /// can disagree on schema mid-reload (a newly added column exists only in reloaded segments), and merging rows of + /// differing width under one schema would corrupt the result rather than fail. + private boolean pullBlock() { + while (true) { + SelectionResultsBlock block = nextBlock(); + if (block == null) { + return false; + } + if (_dataSchema == null) { + _dataSchema = block.getDataSchema(); + } else if (!_dataSchema.equals(block.getDataSchema())) { + recordDataSchemaMismatch(block.getDataSchema()); + return false; + } + List rows = block.getRows(); + if (rows != null && !rows.isEmpty()) { + _rows = rows; + _pos = 0; + _head = rows.get(0); + return true; + } + // Defensive: an unexpected empty (non-null) block. Keep pulling only for streaming children; a single-block + // child yields exactly one block, so treat it as exhausted. + if (!_streamingChild) { + return false; + } + } + } + + private SelectionResultsBlock nextBlock() { + try { + return (SelectionResultsBlock) _operator.nextBlock(); + } catch (RuntimeException e) { + throw wrapOperatorException(_operator, e); + } + } + + /// Returns whether the underlying child operator is the lazy {@link StreamingSelectionOrderByOperator}. Must be + /// called after {@link #acquireSegment()} because materializing the wrapped child runs the plan node, which + /// accesses segment buffers. + private boolean isStreamingChild() { + Operator underlying = _operator; + if (_operator instanceof AcquireReleaseColumnsSegmentOperator) { + AcquireReleaseColumnsSegmentOperator wrapper = (AcquireReleaseColumnsSegmentOperator) _operator; + wrapper.materializeChildOperator(); + underlying = wrapper.getChildOperators().get(0); + } + return underlying instanceof StreamingSelectionOrderByOperator; + } + + private void acquireSegment() { + if (_operator instanceof AcquireReleaseColumnsSegmentOperator) { + ((AcquireReleaseColumnsSegmentOperator) _operator).acquire(); + } + _acquired = true; + } + + /// Releases the segment if still held. Idempotent: safe to call from {@link #exhaust()} and combine cleanup. + private void release() { + if (_acquired) { + if (_operator instanceof AcquireReleaseColumnsSegmentOperator) { + ((AcquireReleaseColumnsSegmentOperator) _operator).release(); + } + _acquired = false; + } + } + + private void exhaust() { + release(); + _rows = null; + _head = null; + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java new file mode 100644 index 000000000000..e66a6ee0391f --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java @@ -0,0 +1,514 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.operator.query; + +import com.google.common.base.CaseFormat; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.PriorityQueue; +import java.util.Set; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.OrderByExpressionContext; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.common.RowBasedBlockValueFetcher; +import org.apache.pinot.core.operator.BaseOperator; +import org.apache.pinot.core.operator.BaseProjectOperator; +import org.apache.pinot.core.operator.BitmapDocIdSetOperator; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.operator.ExecutionStatistics; +import org.apache.pinot.core.operator.ExplainAttributeBuilder; +import org.apache.pinot.core.operator.ProjectionOperator; +import org.apache.pinot.core.operator.ProjectionOperatorUtils; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock; +import org.apache.pinot.core.operator.transform.TransformOperator; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.selection.SelectionOperatorUtils; +import org.apache.pinot.core.query.utils.OrderByComparatorFactory; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.spi.query.QueryScanCostContext; +import org.roaringbitmap.RoaringBitmap; + + +/// Lazy, incremental selection ORDER BY operator for segments that are physically sorted on the first order-by column. +/// +/// Unlike {@link SelectionOrderByOperator} (which materializes the segment's whole top-K in a single block) this +/// operator emits one globally-sorted {@link SelectionResultsBlock} per {@link #getNextBlock()} call and returns +/// {@code null} when the segment is exhausted, so that a downstream k-way-merge combine operator can pull from many +/// segments lazily and stop early. It relies on the underlying project operator iterating the first order-by column in +/// the query order (the caller must guarantee {@code projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))}). +/// +/// It runs in one of two emission modes: +/// +/// - **No tail to sort** ({@code numSortedExpressions == numOrderByExpressions}, e.g. {@code ORDER BY sorted}): +/// rows already arrive from the project operator in final order, so each call emits the next project block +/// (trimmed to the remaining {@code limit + offset} budget). +/// - **Tail to sort** ({@code numSortedExpressions < numOrderByExpressions}, e.g. +/// {@code ORDER BY sorted, other}): +/// each call reads forward until the first order-by value changes (a primary-value "run"), retains the run's top +/// {@code limit + offset} rows by the full comparator, and emits them sorted. This bounds the in-memory run buffer to +/// {@code limit + offset} rows even when the first order-by column is near-constant (very low cardinality). +/// +/// Like {@link SelectionOrderByOperator} it preserves the two-phase projection optimization: when there are output +/// expressions that are not order-by expressions, the forward scan only fetches the order-by expressions plus the +/// document id, and the non-order-by expressions are fetched in a second pass over the retained document ids of each +/// emitted block. +/// +/// This operator is stateful across {@link #getNextBlock()} calls and is **not** thread-safe; a single consumer +/// must drive it. +public class StreamingSelectionOrderByOperator extends BaseOperator { + private static final String EXPLAIN_NAME = "SELECT_ORDERBY_STREAMING"; + + private final IndexSegment _indexSegment; + private final QueryContext _queryContext; + private final boolean _nullHandlingEnabled; + /// Deduped order-by expressions followed by output expressions from SelectionOperatorUtils.extractExpressions() + private final List _expressions; + private final BaseProjectOperator _projectOperator; + private final List _orderByExpressions; + private final ColumnContext[] _orderByColumnContexts; + private final int _numExpressions; + private final int _numOrderByExpressions; + private final int _numRowsToKeep; + /// Whether there are output expressions that are not order-by expressions (requires the two-phase fetch) + private final boolean _twoPhase; + /// Whether the order-by has an unsorted tail that must be sorted in memory per run + private final boolean _tailToSort; + /// Expressions fetched during the forward scan: order-by expressions only when two-phase, otherwise all expressions + private final List _phase1Expressions; + private final int _numPhase1Columns; + private final Comparator _comparator; + /// Compares only the first order-by column; used to detect primary-value run boundaries + private final Comparator _primaryComparator; + /// Pre-allocated run heap (cleared and reused each nextRun() call to avoid per-run allocation) + private final Comparator _reversedComparator; + private final PriorityQueue _runHeap; + + // Pre-computed invariants for the two-phase fetch (null when single-phase) + private final List _nonOrderByExpressions; + private final Map _phase2DataSourceMap; + private final int _phase2NumColumns; + + /// Lazily built and cached; for two-phase it requires the transform operator's result column contexts + private DataSchema _dataSchema; + + // Forward-scan cursor state (used by the tail-to-sort mode) + private ValueBlock _currentBlock; + private RowBasedBlockValueFetcher _currentFetcher; + private int[] _currentDocIds; + private RoaringBitmap[] _currentNullBitmaps; + private int _currentNumDocs; + private int _currentPos; + /// One-row lookahead: the first row of the next run, stashed when a run boundary is crossed + private Object[] _pendingRow; + private boolean _projectExhausted; + + private boolean _exhausted; + private int _numRowsEmitted; + private int _numDocsScanned = 0; + private long _numEntriesScannedPostFilter = 0; + + public StreamingSelectionOrderByOperator(IndexSegment indexSegment, QueryContext queryContext, + List expressions, BaseProjectOperator projectOperator, int numSortedExpressions) { + _indexSegment = indexSegment; + _queryContext = queryContext; + _nullHandlingEnabled = queryContext.isNullHandlingEnabled(); + _expressions = expressions; + _projectOperator = projectOperator; + + _orderByExpressions = queryContext.getOrderByExpressions(); + assert _orderByExpressions != null; + _numExpressions = expressions.size(); + _numOrderByExpressions = _orderByExpressions.size(); + _orderByColumnContexts = new ColumnContext[_numOrderByExpressions]; + for (int i = 0; i < _numOrderByExpressions; i++) { + ExpressionContext expression = _orderByExpressions.get(i).getExpression(); + _orderByColumnContexts[i] = _projectOperator.getResultColumnContext(expression); + } + + _numRowsToKeep = queryContext.getOffset() + queryContext.getLimit(); + _twoPhase = _numExpressions > _numOrderByExpressions; + _tailToSort = numSortedExpressions < _numOrderByExpressions; + _comparator = + OrderByComparatorFactory.getComparator(_orderByExpressions, _orderByColumnContexts, _nullHandlingEnabled); + // The first order-by column is the physically sorted column, so it never contains nulls on this path; comparing + // only index 0 is enough to detect when one primary-value run ends and the next begins. + _primaryComparator = + OrderByComparatorFactory.getComparator(_orderByExpressions, _orderByColumnContexts, _nullHandlingEnabled, 0, 1); + _reversedComparator = _comparator.reversed(); + _runHeap = new PriorityQueue<>( + Math.min(_numRowsToKeep, SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY), _reversedComparator); + + if (_twoPhase) { + _phase1Expressions = new ArrayList<>(_numOrderByExpressions); + for (OrderByExpressionContext orderByExpression : _orderByExpressions) { + _phase1Expressions.add(orderByExpression.getExpression()); + } + _nonOrderByExpressions = _expressions.subList(_numOrderByExpressions, _numExpressions); + Set columns = new HashSet<>(); + for (ExpressionContext expressionContext : _nonOrderByExpressions) { + expressionContext.getColumns(columns); + } + _phase2NumColumns = columns.size(); + _phase2DataSourceMap = new HashMap<>(); + for (String column : columns) { + _phase2DataSourceMap.put(column, _indexSegment.getDataSource(column, _queryContext.getSchema())); + } + } else { + _phase1Expressions = _expressions; + _nonOrderByExpressions = null; + _phase2NumColumns = 0; + _phase2DataSourceMap = null; + // Single-phase: all output expressions are order-by expressions, so their types are known up front. + _dataSchema = buildSinglePhaseDataSchema(); + } + _numPhase1Columns = _phase1Expressions.size(); + } + + @Override + protected SelectionResultsBlock getNextBlock() { + if (_exhausted) { + return null; + } + List rows = _tailToSort ? nextRun() : nextSortedRows(); + if (rows == null || rows.isEmpty()) { + _exhausted = true; + return null; + } + if (_twoPhase) { + fetchNonOrderByColumns(rows); + } + // Single-phase builds the schema in the constructor; two-phase builds it during fetchNonOrderByColumns above. + assert _dataSchema != null; + return new SelectionResultsBlock(_dataSchema, rows, _comparator, _queryContext); + } + + /// No-tail-to-sort mode: the project operator already returns rows in final order, so emit the next project block, + /// trimmed to the remaining {@code limit + offset} budget. Returns {@code null} when exhausted. + @Nullable + private List nextSortedRows() { + int remaining = _numRowsToKeep - _numRowsEmitted; + if (remaining <= 0) { + return null; + } + ValueBlock valueBlock = _projectOperator.nextBlock(); + if (valueBlock == null) { + return null; + } + int numDocsFetched = valueBlock.getNumDocs(); + BlockValSet[] blockValSets = new BlockValSet[_numPhase1Columns]; + for (int i = 0; i < _numPhase1Columns; i++) { + blockValSets[i] = valueBlock.getBlockValueSet(_phase1Expressions.get(i)); + } + RowBasedBlockValueFetcher blockValueFetcher = new RowBasedBlockValueFetcher(blockValSets); + int[] docIds = _twoPhase ? valueBlock.getDocIds() : null; + RoaringBitmap[] nullBitmaps = null; + if (_nullHandlingEnabled) { + nullBitmaps = new RoaringBitmap[_numPhase1Columns]; + for (int i = 0; i < _numPhase1Columns; i++) { + nullBitmaps[i] = blockValSets[i].getNullBitmap(); + } + } + _numDocsScanned += numDocsFetched; + _numEntriesScannedPostFilter += (long) numDocsFetched * _projectOperator.getNumColumnsProjected(); + reportScanCost(numDocsFetched, (long) numDocsFetched * _projectOperator.getNumColumnsProjected()); + + // Rows arrive sorted; we only need the first 'remaining' of them globally. + int numRows = Math.min(numDocsFetched, remaining); + List rows = new ArrayList<>(numRows); + for (int i = 0; i < numRows; i++) { + rows.add(materializeRow(blockValueFetcher, docIds, nullBitmaps, i)); + } + _numRowsEmitted += rows.size(); + return rows; + } + + /// Tail-to-sort mode: read forward until the first order-by value changes, retain the run's top + /// {@code limit + offset} rows by the full comparator, and return them sorted. Returns {@code null} when + /// exhausted. + @Nullable + private List nextRun() { + int remaining = _numRowsToKeep - _numRowsEmitted; + if (remaining <= 0) { + return null; + } + if (_pendingRow == null) { + _pendingRow = nextRow(); + if (_pendingRow == null) { + return null; + } + } + PriorityQueue runHeap = _runHeap; + runHeap.clear(); + Object[] runFirstRow = _pendingRow; + SelectionOperatorUtils.addToPriorityQueue(_pendingRow, runHeap, _numRowsToKeep); + _pendingRow = null; + Object[] row; + while ((row = nextRow()) != null) { + if (_primaryComparator.compare(row, runFirstRow) == 0) { + SelectionOperatorUtils.addToPriorityQueue(row, runHeap, _numRowsToKeep); + } else { + // Run boundary: this row starts the next run, keep it for the next call. + _pendingRow = row; + break; + } + } + List rows = drainAscending(runHeap); + // A segment never contributes more than 'limit + offset' rows to the global result, and they are a prefix of its + // local sorted order, so cap the total emitted across runs at the remaining budget (the rows are ascending, keep + // the smallest 'remaining'). + if (rows.size() > remaining) { + rows = rows.subList(0, remaining); + } + _numRowsEmitted += rows.size(); + return rows; + } + + /// Pulls the next row of the forward scan (across project blocks), materialized as an + /// {@code Object[_numExpressions]}. For two-phase the document id is stashed at index + /// {@code _numOrderByExpressions} (overwritten in the second pass). Returns {@code null} when the project operator + /// is exhausted. + @Nullable + private Object[] nextRow() { + while (true) { + if (_currentBlock == null || _currentPos >= _currentNumDocs) { + if (_projectExhausted) { + return null; + } + _currentBlock = _projectOperator.nextBlock(); + if (_currentBlock == null) { + _projectExhausted = true; + return null; + } + BlockValSet[] blockValSets = new BlockValSet[_numPhase1Columns]; + for (int i = 0; i < _numPhase1Columns; i++) { + blockValSets[i] = _currentBlock.getBlockValueSet(_phase1Expressions.get(i)); + } + _currentFetcher = new RowBasedBlockValueFetcher(blockValSets); + _currentNumDocs = _currentBlock.getNumDocs(); + _currentDocIds = _twoPhase ? _currentBlock.getDocIds() : null; + if (_nullHandlingEnabled) { + _currentNullBitmaps = new RoaringBitmap[_numPhase1Columns]; + for (int i = 0; i < _numPhase1Columns; i++) { + _currentNullBitmaps[i] = blockValSets[i].getNullBitmap(); + } + } + _currentPos = 0; + _numDocsScanned += _currentNumDocs; + _numEntriesScannedPostFilter += (long) _currentNumDocs * _projectOperator.getNumColumnsProjected(); + reportScanCost(_currentNumDocs, (long) _currentNumDocs * _projectOperator.getNumColumnsProjected()); + if (_currentNumDocs == 0) { + _currentBlock = null; + continue; + } + } + int rowId = _currentPos++; + return materializeRow(_currentFetcher, _currentDocIds, _currentNullBitmaps, rowId); + } + } + + /// Materializes a single phase-1 row (deep-copied out of the value block buffers) from the given fetcher. + private Object[] materializeRow(RowBasedBlockValueFetcher fetcher, @Nullable int[] docIds, + @Nullable RoaringBitmap[] nullBitmaps, int rowId) { + Object[] row = new Object[_numExpressions]; + fetcher.getRow(rowId, row, 0); + if (_twoPhase) { + row[_numOrderByExpressions] = docIds[rowId]; + } + if (_nullHandlingEnabled) { + for (int colId = 0; colId < _numPhase1Columns; colId++) { + if (nullBitmaps[colId] != null && nullBitmaps[colId].contains(rowId)) { + row[colId] = null; + } + } + } + return row; + } + + /// Drains a max-heap (created with the reversed comparator) into an ascending list, mutable so the second pass can + /// fill non-order-by values in place. + private List drainAscending(PriorityQueue heap) { + int numRows = heap.size(); + Object[][] sortedRows = new Object[numRows][]; + for (int i = numRows - 1; i >= 0; i--) { + sortedRows[i] = heap.poll(); + } + return Arrays.asList(sortedRows); + } + + /// Second pass of the two-phase fetch: fills the non-order-by expression values for the rows of a single emitted + /// block. + /// The rows keep their final (comparator) order; the fill iterates a document-id-sorted view that shares the same row + /// instances, mirroring {@link SelectionOrderByOperator#computePartiallyOrdered()}. + private void fetchNonOrderByColumns(List rows) { + int numRows = rows.size(); + RoaringBitmap docIds = new RoaringBitmap(); + for (Object[] row : rows) { + docIds.add((int) row[_numOrderByExpressions]); + } + // Document-id-sorted view sharing the same row instances (the bitmap returns docIds in ascending order). + List rowsByDocId = new ArrayList<>(rows); + rowsByDocId.sort(Comparator.comparingInt(o -> (int) o[_numOrderByExpressions])); + + BitmapDocIdSetOperator docIdOperator = BitmapDocIdSetOperator.ascending(docIds, numRows); + try (ProjectionOperator projectionOperator = + ProjectionOperatorUtils.getProjectionOperator(_phase2DataSourceMap, docIdOperator, _queryContext)) { + TransformOperator transformOperator = + new TransformOperator(_queryContext, projectionOperator, _nonOrderByExpressions); + + int numNonOrderByExpressions = _nonOrderByExpressions.size(); + BlockValSet[] blockValSets = new BlockValSet[numNonOrderByExpressions]; + int rowBaseId = 0; + ValueBlock valueBlock; + while ((valueBlock = transformOperator.nextBlock()) != null) { + for (int i = 0; i < numNonOrderByExpressions; i++) { + blockValSets[i] = valueBlock.getBlockValueSet(_nonOrderByExpressions.get(i)); + } + RowBasedBlockValueFetcher blockValueFetcher = new RowBasedBlockValueFetcher(blockValSets); + int numDocsFetched = valueBlock.getNumDocs(); + for (int i = 0; i < numDocsFetched; i++) { + blockValueFetcher.getRow(i, rowsByDocId.get(rowBaseId + i), _numOrderByExpressions); + } + if (_nullHandlingEnabled) { + RoaringBitmap[] nullBitmaps = new RoaringBitmap[numNonOrderByExpressions]; + for (int i = 0; i < numNonOrderByExpressions; i++) { + nullBitmaps[i] = blockValSets[i].getNullBitmap(); + } + for (int i = 0; i < numDocsFetched; i++) { + Object[] values = rowsByDocId.get(rowBaseId + i); + for (int colId = 0; colId < numNonOrderByExpressions; colId++) { + if (nullBitmaps[colId] != null && nullBitmaps[colId].contains(i)) { + values[_numOrderByExpressions + colId] = null; + } + } + } + } + _numEntriesScannedPostFilter += (long) numDocsFetched * _phase2NumColumns; + // Phase 2 re-reads docs already counted in phase 1, so only the extra entries are reported. + reportScanCost(0, (long) numDocsFetched * _phase2NumColumns); + rowBaseId += numDocsFetched; + } + + if (_dataSchema == null) { + _dataSchema = buildTwoPhaseDataSchema(transformOperator); + } + } + } + + private DataSchema buildSinglePhaseDataSchema() { + String[] columnNames = new String[_numExpressions]; + DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[_numExpressions]; + for (int i = 0; i < _numExpressions; i++) { + columnNames[i] = _expressions.get(i).toString(); + columnDataTypes[i] = DataSchema.ColumnDataType.fromDataType(_orderByColumnContexts[i].getDataType(), + _orderByColumnContexts[i].isSingleValue()); + } + return new DataSchema(columnNames, columnDataTypes); + } + + private DataSchema buildTwoPhaseDataSchema(TransformOperator transformOperator) { + int numNonOrderByExpressions = _nonOrderByExpressions.size(); + String[] columnNames = new String[_numExpressions]; + DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[_numExpressions]; + for (int i = 0; i < _numExpressions; i++) { + columnNames[i] = _expressions.get(i).toString(); + } + for (int i = 0; i < _numOrderByExpressions; i++) { + columnDataTypes[i] = DataSchema.ColumnDataType.fromDataType(_orderByColumnContexts[i].getDataType(), + _orderByColumnContexts[i].isSingleValue()); + } + for (int i = 0; i < numNonOrderByExpressions; i++) { + ColumnContext columnContext = transformOperator.getResultColumnContext(_nonOrderByExpressions.get(i)); + columnDataTypes[_numOrderByExpressions + i] = + DataSchema.ColumnDataType.fromDataType(columnContext.getDataType(), columnContext.isSingleValue()); + } + return new DataSchema(columnNames, columnDataTypes); + } + + @Override + public String toExplainString() { + StringBuilder stringBuilder = new StringBuilder(EXPLAIN_NAME).append("(selectList:"); + if (!_expressions.isEmpty()) { + stringBuilder.append(_expressions.get(0)); + for (int i = 1; i < _expressions.size(); i++) { + stringBuilder.append(", ").append(_expressions.get(i)); + } + } + return stringBuilder.append(')').toString(); + } + + @Override + protected String getExplainName() { + return CaseFormat.UPPER_UNDERSCORE.to(CaseFormat.UPPER_CAMEL, EXPLAIN_NAME); + } + + @Override + protected void explainAttributes(ExplainAttributeBuilder attributeBuilder) { + super.explainAttributes(attributeBuilder); + if (_expressions.isEmpty()) { + return; + } + attributeBuilder.putStringList("selectList", + _expressions.stream().map(ExpressionContext::toString).collect(Collectors.toList())); + } + + @Override + public List getChildOperators() { + return Collections.singletonList(_projectOperator); + } + + @Override + public IndexSegment getIndexSegment() { + return _indexSegment; + } + + @Override + public ExecutionStatistics getExecutionStatistics() { + long numEntriesScannedInFilter = _projectOperator.getExecutionStatistics().getNumEntriesScannedInFilter(); + int numTotalDocs = _indexSegment.getSegmentMetadata().getTotalDocs(); + return new ExecutionStatistics(_numDocsScanned, numEntriesScannedInFilter, _numEntriesScannedPostFilter, + numTotalDocs); + } + + /// Reports scan cost to the shared {@link QueryScanCostContext} so scan-based query killing + /// ({@link org.apache.pinot.core.common.Operator#nextBlock()} -> {@code checkScanBasedKilling}) can see this + /// operator's work. Without this the killer is blind on the streaming path, unlike every other selection operator. + private void reportScanCost(int numDocsScanned, long numEntriesScannedPostFilter) { + QueryScanCostContext scanCost = getScanCostContext(); + if (scanCost != null) { + if (numDocsScanned > 0) { + scanCost.addDocsScanned(numDocsScanned); + } + if (numEntriesScannedPostFilter > 0) { + scanCost.addEntriesScannedPostFilter(numEntriesScannedPostFilter); + } + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java index 195b19e90176..aef488174cea 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java @@ -34,6 +34,7 @@ import org.apache.pinot.core.operator.combine.SelectionOrderByCombineOperator; import org.apache.pinot.core.operator.combine.SequentialSortedGroupByCombineOperator; import org.apache.pinot.core.operator.combine.SortedGroupByCombineOperator; +import org.apache.pinot.core.operator.combine.StreamingSelectionOrderByCombineOperator; import org.apache.pinot.core.operator.streaming.StreamingGroupByCombineOperator; import org.apache.pinot.core.operator.streaming.StreamingSelectionOnlyCombineOperator; import org.apache.pinot.core.query.executor.ResultsBlockStreamer; @@ -134,6 +135,17 @@ private BaseCombineOperator getCombineOperator() { // Use streaming operator only for non-empty selection-only query return new StreamingSelectionOnlyCombineOperator(operators, _queryContext, _executorService); } + // Streaming selection order-by (opt-in via the sortedSelectionMergeEnabled hint). Selection-only already + // returned above, so reaching here with a non-empty limit and an order-by present implies selection order-by. + if (_queryContext.isSortedSelectionMergeEnabled() && QueryContextUtils.isSelectionQuery(_queryContext) + && _queryContext.getLimit() != 0) { + List orderByExpressions = _queryContext.getOrderByExpressions(); + if (orderByExpressions != null + && orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) { + return new StreamingSelectionOrderByCombineOperator(operators, _queryContext, _executorService, + true /* streaming */); + } + } int flushThreshold = _queryContext.getStreamingGroupByFlushThreshold(); if (flushThreshold > 0 && QueryContextUtils.isAggregationQuery(_queryContext) && _queryContext.getGroupByExpressions() != null) { @@ -165,6 +177,10 @@ private BaseCombineOperator getCombineOperator() { List orderByExpressions = _queryContext.getOrderByExpressions(); assert orderByExpressions != null; if (orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) { + if (_queryContext.isSortedSelectionMergeEnabled()) { + return new StreamingSelectionOrderByCombineOperator(operators, _queryContext, _executorService, + false /* streaming */); + } return new MinMaxValueBasedSelectionOrderByCombineOperator(operators, _queryContext, _executorService); } else { return new SelectionOrderByCombineOperator(operators, _queryContext, _executorService); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java index fdb68d9866fe..2633820cd870 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java @@ -32,6 +32,7 @@ import org.apache.pinot.core.operator.query.SelectionOrderByOperator; import org.apache.pinot.core.operator.query.SelectionPartiallyOrderedByDescOperation; import org.apache.pinot.core.operator.query.SelectionPartiallyOrderedByLinearOperator; +import org.apache.pinot.core.operator.query.StreamingSelectionOrderByOperator; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.query.selection.SelectionOperatorUtils; import org.apache.pinot.segment.spi.IndexSegment; @@ -90,11 +91,36 @@ public Operator run() { maxDocsPerCall = Math.min(limit + _queryContext.getOffset(), DocIdSetPlanNode.MAX_DOC_PER_CALL); } - BaseProjectOperator projectOperator = getSortedByProject(expressions, maxDocsPerCall, orderByExpressions); boolean asc = orderByExpressions.get(0).isAsc(); // Remember that we cannot use asc == projectOperator.isAscending() because empty operators are considered // both ascending and descending DocIdOrderedOperator.DocIdOrder queryOrder = DocIdOrderedOperator.DocIdOrder.fromAsc(asc); + + // Opt-in streaming path: emit one globally-sorted block at a time so a downstream k-way-merge combine can pull + // lazily. Only build it when the first order-by column is an identifier (kept consistent with the combine-side + // gate) and the forward-scan project is order-compatible; the DESC-incompatible sorted case still falls back to + // the materialized SelectionPartiallyOrderedByDescOperation below so global order stays correct. + if (_queryContext.isSortedSelectionMergeEnabled() + && orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) { + // When there are non-order-by output expressions, only fetch the order-by expressions during the forward scan + // (the streaming operator fetches the rest in a second pass); otherwise fetch all expressions. + List projectExpressions = expressions; + if (expressions.size() > numOrderByExpressions) { + projectExpressions = new ArrayList<>(numOrderByExpressions); + for (OrderByExpressionContext orderByExpression : orderByExpressions) { + projectExpressions.add(orderByExpression.getExpression()); + } + } + BaseProjectOperator streamingProjectOperator = + getSortedByProject(projectExpressions, maxDocsPerCall, orderByExpressions); + if (streamingProjectOperator.isCompatibleWith(queryOrder)) { + return new StreamingSelectionOrderByOperator(_indexSegment, _queryContext, expressions, + streamingProjectOperator, sortedColumnsPrefixSize); + } + // DESC-incompatible: fall through to the materialized fallback (rebuilds the project over all expressions). + } + + BaseProjectOperator projectOperator = getSortedByProject(expressions, maxDocsPerCall, orderByExpressions); if (projectOperator.isCompatibleWith(queryOrder)) { return new SelectionPartiallyOrderedByLinearOperator(_indexSegment, _queryContext, expressions, projectOperator, sortedColumnsPrefixSize); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java index efb7aca3755b..9ef94f62c07c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java @@ -271,6 +271,17 @@ void applyQueryOptions(QueryContext queryContext) { } queryContext.setMaxExecutionThreads(maxExecutionThreads); + // Set streaming selection order-by options (opt-in; gated on selection queries to prevent accidental routing + // if a downstream guard is ever missed) + if (QueryContextUtils.isSelectionQuery(queryContext)) { + queryContext.setSortedSelectionMergeEnabled(QueryOptionsUtils.isSortedSelectionMergeEnabled(queryOptions)); + Integer sortedSelectionMergeBlockSize = + QueryOptionsUtils.getSortedSelectionMergeBlockSize(queryOptions); + if (sortedSelectionMergeBlockSize != null) { + queryContext.setSortedSelectionMergeBlockSize(sortedSelectionMergeBlockSize); + } + } + // Set group-by query options if (QueryContextUtils.isAggregationQuery(queryContext) && queryContext.getGroupByExpressions() != null) { // Set maxInitialResultHolderCapacity diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java index 25c03ffebda3..d989404ba7af 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java @@ -44,6 +44,7 @@ import org.apache.pinot.segment.spi.datasource.DataSource; import org.apache.pinot.spi.config.table.FieldConfig; import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.CommonConstants.Broker; import org.apache.pinot.spi.utils.CommonConstants.Server; @@ -143,6 +144,11 @@ public class QueryContext { private int _effectiveSegmentGroupTrimSize; // Flush threshold for streaming group-by (0 = disabled) private int _streamingGroupByFlushThreshold; + /// Opt-in: use the streaming k-way-merge selection ORDER BY combine over sorted segments + private boolean _sortedSelectionMergeEnabled; + /// Output block size (rows) for the streaming selection ORDER BY combine + private int _sortedSelectionMergeBlockSize = Broker.DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE; + // Whether null handling is enabled private boolean _nullHandlingEnabled; // Whether server returns the final result @@ -547,6 +553,22 @@ public void setStreamingGroupByFlushThreshold(int streamingGroupByFlushThreshold _streamingGroupByFlushThreshold = streamingGroupByFlushThreshold; } + public boolean isSortedSelectionMergeEnabled() { + return _sortedSelectionMergeEnabled; + } + + public void setSortedSelectionMergeEnabled(boolean sortedSelectionMergeEnabled) { + _sortedSelectionMergeEnabled = sortedSelectionMergeEnabled; + } + + public int getSortedSelectionMergeBlockSize() { + return _sortedSelectionMergeBlockSize; + } + + public void setSortedSelectionMergeBlockSize(int sortedSelectionMergeBlockSize) { + _sortedSelectionMergeBlockSize = sortedSelectionMergeBlockSize; + } + public boolean isNullHandlingEnabled() { return _nullHandlingEnabled; } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java index 30d0b37da0e5..140e3a887e54 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java @@ -24,6 +24,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Supplier; @@ -41,6 +42,7 @@ import org.apache.pinot.spi.exception.EarlyTerminationException; import org.apache.pinot.spi.exception.QueryErrorCode; import org.apache.pinot.spi.exception.QueryErrorMessage; +import org.apache.pinot.spi.query.QueryThreadContext; import org.apache.pinot.util.TestUtils; import org.testng.Assert; import org.testng.annotations.AfterClass; @@ -54,6 +56,7 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; /** @@ -133,16 +136,7 @@ public void testCancelSelectionOrderByCombineOperator() { @Test public void testCancelMinMaxValueBasedSelectionOrderByCombineOperator() { CountDownLatch ready = new CountDownLatch(1); - List operators = getOperators(ready, () -> { - IndexSegment seg = mock(IndexSegment.class); - DataSource ds = mock(DataSource.class); - DataSourceMetadata dsmd = mock(DataSourceMetadata.class); - when(dsmd.getMinValue()).thenReturn(100L); - when(dsmd.getMaxValue()).thenReturn(200L); - when(seg.getDataSource(anyString(), any())).thenReturn(ds); - when(ds.getDataSourceMetadata()).thenReturn(dsmd); - return seg; - }); + List operators = getOperators(ready, minMaxSegmentSupplier()); QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column"); queryContext.setEndTimeMs(System.currentTimeMillis() + 10000); MinMaxValueBasedSelectionOrderByCombineOperator combineOperator = @@ -172,13 +166,85 @@ public void testCancelGroupByOrderByCombineOperator() { testCancelCombineOperator(combineOperator, ready); } + @Test + public void testCancelStreamingSelectionOrderByCombineOperator() { + CountDownLatch ready = new CountDownLatch(1); + List operators = getOperators(ready, minMaxSegmentSupplier()); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column"); + queryContext.setEndTimeMs(System.currentTimeMillis() + 10000); + // Single-stage mode: nextBlock() drives the whole merge synchronously on the (cancellable) caller thread. + StreamingSelectionOrderByCombineOperator combineOperator = + new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService, false); + testCancelCombineOperator(combineOperator, ready, operators); + } + + @Test + public void testCancelStreamingSelectionOrderByCombineOperatorStreamingMode() { + CountDownLatch ready = new CountDownLatch(1); + List operators = getOperators(ready, minMaxSegmentSupplier()); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column"); + queryContext.setEndTimeMs(System.currentTimeMillis() + 10000); + // Streaming (MSE-leaf) mode: the first getNextBlock() still drives the merge on the caller thread, so the same + // interrupt-on-cancel path applies and must surface an ExceptionResultsBlock. + StreamingSelectionOrderByCombineOperator combineOperator = + new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService, true); + testCancelCombineOperator(combineOperator, ready, operators); + } + + /// The merge loop drains its heap on the caller thread and, for single-block cursors, may never re-enter a child + /// operator, so it must check the query deadline itself. With an already-expired deadline the operator must surface a + /// timeout before activating any child - asserted via {@code _operationInProgress}, which also keeps the test + /// from passing vacuously if the timeout came from somewhere else. + @Test + public void testStreamingSelectionOrderByCombineOperatorHonorsDeadline() { + List operators = getOperators(null, minMaxSegmentSupplier()); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column"); + queryContext.setEndTimeMs(System.currentTimeMillis() - 1); + StreamingSelectionOrderByCombineOperator combineOperator = + new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService, false); + try (QueryThreadContext ignore = QueryThreadContext.openForSseTest()) { + BaseResultsBlock resultsBlock = combineOperator.nextBlock(); + assertTrue(resultsBlock instanceof ExceptionResultsBlock, + "Expired deadline must surface as an ExceptionResultsBlock, got: " + resultsBlock.getClass().getName()); + } + for (Operator operator : operators) { + assertFalse(((SlowOperator) operator)._operationInProgress.get(), + "Deadline must be observed before any child operator is driven"); + } + } + + /// A segment whose first order-by column reports a non-null min/max, as required by the min/max-based operators. + private static Supplier minMaxSegmentSupplier() { + return () -> { + IndexSegment seg = mock(IndexSegment.class); + DataSource ds = mock(DataSource.class); + DataSourceMetadata dsmd = mock(DataSourceMetadata.class); + when(dsmd.getMinValue()).thenReturn(100L); + when(dsmd.getMaxValue()).thenReturn(200L); + when(seg.getDataSource(anyString(), any())).thenReturn(ds); + when(ds.getDataSourceMetadata()).thenReturn(dsmd); + return seg; + }; + } + private void testCancelCombineOperator(BaseCombineOperator combineOperator, CountDownLatch ready) { + testCancelCombineOperator(combineOperator, ready, List.of()); + } + + /// Submits the combine operator on a separate thread, waits for a child operator to start, cancels the future, and + /// asserts the operator surfaces an {@link ExceptionResultsBlock}. When {@code operatorsToVerify} is non-empty + /// (the single-threaded streaming combine, whose merge runs the children on the caller thread), it additionally + /// asserts the cancellation was genuine - a child actually started and none completed normally - so the test + /// cannot pass vacuously on an unrelated failure or by never exercising the cancel path. + private void testCancelCombineOperator(BaseCombineOperator combineOperator, CountDownLatch ready, + List operatorsToVerify) { AtomicReference resultsBlock = new AtomicReference<>(); // Avoid early finalization by not using Executors.newSingleThreadExecutor (java <= 20, JDK-8145304) ExecutorService combineExecutor = Executors.newFixedThreadPool(1); try { Future future = combineExecutor.submit(() -> resultsBlock.set(combineOperator.nextBlock())); - ready.await(); + // Bound the wait so a regression where no child operator ever starts fails fast here instead of hanging. + assertTrue(ready.await(10, TimeUnit.SECONDS), "Expected a child operator to start before cancellation"); // At this point, the combineOperator is or will be waiting on future.get() for all sub operators, and the // waiting can be cancelled as below. future.cancel(true); @@ -189,6 +255,12 @@ private void testCancelCombineOperator(BaseCombineOperator combineOperator, C } TestUtils.waitForCondition((aVoid) -> resultsBlock.get() instanceof ExceptionResultsBlock, 10_000, "Should have been cancelled"); + // Genuine-cancellation check: no child ran to normal completion (each was interrupted or never started), so the + // ExceptionResultsBlock above is the cancellation outcome rather than an unrelated error. + for (Operator operator : operatorsToVerify) { + assertFalse(((SlowOperator) operator)._notInterrupted.get(), + "No operator should have completed normally after cancellation"); + } } /** diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java new file mode 100644 index 000000000000..e2d9c9f8a9be --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java @@ -0,0 +1,553 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.operator.combine; + +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.stream.Collectors; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.request.context.OrderByExpressionContext; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock; +import org.apache.pinot.core.operator.blocks.results.MetadataResultsBlock; +import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock; +import org.apache.pinot.core.plan.CombinePlanNode; +import org.apache.pinot.core.plan.PlanNode; +import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2; +import org.apache.pinot.core.plan.maker.PlanMaker; +import org.apache.pinot.core.query.executor.ResultsBlockStreamer; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.core.query.utils.OrderByComparatorFactory; +import org.apache.pinot.core.util.QueryMultiThreadingUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.CommonConstants.Server; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.intellij.lang.annotations.Language; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/// Combine-level tests for {@link StreamingSelectionOrderByCombineOperator} (step-3 operator) and its wiring into +/// {@link CombinePlanNode#getCombineOperator()} (step-4). +/// +///

The streaming combine must return the same globally-sorted top-K rows as the default +/// {@link MinMaxValueBasedSelectionOrderByCombineOperator}, only (in streaming mode) spread across several bounded +/// blocks. Each functional test therefore asserts streaming-vs-non-streaming parity: it runs the identical query +/// twice over the same in-memory segments - once with {@code sortedSelectionMergeEnabled=true} (asserting the new +/// operator was actually selected) and once with the hint off (asserting the {@code MinMax} operator was selected) - +/// then checks the two row sets are equal as a multiset and that the streaming output is fully sorted by the order-by +/// comparator. +/// +///

To keep the top-K boundary unambiguous (operators may legitimately disagree on which of several rows that tie on +/// every order-by key fall inside the limit) every parity query ends its ORDER BY with the globally-unique +/// {@code valCol} +/// so the comparator is a total order; the merge still genuinely interleaves segments because the primary sort column +/// ({@code sortedCol}) overlaps across segments. Multiset (rather than positional) comparison then tolerates only the +/// harmless reordering of fully-equal projected rows. +public class StreamingSelectionOrderByCombineOperatorTest { + private static final File TEMP_DIR = + new File(FileUtils.getTempDirectory(), "StreamingSelectionOrderByCombineOperatorTest"); + private static final String RAW_TABLE_NAME = "testTable"; + + private static final String SORTED_COL = "sortedCol"; + private static final String TAIL_COL = "tailCol"; + private static final String VAL_COL = "valCol"; + private static final String NULLABLE_COL = "nullableCol"; + /// Non-INT projected columns so the parity assertion can catch a stored-type / boxing regression (e.g. LONG emitted + /// where MinMax emits INT), which an all-INT suite cannot observe. + private static final String LONG_COL = "longCol"; + private static final String STR_COL = "strCol"; + + /// Create (MAX_NUM_THREADS_PER_QUERY * 2) sorted segments so the leaf runs plan nodes across multiple threads. + private static final int NUM_SEGMENTS = QueryMultiThreadingUtils.MAX_NUM_THREADS_PER_QUERY * 2; + private static final int NUM_RECORDS_PER_SEGMENT = 100; + + private static final TableConfig SORTED_TABLE_CONFIG = + new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setSortedColumn(SORTED_COL).build(); + private static final TableConfig UNSORTED_TABLE_CONFIG = + new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build(); + private static final Schema SCHEMA = new Schema.SchemaBuilder() + .addSingleValueDimension(SORTED_COL, FieldSpec.DataType.INT) + .addSingleValueDimension(TAIL_COL, FieldSpec.DataType.INT) + .addSingleValueDimension(VAL_COL, FieldSpec.DataType.INT) + .addSingleValueDimension(NULLABLE_COL, FieldSpec.DataType.INT) + .addSingleValueDimension(LONG_COL, FieldSpec.DataType.LONG) + .addSingleValueDimension(STR_COL, FieldSpec.DataType.STRING) + .build(); + + private static final PlanMaker PLAN_MAKER = new InstancePlanMakerImplV2(); + private static final ExecutorService EXECUTOR = Executors.newCachedThreadPool(); + + /// Sorted segments with overlapping primary-column ranges, so the k-way merge interleaves them (genuine merge rather + /// than concatenation). Built with null handling on so the null-handling test sees real nulls in NULLABLE_COL; reads + /// with null handling off fall back to the column default. + private List _sortedSegments; + /// Sorted segments with disjoint, globally-increasing ranges: ORDER BY sortedCol with a small LIMIT drains only the + /// lowest segment, so min/max pruning must skip the rest (none acquired/scanned). + private List _disjointSegments; + /// A mix of sorted (streaming child) and physically-unsorted (single materialized top-K block child) segments, + /// exercising both SegmentCursor backings in one merge. + private List _mixedSegments; + /// Very low cardinality primary column (4 distinct values across 100 rows) so each value is a long run: exercises the + /// run/heap path in the streaming children and ties on the primary key at the prune boundary. + private List _lowCardSegments; + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteDirectory(TEMP_DIR); + + _sortedSegments = new ArrayList<>(NUM_SEGMENTS); + for (int i = 0; i < NUM_SEGMENTS; i++) { + _sortedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "sorted_" + i, buildOverlappingSortedRecords(i), true)); + } + + _disjointSegments = new ArrayList<>(NUM_SEGMENTS); + for (int i = 0; i < NUM_SEGMENTS; i++) { + _disjointSegments.add(buildSegment(SORTED_TABLE_CONFIG, "disjoint_" + i, buildDisjointSortedRecords(i), false)); + } + + _lowCardSegments = new ArrayList<>(NUM_SEGMENTS); + for (int i = 0; i < NUM_SEGMENTS; i++) { + _lowCardSegments.add(buildSegment(SORTED_TABLE_CONFIG, "lowCard_" + i, buildLowCardinalityRecords(i), false)); + } + + // Two sorted + two unsorted segments, globally-unique valCol across all four so the multiset comparison is exact. + _mixedSegments = new ArrayList<>(4); + _mixedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "mixedSorted_0", buildOverlappingSortedRecords(0), false)); + _mixedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "mixedSorted_1", buildOverlappingSortedRecords(1), false)); + _mixedSegments.add(buildSegment(UNSORTED_TABLE_CONFIG, "mixedUnsorted_0", buildUnsortedRecords(2), false)); + _mixedSegments.add(buildSegment(UNSORTED_TABLE_CONFIG, "mixedUnsorted_1", buildUnsortedRecords(3), false)); + } + + private static List buildOverlappingSortedRecords(int index) { + int baseValue = index * NUM_RECORDS_PER_SEGMENT / 2; + List records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT); + for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) { + GenericRow record = new GenericRow(); + record.putValue(SORTED_COL, baseValue + i); + record.putValue(TAIL_COL, NUM_RECORDS_PER_SEGMENT - i); + // Globally unique across all segments -> a total order when used as the final order-by key. + record.putValue(VAL_COL, index * 1_000_000 + i); + // Beyond the int range so a regression that narrows LONG -> INT would change the boxed value. + record.putValue(LONG_COL, 10_000_000_000L + index * 1_000_000L + i); + record.putValue(STR_COL, "s_" + index + "_" + i); + // Every 7th row is null so the null-handling test exercises null projection. + if (i % 7 == 0) { + record.addNullValueField(NULLABLE_COL); + } else { + record.putValue(NULLABLE_COL, i); + } + records.add(record); + } + return records; + } + + private static List buildDisjointSortedRecords(int index) { + List records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT); + int baseValue = index * 1000; + for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) { + GenericRow record = new GenericRow(); + record.putValue(SORTED_COL, baseValue + i); + record.putValue(TAIL_COL, i); + record.putValue(VAL_COL, baseValue + i); + record.putValue(NULLABLE_COL, i); + record.putValue(LONG_COL, 10_000_000_000L + baseValue + i); + record.putValue(STR_COL, "d_" + index + "_" + i); + records.add(record); + } + return records; + } + + private static List buildLowCardinalityRecords(int index) { + List records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT); + for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) { + GenericRow record = new GenericRow(); + // 4 distinct values per segment, non-decreasing so the segment is physically sorted on SORTED_COL. + record.putValue(SORTED_COL, i / 25); + record.putValue(TAIL_COL, NUM_RECORDS_PER_SEGMENT - i); + record.putValue(VAL_COL, index * 1_000_000 + i); + record.putValue(NULLABLE_COL, i); + record.putValue(LONG_COL, 10_000_000_000L + index * 1_000_000L + i); + record.putValue(STR_COL, "l_" + index + "_" + i); + records.add(record); + } + return records; + } + + private static List buildUnsortedRecords(int index) { + List records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT); + for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) { + GenericRow record = new GenericRow(); + // A non-monotonic permutation of [0, NUM_RECORDS_PER_SEGMENT) (7919 is prime and coprime with 100), so the + // column is genuinely not physically sorted and SelectionPlanNode falls back to a materialized top-K block. + record.putValue(SORTED_COL, (i * 7919) % NUM_RECORDS_PER_SEGMENT); + record.putValue(TAIL_COL, i); + record.putValue(VAL_COL, index * 1_000_000 + i); + record.putValue(NULLABLE_COL, i); + record.putValue(LONG_COL, 10_000_000_000L + index * 1_000_000L + i); + record.putValue(STR_COL, "u_" + index + "_" + i); + records.add(record); + } + return records; + } + + private static IndexSegment buildSegment(TableConfig tableConfig, String segmentName, List records, + boolean nullHandling) + throws Exception { + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, SCHEMA); + segmentGeneratorConfig.setTableName(RAW_TABLE_NAME); + segmentGeneratorConfig.setSegmentName(segmentName); + segmentGeneratorConfig.setDefaultNullHandlingEnabled(nullHandling); + segmentGeneratorConfig.setOutDir(TEMP_DIR.getPath()); + + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records)); + driver.build(); + + return ImmutableSegmentLoader.load(new File(TEMP_DIR, segmentName), ReadMode.mmap); + } + + @Test + public void testAscendingParity() { + assertParity(_sortedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", false); + } + + @Test + public void testDescendingParity() { + // Reverse order must be allowed for the per-segment forward-scan to iterate sortedCol descending; otherwise the + // segment falls back to the materialized DESC operator (covered separately by testDescIncompatibleFallbackParity). + assertParity(_sortedSegments, + "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC " + + "LIMIT 50", false); + } + + @Test + public void testDescIncompatibleFallbackParity() { + // allowReverseOrder=false + DESC -> the streaming child cannot scan descending, so SelectionPlanNode emits the + // materialized DESC top-K block; the combine still routes to the streaming combine and merges single-block cursors. + assertParity(_sortedSegments, + "SET allowReverseOrder=false; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC " + + "LIMIT 50", false); + } + + @Test + public void testLimitOffsetParity() { + // The server retains limit + offset rows; the broker applies the offset later, so both operators keep 40 rows. + assertParity(_sortedSegments, + "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 30 OFFSET 10", false); + } + + @Test + public void testLowCardinalityMultiColumnParity() { + // Low-cardinality primary column => long runs and many sortedCol ties at the prune boundary; valCol breaks ties. + assertParity(_lowCardSegments, + "SELECT sortedCol, tailCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 40", false); + } + + @Test + public void testTwoPhaseSelectNonOrderByParity() { + // tailCol is selected but not an order-by key -> the streaming children take the two-phase (order-by-then-fetch) + // path. valCol is order-by-only, exercising the phase-1 projection of a non-selected order-by column. + assertParity(_sortedSegments, "SELECT tailCol, sortedCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", + false); + } + + @Test + public void testNonIntProjectionParity() { + // Projects a LONG and a STRING column so the multiset comparison would catch a stored-type / boxing regression that + // an all-INT projection cannot observe. + assertParity(_sortedSegments, + "SELECT strCol, longCol, sortedCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", false); + } + + @Test + public void testMixedSortedAndUnsortedSegmentsParity() { + assertParity(_mixedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", false); + } + + @Test + public void testNullHandlingEnabledParity() { + // Null handling on disables min/max pruning (the combine activates every segment); nullableCol carries real nulls. + assertParity(_sortedSegments, + "SELECT nullableCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", true); + } + + @Test + public void testNullHandlingDisabledParity() { + // Same segments/query as the enabled case but null handling off: nulls read back as the column default, pruning on. + assertParity(_sortedSegments, + "SELECT nullableCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", false); + } + + @Test + public void testStreamingMultiBlockExactCount() { + // A precise check on the bounded-flush behavior: 20 rows flushed in blocks of 3 yields ceil(20/3) = 7 data blocks. + int blockSize = 3; + int limit = 20; + @Language("sql") String query = + "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT " + limit; + Result streaming = run(_sortedSegments, query, true, false, true, blockSize); + assertEquals(streaming._combineOperator.getClass(), StreamingSelectionOrderByCombineOperator.class); + assertEquals(streaming._rows.size(), limit); + assertEquals(streaming._numBlocks, (limit + blockSize - 1) / blockSize, "Unexpected number of streamed blocks"); + assertSorted(streaming._rows, orderByComparator(query, false)); + assertMultisetEquals(streaming._rows, run(_sortedSegments, query, false, false, false, 0)._rows); + } + + @Test + public void testPruningSkipsOutOfTopKSegments() { + // Disjoint, globally-increasing ranges + small LIMIT: only the lowest segment can contribute, the rest are pruned + // (never acquired or scanned), so far fewer than all docs are scanned. + Result result = run(_disjointSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5", + true, false, false, 0); + assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator); + assertEquals(result._rows.size(), 5); + for (int i = 0; i < 5; i++) { + assertEquals((int) result._rows.get(i)[0], i, "Unexpected value at position " + i); + } + int totalDocs = NUM_SEGMENTS * NUM_RECORDS_PER_SEGMENT; + assertTrue(result._numDocsScanned < totalDocs, "Pruning should avoid scanning every doc, scanned: " + + result._numDocsScanned + " of " + totalDocs); + assertTrue(result._numDocsScanned <= NUM_RECORDS_PER_SEGMENT, + "Only the lowest-range segment should be scanned, but docs scanned was: " + result._numDocsScanned); + } + + @Test + public void testEmptyResultSchemaFallback() { + // A filter that matches nothing: every streaming child returns no rows, so the combine rebuilds the result schema + // from the segment metadata (order-by expressions first) rather than from a child block. + Result result = run(_sortedSegments, + "SELECT sortedCol, valCol FROM testTable WHERE sortedCol < 0 ORDER BY sortedCol, valCol LIMIT 10", true, false, + false, 0); + assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator); + assertTrue(result._rows.isEmpty(), "Expected an empty result, got: " + result._rows.size() + " rows"); + assertEquals(result._schema, new DataSchema(new String[]{SORTED_COL, VAL_COL}, + new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.INT})); + } + + @Test + public void testHintOffSelectsMinMaxOperator() { + // Default behavior is unchanged when the hint is off: the classic MinMax operator is still selected. + Result result = run(_sortedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", + false, false, false, 0); + assertTrue(result._combineOperator instanceof MinMaxValueBasedSelectionOrderByCombineOperator, + "Hint off must keep the default MinMax combine operator, got: " + + result._combineOperator.getClass().getSimpleName()); + } + + @Test + public void testNonIdentifierOrderByFallsBack() { + // Even with the hint on, a non-identifier first order-by expression falls back to SelectionOrderByCombineOperator + // (the streaming operator and its segment-level counterpart only support a leading identifier). + Result result = run(_sortedSegments, + "SELECT sortedCol, valCol FROM testTable ORDER BY ADD(sortedCol, 1), valCol LIMIT 50", true, false, false, 0); + assertEquals(result._combineOperator.getClass(), SelectionOrderByCombineOperator.class, + "Non-identifier first order-by must fall back to SelectionOrderByCombineOperator, got: " + + result._combineOperator.getClass().getSimpleName()); + } + + /// Asserts streaming-vs-non-streaming parity for {@code query}. Runs the MinMax combine (hint off) as the reference, + /// then runs the streaming combine in BOTH single-block mode and bounded multi-block streaming mode, asserting each + /// selects the streaming operator and produces rows that are sorted by the order-by comparator and equal the MinMax + /// rows as a multiset. The streaming variant additionally checks the bounded-flush invariants. + private void assertParity(List segments, @Language("sql") String query, boolean nullHandling) { + Result baseline = run(segments, query, false, nullHandling, false, 0); + assertEquals(baseline._combineOperator.getClass(), MinMaxValueBasedSelectionOrderByCombineOperator.class, + "Baseline must be the MinMax combine operator, got: " + baseline._combineOperator.getClass().getSimpleName()); + Comparator comparator = orderByComparator(query, nullHandling); + + // Classic single-stage path (null streamer): the streaming combine flushes the whole merge as one block. + Result singleStage = run(segments, query, true, nullHandling, false, 0); + assertStreamingParity(singleStage, baseline, comparator, query, false, 0); + + // MSE leaf path (non-null streamer): a small block size forces several bounded data blocks before the metadata + // block, genuinely exercising the streaming flush path rather than a single trimmed block. + int blockSize = 3; + Result streamed = run(segments, query, true, nullHandling, true, blockSize); + assertStreamingParity(streamed, baseline, comparator, query, true, blockSize); + } + + private void assertStreamingParity(Result result, Result baseline, Comparator comparator, + @Language("sql") String query, boolean streaming, int blockSize) { + assertEquals(result._combineOperator.getClass(), StreamingSelectionOrderByCombineOperator.class, + "Expected the streaming combine operator for query: " + query); + assertEquals(result._schema, baseline._schema, "Schema mismatch for query: " + query); + assertSorted(result._rows, comparator); + assertMultisetEquals(result._rows, baseline._rows); + if (streaming) { + int total = 0; + for (int size : result._blockSizes) { + assertTrue(size > 0 && size <= blockSize, + "Streamed block size out of range (0, " + blockSize + "] for query " + query + ": " + size); + total += size; + } + assertEquals(total, result._rows.size(), "Streamed block sizes must sum to the row count for query: " + query); + if (result._rows.size() > blockSize) { + assertTrue(result._numBlocks >= 2, "Expected multiple streamed blocks for query: " + query); + } + } + } + + /// Runs one combine over {@code segments} and collects its rows, blocks, schema and docs-scanned stat. + private Result run(List segments, @Language("sql") String query, boolean hintOn, boolean nullHandling, + boolean streaming, int blockSize) { + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query); + queryContext.setNullHandlingEnabled(nullHandling); + if (hintOn) { + queryContext.setSortedSelectionMergeEnabled(true); + if (blockSize > 0) { + queryContext.setSortedSelectionMergeBlockSize(blockSize); + } + } + queryContext.setEndTimeMs(System.currentTimeMillis() + Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS); + + List planNodes = new ArrayList<>(segments.size()); + for (IndexSegment segment : segments) { + SegmentContext segmentContext = new SegmentContext(segment); + planNodes.add(streaming ? PLAN_MAKER.makeStreamingSegmentPlanNode(segmentContext, queryContext) + : PLAN_MAKER.makeSegmentPlanNode(segmentContext, queryContext)); + } + ResultsBlockStreamer streamer = streaming ? block -> { + } : null; + CombinePlanNode combinePlanNode = new CombinePlanNode(planNodes, queryContext, EXECUTOR, streamer); + + Result result = new Result(); + Operator combineOperator = combinePlanNode.run(); + result._combineOperator = combineOperator; + result._rows = new ArrayList<>(); + result._blockSizes = new ArrayList<>(); + if (streaming) { + // Drive the streaming combine: collect bounded data blocks until the terminal metadata block, which carries the + // aggregated execution stats. + while (true) { + BaseResultsBlock block = (BaseResultsBlock) combineOperator.nextBlock(); + if (block instanceof MetadataResultsBlock) { + if (result._schema == null) { + result._schema = block.getDataSchema(); + } + result._numDocsScanned = block.getNumDocsScanned(); + break; + } + SelectionResultsBlock dataBlock = (SelectionResultsBlock) block; + if (result._schema == null) { + result._schema = dataBlock.getDataSchema(); + } + List rows = dataBlock.getRows(); + assertNotNull(rows); + result._rows.addAll(rows); + result._blockSizes.add(rows.size()); + result._numBlocks++; + assertTrue(result._numBlocks < 1_000_000, "Streaming combine did not terminate"); + } + } else { + SelectionResultsBlock block = (SelectionResultsBlock) combineOperator.nextBlock(); + result._schema = block.getDataSchema(); + List rows = block.getRows(); + assertNotNull(rows); + result._rows.addAll(rows); + result._blockSizes.add(rows.size()); + result._numBlocks = 1; + result._numDocsScanned = block.getNumDocsScanned(); + } + return result; + } + + private static Comparator orderByComparator(@Language("sql") String query, boolean nullHandling) { + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query); + List orderByExpressions = queryContext.getOrderByExpressions(); + assertNotNull(orderByExpressions); + return OrderByComparatorFactory.getComparator(orderByExpressions, nullHandling); + } + + private static void assertSorted(List rows, Comparator comparator) { + for (int i = 1; i < rows.size(); i++) { + assertTrue(comparator.compare(rows.get(i - 1), rows.get(i)) <= 0, + "Rows not sorted by the order-by comparator at position " + i); + } + } + + /// Asserts the two row lists contain the same rows, independent of the ordering of fully-equal projected rows. + private static void assertMultisetEquals(List actual, List expected) { + assertEquals(toCanonical(actual), toCanonical(expected), "Row multisets differ"); + } + + /// Canonicalizes rows for multiset comparison. Each cell is encoded with its runtime class so a stored-type / boxing + /// regression (e.g. a LONG emitted where the reference emits INT) changes the encoding and fails the assertion, which + /// a plain {@code Arrays.toString} (type-blind) comparison would miss. + private static List toCanonical(List rows) { + return rows.stream().map(row -> { + StringBuilder sb = new StringBuilder("["); + for (int i = 0; i < row.length; i++) { + if (i > 0) { + sb.append(", "); + } + Object cell = row[i]; + sb.append(cell == null ? "null" : cell.getClass().getSimpleName() + ":" + cell); + } + return sb.append(']').toString(); + }).sorted().collect(Collectors.toList()); + } + + @AfterClass + public void tearDown() + throws IOException { + EXECUTOR.shutdownNow(); + for (List segments : List.of(_sortedSegments, _disjointSegments, _mixedSegments, _lowCardSegments)) { + for (IndexSegment segment : segments) { + segment.destroy(); + } + } + FileUtils.deleteDirectory(TEMP_DIR); + } + + /// Captured output of a single combine run. + private static class Result { + private Operator _combineOperator; + private DataSchema _schema; + private List _rows; + private List _blockSizes; + private int _numBlocks; + private long _numDocsScanned; + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java new file mode 100644 index 000000000000..73a071f673a0 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java @@ -0,0 +1,406 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.operator.query; + +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock; +import org.apache.pinot.core.plan.SelectionPlanNode; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.intellij.lang.annotations.Language; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/// Segment-level tests for {@link StreamingSelectionOrderByOperator}. +/// +///

The operator emits the same globally-sorted rows as the existing materialized selection ORDER BY operators, only +/// spread across many lazily-produced blocks. Each test therefore asserts stream-vs-materialized parity: it +/// drives the streaming operator through {@link SelectionPlanNode} with {@code sortedSelectionMergeEnabled=true}, +/// concatenates +/// every {@link Operator#nextBlock()} output until {@code null}, and asserts the concatenation equals the single block +/// the materialized operator ({@link SelectionPartiallyOrderedByLinearOperator} / {@link SelectionOrderByOperator}) +/// produces for the identical query with the hint off. +/// +///

To keep element-wise comparison deterministic (priority-queue draining is not stable for rows that tie on every +/// order-by column) each fixture makes the order-by column tuple unique per row: the {@code _segment} fixture has a +/// unique sorted column, while the {@code _dupSegment} / {@code _largeSegment} fixtures repeat the sorted column but +/// pair it with a unique {@code TAIL_COL} order-by tail. The all-order-by-columns-tie case (where stream and +/// materialized output may legitimately differ in row order) is intentionally out of scope here and is covered by the +/// combine-level test via multiset comparison. +public class StreamingSelectionOrderByOperatorTest { + private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "StreamingSelectionOrderByOperatorTest"); + private static final String RAW_TABLE_NAME = "testTable"; + + private static final String SORTED_COL = "sortedCol"; + private static final String TAIL_COL = "tailCol"; + private static final String VAL_COL = "valCol"; + private static final String NULLABLE_COL = "nullableCol"; + + private static final TableConfig TABLE_CONFIG = + new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setSortedColumn(SORTED_COL).build(); + private static final Schema SCHEMA = new Schema.SchemaBuilder() + .addSingleValueDimension(SORTED_COL, FieldSpec.DataType.INT) + .addSingleValueDimension(TAIL_COL, FieldSpec.DataType.INT) + .addSingleValueDimension(VAL_COL, FieldSpec.DataType.INT) + .addSingleValueDimension(NULLABLE_COL, FieldSpec.DataType.INT) + .build(); + + /// Unique sorted column, no nulls. Exercises the no-tail emission mode. + private static final int NUM_RECORDS = 30; + /// Repeated sorted column (RUN_SIZE rows per value) paired with a unique tail. Exercises the run / heap path, and + /// carries nulls in NULLABLE_COL for the null-handling cases. + private static final int NUM_DISTINCT_SORTED = 10; + private static final int RUN_SIZE = 4; + private static final int NUM_DUP_RECORDS = NUM_DISTINCT_SORTED * RUN_SIZE; + /// A single primary-value run larger than one project block (DocIdSetPlanNode.MAX_DOC_PER_CALL == 10_000), so the + /// forward scan crosses a block boundary within one run. + private static final int NUM_LARGE_RECORDS = 12_000; + /// Two runs where the first is larger than one project block: with a limit greater than the first run's size the + /// operator both crosses a project-block boundary mid-run AND carries _pendingRow across a run boundary to emit a + /// second block. + private static final int MULTI_RUN0_SIZE = 10_500; + private static final int MULTI_RUN1_SIZE = 1_000; + private static final int NUM_MULTI_RUN_RECORDS = MULTI_RUN0_SIZE + MULTI_RUN1_SIZE; + + private IndexSegment _segment; + private IndexSegment _dupSegment; + private IndexSegment _largeSegment; + private IndexSegment _multiRunSegment; + /// Carries nulls in an order-by (tail) column so the streaming operator's own null path (phase-1 materialization and + /// the null-aware comparator) is exercised, not just the shared two-phase fetch. + private IndexSegment _nullTailSegment; + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteDirectory(TEMP_DIR); + _segment = buildSegment("uniqueSorted", buildUniqueSortedRecords(), false); + _dupSegment = buildSegment("dupSorted", buildDupSortedRecords(), true); + _largeSegment = buildSegment("largeRun", buildLargeRunRecords(), false); + _multiRunSegment = buildSegment("multiRun", buildMultiRunRecords(), false); + _nullTailSegment = buildSegment("nullTail", buildNullTailRecords(), true); + } + + private static List buildUniqueSortedRecords() { + List records = new ArrayList<>(NUM_RECORDS); + for (int i = 0; i < NUM_RECORDS; i++) { + GenericRow record = new GenericRow(); + record.putValue(SORTED_COL, i); + record.putValue(TAIL_COL, NUM_RECORDS - i); + record.putValue(VAL_COL, i * 3); + record.putValue(NULLABLE_COL, i); + records.add(record); + } + return records; + } + + private static List buildDupSortedRecords() { + List records = new ArrayList<>(NUM_DUP_RECORDS); + for (int i = 0; i < NUM_DUP_RECORDS; i++) { + GenericRow record = new GenericRow(); + record.putValue(SORTED_COL, i / RUN_SIZE); + // Tail resets to a descending sequence within each run so the column is NOT globally sorted (the run/heap path is + // only taken when the tail is unsorted); the (sortedCol, tailCol) tuple is still unique per row. + record.putValue(TAIL_COL, RUN_SIZE - 1 - (i % RUN_SIZE)); + record.putValue(VAL_COL, i * 2); + // Every third row carries a null so the two-phase fetch is exercised with and without null handling. + if (i % 3 == 0) { + record.addNullValueField(NULLABLE_COL); + } else { + record.putValue(NULLABLE_COL, i); + } + records.add(record); + } + return records; + } + + private static List buildLargeRunRecords() { + List records = new ArrayList<>(NUM_LARGE_RECORDS); + for (int i = 0; i < NUM_LARGE_RECORDS; i++) { + GenericRow record = new GenericRow(); + // All rows share one sorted value, so they form a single run spanning multiple project blocks. + record.putValue(SORTED_COL, 0); + // A non-monotonic permutation of [0, NUM_LARGE_RECORDS) (7919 is prime and coprime with NUM_LARGE_RECORDS, so the + // mapping is a bijection): unique tail values that are not physically sorted, forcing the run/heap path. + record.putValue(TAIL_COL, (i * 7919) % NUM_LARGE_RECORDS); + record.putValue(VAL_COL, i); + record.putValue(NULLABLE_COL, i); + records.add(record); + } + return records; + } + + private static List buildMultiRunRecords() { + List records = new ArrayList<>(NUM_MULTI_RUN_RECORDS); + appendRun(records, 0, MULTI_RUN0_SIZE); + appendRun(records, 1, MULTI_RUN1_SIZE); + return records; + } + + /// Appends one run of {@code runSize} rows all sharing {@code sortedValue}, with a tail that descends within the run + /// (so the tail column is not globally sorted) and is unique per (sortedCol, tailCol) tuple. + private static void appendRun(List records, int sortedValue, int runSize) { + for (int j = 0; j < runSize; j++) { + GenericRow record = new GenericRow(); + record.putValue(SORTED_COL, sortedValue); + record.putValue(TAIL_COL, runSize - 1 - j); + record.putValue(VAL_COL, sortedValue * 1_000_000 + j); + record.putValue(NULLABLE_COL, j); + records.add(record); + } + } + + private static List buildNullTailRecords() { + List records = new ArrayList<>(); + // Three runs, each with exactly one null tail and two distinct non-null tails, so ordering stays deterministic + // (no all-order-by-columns tie) while a null flows through the order-by tail column. + for (int g = 0; g < 3; g++) { + for (int j = 0; j < 3; j++) { + GenericRow record = new GenericRow(); + record.putValue(SORTED_COL, g); + if (j == 0) { + record.addNullValueField(TAIL_COL); + } else { + record.putValue(TAIL_COL, g * 10 + j); + } + record.putValue(VAL_COL, g * 100 + j); + record.putValue(NULLABLE_COL, g * 100 + j); + records.add(record); + } + } + return records; + } + + private static IndexSegment buildSegment(String segmentName, List records, boolean nullHandling) + throws Exception { + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA); + segmentGeneratorConfig.setTableName(RAW_TABLE_NAME); + segmentGeneratorConfig.setSegmentName(segmentName); + segmentGeneratorConfig.setDefaultNullHandlingEnabled(nullHandling); + segmentGeneratorConfig.setOutDir(TEMP_DIR.getPath()); + + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records)); + driver.build(); + + return ImmutableSegmentLoader.load(new File(TEMP_DIR, segmentName), ReadMode.mmap); + } + + @Test + public void testSingleSortedColumnAscending() { + assertParity(_segment, "SELECT sortedCol FROM testTable ORDER BY sortedCol", false, 1); + } + + @Test + public void testSingleSortedColumnDescending() { + // Reverse order must be allowed for the forward-scan project to iterate the sorted column descending; otherwise + // SelectionPlanNode falls back to the materialized DESC operator and the streaming operator is never built. + assertParity(_segment, "SET allowReverseOrder=true; SELECT sortedCol FROM testTable ORDER BY sortedCol DESC", false, + 1); + } + + @Test + public void testSortedPrefixWithUnsortedTail() { + // Repeated sorted value + unique tail exercises nextRun(): per-run top-K heap and the one-row lookahead. With the + // default LIMIT 10 spanning ~3 runs of 4, the streaming operator emits more than one block. + assertParity(_dupSegment, "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol", false, 2); + } + + @Test + public void testTwoPhaseSingleSortedColumn() { + // Output has a non-order-by column (valCol) -> two-phase fetch. Unique sorted column keeps ordering deterministic. + assertParity(_segment, "SELECT valCol, sortedCol FROM testTable ORDER BY sortedCol", false, 1); + } + + @Test + public void testTwoPhaseWithUnsortedTail() { + // Two-phase fetch combined with the run/heap path (sorted prefix + unsorted tail). + assertParity(_dupSegment, "SELECT valCol FROM testTable ORDER BY sortedCol, tailCol", false, 1); + } + + @Test + public void testLimitOffsetNoTail() { + // Server retains limit + offset rows; the broker applies the offset later. + assertParity(_segment, "SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 5 OFFSET 3", false, 1); + } + + @Test + public void testLimitOffsetWithTail() { + assertParity(_dupSegment, "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 7 OFFSET 5", + false, 1); + } + + @Test + public void testRunSpanningMultipleProjectBlocks() { + // One run of 12_000 rows (> one 10k project block); the heap caps at limit + offset while the scan crosses the + // block boundary. A single primary value means a single run, hence a single emitted block. + StreamingSelectionOrderByOperator operator = assertParity(_largeSegment, + "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 25", false, 1); + // The whole run is scanned to find its top-K, which only happens if the forward scan pulled every project block + // (proving the scan genuinely crossed the 10k boundary rather than stopping at the first block). + assertEquals(operator.getExecutionStatistics().getNumDocsScanned(), NUM_LARGE_RECORDS); + } + + @Test + public void testRunBoundaryAcrossProjectBlocks() { + // First run (10_500 rows) is larger than one project block, and the limit (10_600) exceeds it, so the operator + // crosses a project-block boundary mid-run AND carries _pendingRow across the run boundary to emit a second block. + StreamingSelectionOrderByOperator operator = assertParity(_multiRunSegment, + "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 10600", false, 2); + assertEquals(operator.getExecutionStatistics().getNumDocsScanned(), NUM_MULTI_RUN_RECORDS); + } + + @Test + public void testNullHandlingEnabled() { + assertParity(_dupSegment, "SELECT nullableCol, sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol", true, + 1); + } + + @Test + public void testNullHandlingDisabled() { + assertParity(_dupSegment, "SELECT nullableCol, sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol", + false, 1); + } + + @Test + public void testNullInOrderByColumnWithNullHandling() { + // Drives a null through the order-by tail column (not just the carried non-order-by column), exercising the + // streaming operator's null-aware comparator and phase-1 null materialization. + String query = "SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol"; + assertParity(_nullTailSegment, query, true, 1); + // Absolute anchor: with null handling on, the order-by tail column (index 1 after extractExpressions) actually + // carries nulls through to the output. Guards against both operators substituting a default value identically. + List rows = collectStreamingRows(_nullTailSegment, query, true); + assertTrue(rows.stream().anyMatch(row -> row[1] == null), "Expected a null in the order-by tail column"); + } + + @Test + public void testNullInOrderByColumnWithoutNullHandling() { + String query = "SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol"; + assertParity(_nullTailSegment, query, false, 1); + // With null handling off the null reads back as the column's default value, so no output cell is null. + List rows = collectStreamingRows(_nullTailSegment, query, false); + assertTrue(rows.stream().noneMatch(row -> row[1] == null), + "Expected no nulls in the output when null handling is disabled"); + } + + /// Runs {@code query} twice over {@code segment} - once with the streaming hint on, once off - and asserts the + /// concatenated streaming blocks equal the materialized operator's single block, cell by cell. Returns the (now + /// exhausted) streaming operator so callers can make extra assertions on its execution statistics. + /// + /// @param expectedMinBlocks the minimum number of non-null blocks the streaming operator must emit. Tail-mode + /// cases that span multiple runs pass {@code >= 2} to prove the output is genuinely streamed; cases whose + /// result fits in a single trimmed block pass {@code 1}. + private StreamingSelectionOrderByOperator assertParity(IndexSegment segment, @Language("sql") String query, + boolean nullHandling, int expectedMinBlocks) { + // Streaming path. + QueryContext streamingContext = QueryContextConverterUtils.getQueryContext(query); + streamingContext.setNullHandlingEnabled(nullHandling); + streamingContext.setSortedSelectionMergeEnabled(true); + Operator streamingOperator = + new SelectionPlanNode(new SegmentContext(segment), streamingContext).run(); + assertTrue(streamingOperator instanceof StreamingSelectionOrderByOperator, + "Expected the streaming operator to be built, got: " + streamingOperator.getClass().getSimpleName()); + + List streamingRows = new ArrayList<>(); + DataSchema streamingSchema = null; + int numBlocks = 0; + SelectionResultsBlock block; + while ((block = streamingOperator.nextBlock()) != null) { + numBlocks++; + if (streamingSchema == null) { + streamingSchema = block.getDataSchema(); + } + streamingRows.addAll(block.getRows()); + } + assertTrue(numBlocks >= expectedMinBlocks, + "Expected at least " + expectedMinBlocks + " streaming block(s), got: " + numBlocks); + + // Materialized path (hint off). + QueryContext materializedContext = QueryContextConverterUtils.getQueryContext(query); + materializedContext.setNullHandlingEnabled(nullHandling); + Operator materializedOperator = + new SelectionPlanNode(new SegmentContext(segment), materializedContext).run(); + assertFalse(materializedOperator instanceof StreamingSelectionOrderByOperator, + "Materialized baseline must not be the streaming operator"); + SelectionResultsBlock materializedBlock = materializedOperator.nextBlock(); + assertNotNull(materializedBlock); + List expectedRows = materializedBlock.getRows(); + + assertEquals(streamingSchema, materializedBlock.getDataSchema(), "Schema mismatch for query: " + query); + assertEquals(streamingRows.size(), expectedRows.size(), "Row count mismatch for query: " + query); + for (int i = 0; i < expectedRows.size(); i++) { + assertEquals(streamingRows.get(i), expectedRows.get(i), "Row " + i + " mismatch for query: " + query); + } + return (StreamingSelectionOrderByOperator) streamingOperator; + } + + /// Drains the streaming operator for {@code query} and returns all rows concatenated in emission order. + private List collectStreamingRows(IndexSegment segment, @Language("sql") String query, + boolean nullHandling) { + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query); + queryContext.setNullHandlingEnabled(nullHandling); + queryContext.setSortedSelectionMergeEnabled(true); + Operator operator = new SelectionPlanNode(new SegmentContext(segment), queryContext).run(); + List rows = new ArrayList<>(); + SelectionResultsBlock block; + while ((block = operator.nextBlock()) != null) { + rows.addAll(block.getRows()); + } + return rows; + } + + @AfterClass + public void tearDown() + throws IOException { + for (IndexSegment segment : new IndexSegment[]{_segment, _dupSegment, _largeSegment, _multiRunSegment, + _nullTailSegment}) { + if (segment != null) { + segment.destroy(); + } + } + FileUtils.deleteDirectory(TEMP_DIR); + } +} diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/StreamingSortedMailboxReceiveTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/StreamingSortedMailboxReceiveTest.java new file mode 100644 index 000000000000..0426edef50ea --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/StreamingSortedMailboxReceiveTest.java @@ -0,0 +1,298 @@ +/** + * 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.integration.tests.custom; + +import com.fasterxml.jackson.databind.JsonNode; +import java.io.File; +import java.util.ArrayList; +import java.util.List; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.CommonConstants; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/** + * Integration parity test for the streaming k-way merge in {@code SortedMailboxReceiveOperator}. + * + *

Verifies that a {@code SELECT ... ORDER BY LIMIT } query returns identical row sets AND identical order + * with the k-way merge on and off. The merge needs both {@code sortedSelectionMergeEnabled} (planner marks the receive + * node {@code sortedOnSender}) and {@code streamingSortedMailboxReceive} (runtime activates the merge given that + * marking), so each parity test compares an all-off arm against an all-on arm; see {@link #mergeOptions}. + * + *

Also asserts that the streaming sorted leaf — the precondition behind the {@code sortedOnSender} marking — is + * opt-in, appearing in {@code EXPLAIN PLAN FOR} output only when the option is set. + * + *

The base cluster starts two servers and the data is split across two segments, so the receive node merges multiple + * pre-sorted sender streams. + */ +@Test(suiteName = "CustomClusterIntegrationTest") +public class StreamingSortedMailboxReceiveTest extends CustomDataQueryClusterIntegrationTest { + private static final String DEFAULT_TABLE_NAME = "StreamingSortedMailboxReceiveTest"; + private static final int NUM_TOTAL_DOCS = 1000; + private static final String KEY_INT = "keyInt"; + private static final String KEY_STR = "keyStr"; + private static final String PAYLOAD = "payload"; + + @Override + public String getTableName() { + return DEFAULT_TABLE_NAME; + } + + @Override + public Schema createSchema() { + return new Schema.SchemaBuilder().setSchemaName(getTableName()) + .addSingleValueDimension(KEY_INT, FieldSpec.DataType.INT) + .addSingleValueDimension(KEY_STR, FieldSpec.DataType.STRING) + .addSingleValueDimension(PAYLOAD, FieldSpec.DataType.LONG) + .build(); + } + + @Override + protected long getCountStarResult() { + return NUM_TOTAL_DOCS; + } + + @Override + public List createAvroFiles() + throws Exception { + org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false); + avroSchema.setFields(List.of( + new org.apache.avro.Schema.Field(KEY_INT, org.apache.avro.Schema.create(org.apache.avro.Schema.Type.INT), + null, null), + new org.apache.avro.Schema.Field(KEY_STR, org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING), + null, null), + new org.apache.avro.Schema.Field(PAYLOAD, org.apache.avro.Schema.create(org.apache.avro.Schema.Type.LONG), + null, null))); + + try (AvroFilesAndWriters avroFilesAndWriters = createAvroFilesAndWriters(avroSchema)) { + List> writers = avroFilesAndWriters.getWriters(); + for (int i = 0; i < NUM_TOTAL_DOCS; i++) { + GenericData.Record record = new GenericData.Record(avroSchema); + // Unique keys so ORDER BY produces a single deterministic ordering (no ties to differ across paths). + record.put(KEY_INT, i); + record.put(KEY_STR, String.format("key-%05d", i)); + record.put(PAYLOAD, (long) i * 7); + // Round-robin across the avro files (segments) so each segment holds an interleaved key range. + writers.get(i % getNumAvroFiles()).append(record); + } + return avroFilesAndWriters.getAvroFiles(); + } + } + + @Test + public void testOrderByLimitParityWithAndWithoutStreamingMerge() + throws Exception { + setUseMultiStageQueryEngine(true); + // LIMIT below the total doc count so a leaf ORDER BY LIMIT Sort is pushed to the senders. + String baseQuery = String.format("SELECT %s, %s, %s FROM %s ORDER BY %s LIMIT 50", KEY_INT, KEY_STR, PAYLOAD, + getTableName(), KEY_INT); + + JsonNode baselineResponse = runAndGetResponse(mergeOptions(false) + baseQuery); + JsonNode mergeResponse = runAndGetResponse(mergeOptions(true) + baseQuery); + JsonNode baselineRows = baselineResponse.get("resultTable").get("rows"); + JsonNode mergeRows = mergeResponse.get("resultTable").get("rows"); + assertRowsIdenticalInOrder(baselineRows, mergeRows); + // The parity assertion above holds trivially if both arms took the accumulate path, so pin the paths apart. + assertKWayMergeUsed(baselineResponse, false); + assertKWayMergeUsed(mergeResponse, true); + assertEquals(mergeRows.size(), 50, "LIMIT must be honored"); + // Sanity: results are actually globally sorted ascending by keyInt. + for (int i = 1; i < mergeRows.size(); i++) { + assertTrue(mergeRows.get(i).get(0).asInt() >= mergeRows.get(i - 1).get(0).asInt(), + "merge output must be globally sorted ascending by " + KEY_INT); + } + } + + @Test + public void testOrderByLimitParityDescending() + throws Exception { + setUseMultiStageQueryEngine(true); + String baseQuery = String.format("SELECT %s, %s FROM %s ORDER BY %s DESC LIMIT 25", KEY_INT, PAYLOAD, + getTableName(), KEY_INT); + + JsonNode baselineRows = runAndGetRows(mergeOptions(false) + baseQuery); + JsonNode mergeResponse = runAndGetResponse(mergeOptions(true) + baseQuery); + JsonNode mergeRows = mergeResponse.get("resultTable").get("rows"); + assertRowsIdenticalInOrder(baselineRows, mergeRows); + assertKWayMergeUsed(mergeResponse, true); + assertEquals(mergeRows.size(), 25, "LIMIT must be honored"); + // Sanity: results are actually globally sorted descending by keyInt (DESC is where an inverted comparator or wrong + // null direction in the merge would surface, since both paths could otherwise share the same bug undetected). + for (int i = 1; i < mergeRows.size(); i++) { + assertTrue(mergeRows.get(i).get(0).asInt() <= mergeRows.get(i - 1).get(0).asInt(), + "merge output must be globally sorted descending by " + KEY_INT); + } + } + + @Test + public void testPlannerAutoActivationParity() + throws Exception { + setUseMultiStageQueryEngine(true); + String baseQuery = String.format("SELECT %s, %s, %s FROM %s ORDER BY %s LIMIT 50", KEY_INT, KEY_STR, PAYLOAD, + getTableName(), KEY_INT); + + // Baseline: both options off, so the receive accumulates and sorts. + JsonNode baselineRows = runAndGetRows(mergeOptions(false) + baseQuery); + // Merge arm: sortedSelectionMergeEnabled marks the receive node sortedOnSender AND streamingSortedMailboxReceive + // activates the k-way merge given that marking. + JsonNode autoResponse = runAndGetResponse(mergeOptions(true) + baseQuery); + assertRowsIdenticalInOrder(baselineRows, autoResponse.get("resultTable").get("rows")); + // Anchors: the merge arm must actually have run the merge, not silently fallen back. The runtime evidence is the + // kWayMergeUsed stat; the planner-side precondition is the leaf running the streaming sorted selection combine. + assertKWayMergeUsed(autoResponse, true); + assertStreamingSortedLeafPlanned(baseQuery); + } + + @Test + public void testExplainShowsStreamingSortedLeafWithStep1Hint() + throws Exception { + setUseMultiStageQueryEngine(true); + // The planner sets MailboxReceiveNode.sortedOnSender during fragmentation when the step-1 hint is on and the + // sender fragment is a leaf selection ORDER BY. That internal flag is NOT surfaced as text by any EXPLAIN mode + // today (the asking-servers explain renders only the leaf stage via PlanNodeToRelConverter; the intermediate + // exchange stage stays a logical PinotLogicalExchange), so we cannot assert the flag string directly here. + // Its runtime effect is proven by testPlannerAutoActivationParity. What the asking-servers explain DOES show + // stably is the precondition the flag encodes: with the step-1 hint on, the leaf runs the streaming sorted + // selection combine (SelectOrderbyStreaming) under the exchange, i.e. each sender stream is globally sorted. + String query = String.format( + "SET %s=true; SET %s=true; EXPLAIN PLAN FOR SELECT %s, %s FROM %s ORDER BY %s LIMIT 50", + CommonConstants.Broker.Request.QueryOptionKey.SORTED_SELECTION_MERGE_ENABLED, + CommonConstants.Broker.Request.QueryOptionKey.EXPLAIN_ASKING_SERVERS, KEY_INT, PAYLOAD, getTableName(), + KEY_INT); + JsonNode plan = postQuery(query); + assertEquals(plan.get("exceptions").size(), 0, "EXPLAIN produced exceptions: " + plan.get("exceptions")); + String planText = plan.toString(); + assertTrue(planText.contains("SelectOrderbyStreaming"), + "Step-1 hint should activate the streaming sorted leaf selection. Plan: " + plan); + assertTrue(planText.contains("PinotLogicalExchange"), + "Plan should retain the exchange feeding the sorted receiver. Plan: " + plan); + // Negative control: without the option the leaf must NOT plan the streaming sorted selection, so the assertion + // above cannot pass vacuously on a plan that always contains that string. + JsonNode plainPlan = postQuery(String.format( + "SET %s=true; EXPLAIN PLAN FOR SELECT %s, %s FROM %s ORDER BY %s LIMIT 50", + CommonConstants.Broker.Request.QueryOptionKey.EXPLAIN_ASKING_SERVERS, KEY_INT, PAYLOAD, getTableName(), + KEY_INT)); + assertEquals(plainPlan.get("exceptions").size(), 0, + "EXPLAIN produced exceptions: " + plainPlan.get("exceptions")); + assertFalse(plainPlan.toString().contains("SelectOrderbyStreaming"), + "Streaming sorted leaf must be opt-in. Plan: " + plainPlan); + } + + /** + * Query-option prefix for one arm of a parity test. For the plain leaf-selection shape these tests use, the k-way + * merge needs BOTH options: {@code sortedSelectionMergeEnabled} is what lets the planner mark the receive node + * {@code sortedOnSender} (there is no rel-level sort exchange here to declare it), and {@code + * streamingSortedMailboxReceive} is what turns the merge on given that marking. Setting only the latter is a silent + * no-op, so the "merge" arm must set both or the test compares the accumulate path against itself. + */ + private static String mergeOptions(boolean enabled) { + return "SET " + CommonConstants.Broker.Request.QueryOptionKey.SORTED_SELECTION_MERGE_ENABLED + "=" + enabled + "; " + + "SET " + CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE + "=" + enabled + + "; "; + } + + private JsonNode runAndGetRows(String query) + throws Exception { + return runAndGetResponse(query).get("resultTable").get("rows"); + } + + private JsonNode runAndGetResponse(String query) + throws Exception { + JsonNode response = postQuery(query); + assertEquals(response.get("exceptions").size(), 0, "Query produced exceptions: " + response.get("exceptions")); + return response; + } + + /** + * Asserts whether any MAILBOX_RECEIVE node in the response {@code stageStats} tree reports {@code kWayMergeUsed}. + * + *

This is the runtime half of the anchor (see {@link #assertStreamingSortedLeafPlanned} for the planner half), + * and the only direct evidence that the k-way merge actually ran rather than silently falling back to + * accumulate-then-sort. Reporting is presence-based: the stat is rendered only when the merge was used, so its + * absence across every receive node means every receive took the accumulate path. + */ + private void assertKWayMergeUsed(JsonNode response, boolean expected) { + JsonNode stageStats = response.get("stageStats"); + assertNotNull(stageStats, "Response should carry stageStats: " + response); + assertEquals(findKWayMergeUsed(stageStats), expected, + "Unexpected kWayMergeUsed in stageStats: " + stageStats); + } + + private static boolean findKWayMergeUsed(JsonNode node) { + if (node == null || !node.isObject()) { + return false; + } + JsonNode used = node.get("kWayMergeUsed"); + if (used != null && used.asBoolean()) { + return true; + } + JsonNode children = node.get("children"); + if (children != null) { + for (JsonNode child : children) { + if (findKWayMergeUsed(child)) { + return true; + } + } + } + return false; + } + + /** + * Asserts that, under the merge options, the leaf stage really does run the streaming sorted selection combine. + * + *

That is the precondition the planner requires before it marks the receive node {@code sortedOnSender}, which in + * turn is what activates the k-way merge. The {@code sortedOnSender} flag itself is not rendered by any EXPLAIN mode + * (the asking-servers explain renders only the leaf stage; the intermediate exchange stage stays a logical + * {@code PinotLogicalExchange}), so this covers the planner half; {@link #assertKWayMergeUsed} covers the runtime + * half by reading the {@code kWayMergeUsed} stat out of the response {@code stageStats}. Without both anchors a + * parity test passes just as happily when both arms silently take the accumulate path. + */ + private void assertStreamingSortedLeafPlanned(String baseQuery) + throws Exception { + JsonNode plan = postQuery(mergeOptions(true) + "SET " + + CommonConstants.Broker.Request.QueryOptionKey.EXPLAIN_ASKING_SERVERS + "=true; EXPLAIN PLAN FOR " + + baseQuery); + assertEquals(plan.get("exceptions").size(), 0, "EXPLAIN produced exceptions: " + plan.get("exceptions")); + String planText = plan.toString(); + assertTrue(planText.contains("SelectOrderbyStreaming"), + "Merge arm did not plan the streaming sorted leaf, so the k-way merge could not have activated. Plan: " + plan); + assertTrue(planText.contains("PinotLogicalExchange"), + "Plan should retain the exchange feeding the sorted receiver. Plan: " + plan); + } + + private static void assertRowsIdenticalInOrder(JsonNode expected, JsonNode actual) { + assertEquals(actual.size(), expected.size(), "row count mismatch"); + List expectedRows = new ArrayList<>(); + List actualRows = new ArrayList<>(); + for (int i = 0; i < expected.size(); i++) { + expectedRows.add(expected.get(i).toString()); + actualRows.add(actual.get(i).toString()); + } + // Order-sensitive comparison: row i must match across both result sets. + assertEquals(actualRows, expectedRows, "row sets / order differ between streaming-merge and baseline"); + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java index 193581dcdbb0..701cadd45cb8 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java @@ -535,7 +535,8 @@ private DispatchableSubPlan toDispatchableSubPlan(RelRoot relRoot, PlannerContex return pinotDispatchPlanner.createDispatchableSubPlanV2(plan.getLeft(), plan.getRight()); } SubPlan plan = PinotLogicalQueryPlanner.makePlan(relRoot, tracker, useSpools(plannerContext.getOptions()), - _envConfig.defaultHashFunction(), pruneUnnestColumns(plannerContext.getOptions())); + _envConfig.defaultHashFunction(), pruneUnnestColumns(plannerContext.getOptions()), + QueryOptionsUtils.isSortedSelectionMergeEnabled(plannerContext.getOptions()), _envConfig.getTableCache()); PinotDispatchPlanner pinotDispatchPlanner = new PinotDispatchPlanner(plannerContext, _envConfig.getWorkerManager(), _envConfig.getRequestId(), _envConfig.getTableCache()); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java index eceb18102d97..7a02ba36af2b 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java @@ -34,6 +34,7 @@ import org.apache.calcite.rel.RelRoot; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType; +import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.query.context.PhysicalPlannerContext; import org.apache.pinot.query.planner.PlanFragment; import org.apache.pinot.query.planner.SubPlan; @@ -61,11 +62,13 @@ private PinotLogicalQueryPlanner() { */ public static SubPlan makePlan(RelRoot relRoot, @Nullable TransformationTracker.Builder tracker, boolean useSpools, - String hashFunction, boolean pruneUnnestColumns) { + String hashFunction, boolean pruneUnnestColumns, boolean sortedSelectionMergeEnabled, + @Nullable TableCache tableCache) { PlanNode rootNode = new RelToPlanNodeConverter(tracker, hashFunction, !CommonConstants.Helix.DEFAULT_ENABLE_CASE_INSENSITIVE, pruneUnnestColumns).toPlanNode(relRoot.rel); - PlanFragment rootFragment = planNodeToPlanFragment(rootNode, tracker, useSpools, hashFunction); + PlanFragment rootFragment = + planNodeToPlanFragment(rootNode, tracker, useSpools, hashFunction, sortedSelectionMergeEnabled, tableCache); return new SubPlan(rootFragment, new SubPlanMetadata(RelToPlanNodeConverter.getTableNamesFromRelRoot(relRoot.rel), relRoot.fields), List.of()); @@ -111,8 +114,8 @@ public static Pair makePlanV2( private static PlanFragment planNodeToPlanFragment( PlanNode node, @Nullable TransformationTracker.Builder tracker, boolean useSpools, - String hashFunction) { - PlanFragmenter fragmenter = new PlanFragmenter(); + String hashFunction, boolean sortedSelectionMergeEnabled, @Nullable TableCache tableCache) { + PlanFragmenter fragmenter = new PlanFragmenter(sortedSelectionMergeEnabled, tableCache); PlanFragmenter.Context fragmenterContext = fragmenter.createContext(); node = node.visit(fragmenter, fragmenterContext); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java index b384ef275e0f..a0b33eb61a26 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java @@ -24,8 +24,11 @@ import java.util.ArrayList; import java.util.IdentityHashMap; 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.config.provider.TableCache; import org.apache.pinot.query.planner.PlanFragment; import org.apache.pinot.query.planner.SubPlan; import org.apache.pinot.query.planner.plannode.AggregateNode; @@ -45,6 +48,8 @@ import org.apache.pinot.query.planner.plannode.UnnestNode; import org.apache.pinot.query.planner.plannode.ValueNode; import org.apache.pinot.query.planner.plannode.WindowNode; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.utils.builder.TableNameBuilder; /** @@ -70,6 +75,24 @@ public class PlanFragmenter implements PlanNodeVisitorleaf selection ORDER BY over a single + * physical table, i.e. a {@link SortNode} whose single-input chain down to the leaf consists solely of + * {@link ProjectNode} and {@link TableScanNode} nodes, bottoms out at a {@link TableScanNode}, and that scan resolves + * to exactly one physical table (see {@link #resolvesToSinglePhysicalTable}). + * + *

Any branching node (input count != 1 that is not the leaf scan) or any node that breaks the single-table leaf + * shape (Join, Aggregate, MailboxReceive/Exchange, Window, SetOp, etc.) makes this return {@code false}. This is the + * shape for which the k-way merge in {@code SortedMailboxReceiveOperator} can be safely auto-activated. + */ + private boolean isLeafSelectionOrderBy(PlanNode root) { + if (!(root instanceof SortNode)) { + return false; + } + PlanNode current = root; + while (true) { + if (current instanceof TableScanNode) { + return resolvesToSinglePhysicalTable(((TableScanNode) current).getTableName()); + } + // Only SortNode (root), ProjectNode and TableScanNode are allowed in the chain. + if (!(current instanceof SortNode) && !(current instanceof ProjectNode)) { + return false; + } + List inputs = current.getInputs(); + if (inputs.size() != 1) { + return false; + } + current = inputs.get(0); + } + } + + /** + * Returns {@code true} only if the scanned table is guaranteed to be served by exactly one physical table. + * + *

This is a hard precondition for the k-way merge: a scan over a hybrid (OFFLINE + REALTIME) table is compiled + * into two {@code ServerQueryRequest}s that {@code LeafOperator} runs concurrently, pushing both result sets into the + * same mailbox with no cross-request merge. That mailbox stream is the concatenation of two independently sorted + * runs, not a sorted stream, and the merge would silently emit rows in the wrong order. The same applies to a logical + * table, which can fan out to several physical tables. + * + *

Fails closed: an unknown table or a missing {@link TableCache} yields {@code false}, so the receiver keeps the + * accumulate-then-sort path. + */ + private boolean resolvesToSinglePhysicalTable(String tableName) { + if (_tableCache == null) { + return false; + } + // An explicit type suffix (t_OFFLINE / t_REALTIME) already pins the scan to one physical table. + if (TableNameBuilder.getTableTypeFromTableName(tableName) != null) { + return true; + } + String actualTableName = _tableCache.getActualTableName(tableName); + if (actualTableName == null) { + // Unknown to the table cache: it may be a logical table or simply absent. Either way, do not mark. + return false; + } + if (TableNameBuilder.getTableTypeFromTableName(actualTableName) != null) { + return true; + } + if (_tableCache.isLogicalTable(actualTableName)) { + return false; + } + boolean hasOffline = + _tableCache.getTableConfig(TableNameBuilder.forType(TableType.OFFLINE).tableNameWithType(actualTableName)) + != null; + boolean hasRealtime = + _tableCache.getTableConfig(TableNameBuilder.forType(TableType.REALTIME).tableNameWithType(actualTableName)) + != null; + return hasOffline != hasRealtime; + } + + /** + * Returns {@code true} if the two collation lists are equivalent for sorted-merge purposes: same size and, for each + * position, equal field index and equal direction (and equal null direction). A {@code null} list (e.g. a plain, + * non-sorted exchange) is treated as "not a sorted collation" and yields {@code false}. + * + *

An empty list is likewise rejected, and that case is load-bearing rather than cosmetic. A plain + * {@code SELECT ... LIMIT n} with no ORDER BY compiles to a collation-less {@code LogicalSort} (fetch only) under a + * collation-less sort exchange, so both lists are empty and an element-wise comparison alone would call them + * "matching". That would mark the receive as sorted-on-sender, and {@code SortedMailboxReceiveOperator} rejects an + * empty collation outright, failing every such query. There is also nothing to merge on without a collation. + */ + private static boolean collationsMatch(@Nullable List a, @Nullable List b) { + if (a == null || b == null || a.isEmpty() || a.size() != b.size()) { + return false; + } + for (int i = 0; i < a.size(); i++) { + RelFieldCollation ca = a.get(i); + RelFieldCollation cb = b.get(i); + if (ca.getFieldIndex() != cb.getFieldIndex() || ca.getDirection() != cb.getDirection() + || ca.nullDirection != cb.nullDirection) { + return false; + } + } + return true; + } + public static class Context { private final int _currentPlanFragmentId; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/PlanFragmenterTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/PlanFragmenterTest.java new file mode 100644 index 000000000000..b577378ca5a2 --- /dev/null +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/PlanFragmenterTest.java @@ -0,0 +1,274 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.planner.logical; + +import java.util.ArrayList; +import java.util.List; +import java.util.function.Function; +import java.util.function.Supplier; +import javax.annotation.Nullable; +import org.apache.calcite.rel.RelDistribution; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.calcite.rel.core.JoinRelType; +import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType; +import org.apache.pinot.common.config.provider.TableCache; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.planner.plannode.AggregateNode; +import org.apache.pinot.query.planner.plannode.ExchangeNode; +import org.apache.pinot.query.planner.plannode.FilterNode; +import org.apache.pinot.query.planner.plannode.JoinNode; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; +import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.planner.plannode.ProjectNode; +import org.apache.pinot.query.planner.plannode.SortNode; +import org.apache.pinot.query.planner.plannode.TableScanNode; +import org.apache.pinot.query.planner.plannode.WindowNode; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.mockito.Mockito; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +/** + * Tests the {@code sortedOnSender} marking that {@link PlanFragmenter} applies when the + * {@code sortedSelectionMergeEnabled} query option is on. + * + *

This marking is the only gate between the query option and the k-way merge in + * {@code SortedMailboxReceiveOperator}, and a false positive is not a crash: the merge assumes each mailbox stream is + * sorted, nothing downstream re-sorts, and under {@code LIMIT}/{@code OFFSET} the query simply returns the wrong rows. + * So every reject branch is pinned here, not just the accept path. + */ +public class PlanFragmenterTest { + private static final String OFFLINE_ONLY_TABLE = "offlineOnlyTable"; + private static final String REALTIME_ONLY_TABLE = "realtimeOnlyTable"; + private static final String HYBRID_TABLE = "hybridTable"; + private static final String LOGICAL_TABLE = "logicalTable"; + private static final String UNKNOWN_TABLE = "unknownTable"; + + private static final DataSchema SCHEMA = + new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{ + DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING + }); + private static final List COLLATION = List.of(new RelFieldCollation(0)); + + // -------------------------------------------------------------------------------------------------------------- + // Sender-fragment shapes + // -------------------------------------------------------------------------------------------------------------- + + @DataProvider(name = "senderShapes") + public Object[][] senderShapes() { + return new Object[][]{ + // {description, sender fragment root, expected sortedOnSender when the option is ON} + {"sort over scan", (Supplier) () -> sort(scan(OFFLINE_ONLY_TABLE)), true}, + {"sort over project over scan", (Supplier) () -> sort(project(scan(OFFLINE_ONLY_TABLE))), true}, + {"sort over two projects over scan", + (Supplier) () -> sort(project(project(scan(OFFLINE_ONLY_TABLE)))), true}, + {"realtime-only table", (Supplier) () -> sort(scan(REALTIME_ONLY_TABLE)), true}, + {"explicit OFFLINE type suffix", (Supplier) () -> sort(scan(HYBRID_TABLE + "_OFFLINE")), true}, + {"explicit REALTIME type suffix", (Supplier) () -> sort(scan(HYBRID_TABLE + "_REALTIME")), true}, + + // Rejected: not a leaf selection ORDER BY. + {"bare scan (no sort)", (Supplier) () -> scan(OFFLINE_ONLY_TABLE), false}, + {"project root", (Supplier) () -> project(scan(OFFLINE_ONLY_TABLE)), false}, + {"filter in the chain", (Supplier) () -> sort(filter(scan(OFFLINE_ONLY_TABLE))), false}, + {"aggregate below sort", (Supplier) () -> sort(aggregate(scan(OFFLINE_ONLY_TABLE))), false}, + {"window below sort", (Supplier) () -> sort(window(scan(OFFLINE_ONLY_TABLE))), false}, + {"join below sort", + (Supplier) () -> sort(join(scan(OFFLINE_ONLY_TABLE), scan(REALTIME_ONLY_TABLE))), false}, + + // Rejected: the scan does not resolve to exactly one physical table, so a single mailbox stream can carry two + // independently sorted runs (hybrid) or several tables (logical). + {"hybrid table", (Supplier) () -> sort(scan(HYBRID_TABLE)), false}, + {"logical table", (Supplier) () -> sort(scan(LOGICAL_TABLE)), false}, + {"unknown table", (Supplier) () -> sort(scan(UNKNOWN_TABLE)), false} + }; + } + + @Test(dataProvider = "senderShapes") + public void testSortedOnSenderMarkingByShape(String description, Supplier senderRoot, + boolean expectedWhenEnabled) { + assertEquals(fragmentAndGetSortedOnSender(senderRoot.get(), COLLATION, true), expectedWhenEnabled, description); + } + + @Test(dataProvider = "senderShapes") + public void testNoMarkingWhenOptionDisabled(String description, Supplier senderRoot, + boolean expectedWhenEnabled) { + // With the option off the flag must always equal the exchange's own sortOnSender (false here), regardless of shape. + assertFalse(fragmentAndGetSortedOnSender(senderRoot.get(), COLLATION, false), description); + } + + // -------------------------------------------------------------------------------------------------------------- + // Collation matching + // -------------------------------------------------------------------------------------------------------------- + + @DataProvider(name = "collations") + public Object[][] collations() { + RelFieldCollation asc0 = new RelFieldCollation(0); + RelFieldCollation desc0 = + new RelFieldCollation(0, RelFieldCollation.Direction.DESCENDING, RelFieldCollation.NullDirection.LAST); + RelFieldCollation asc0NullsFirst = + new RelFieldCollation(0, RelFieldCollation.Direction.ASCENDING, RelFieldCollation.NullDirection.FIRST); + RelFieldCollation asc1 = new RelFieldCollation(1); + return new Object[][]{ + {"identical single key", List.of(asc0), List.of(asc0), true}, + {"identical two keys", List.of(asc0, asc1), List.of(asc0, asc1), true}, + {"identical descending", List.of(desc0), List.of(desc0), true}, + {"different field index", List.of(asc0), List.of(asc1), false}, + {"different direction", List.of(asc0), List.of(desc0), false}, + {"different null direction", List.of(asc0), List.of(asc0NullsFirst), false}, + {"different size", List.of(asc0, asc1), List.of(asc0), false}, + {"exchange has no collation", List.of(asc0), null, false}, + // A collation-less LogicalSort (plain `SELECT ... LIMIT n`, no ORDER BY) below a collation-less sort exchange. + // Both lists are empty, so a naive element-wise comparison would call them "matching" and mark the receive as + // sorted-on-sender; SortedMailboxReceiveOperator then rejects the empty collation and the query fails. + {"both collations empty", List.of(), List.of(), false}, + {"sort has no collation, exchange does", List.of(), List.of(asc0), false} + }; + } + + @Test(dataProvider = "collations") + public void testSortedOnSenderMarkingByCollation(String description, List sortCollation, + @Nullable List exchangeCollation, boolean expected) { + PlanNode senderRoot = new SortNode(0, SCHEMA, PlanNode.NodeHint.EMPTY, mutable(scan(OFFLINE_ONLY_TABLE)), + sortCollation, 10, 0); + assertEquals(fragmentAndGetSortedOnSender(senderRoot, exchangeCollation, true), expected, description); + } + + /** + * When the exchange itself already declares {@code sortOnSender}, the flag must stay set no matter what the option or + * the shape gate says: the marking is an additional source of truth, never a filter on the existing one. + */ + @Test + public void testExistingSortOnSenderIsPreserved() { + // A shape the gate rejects (hybrid table), with sortOnSender declared on the exchange. + assertTrue(fragmentAndGetSortedOnSender(sort(scan(HYBRID_TABLE)), COLLATION, false, true)); + assertTrue(fragmentAndGetSortedOnSender(sort(scan(HYBRID_TABLE)), COLLATION, true, true)); + } + + /** + * A null {@link TableCache} must fail closed: no table can be proven single-physical, so nothing is marked. + */ + @Test + public void testNullTableCacheFailsClosed() { + PlanNode receiverRoot = exchange(sort(scan(OFFLINE_ONLY_TABLE)), COLLATION, false); + PlanFragmenter fragmenter = new PlanFragmenter(true, null); + PlanNode result = receiverRoot.visit(fragmenter, fragmenter.createContext()); + assertFalse(((MailboxReceiveNode) result).isSortedOnSender()); + } + + // -------------------------------------------------------------------------------------------------------------- + // Helpers + // -------------------------------------------------------------------------------------------------------------- + + private boolean fragmentAndGetSortedOnSender(PlanNode senderRoot, @Nullable List exchangeCollation, + boolean optionEnabled) { + return fragmentAndGetSortedOnSender(senderRoot, exchangeCollation, optionEnabled, false); + } + + private boolean fragmentAndGetSortedOnSender(PlanNode senderRoot, @Nullable List exchangeCollation, + boolean optionEnabled, boolean exchangeSortOnSender) { + PlanNode receiverRoot = exchange(senderRoot, exchangeCollation, exchangeSortOnSender); + PlanFragmenter fragmenter = new PlanFragmenter(optionEnabled, mockTableCache()); + PlanNode result = receiverRoot.visit(fragmenter, fragmenter.createContext()); + return ((MailboxReceiveNode) result).isSortedOnSender(); + } + + private static ExchangeNode exchange(PlanNode input, @Nullable List collations, + boolean sortOnSender) { + return new ExchangeNode(0, SCHEMA, mutable(input), PinotRelExchangeType.getDefaultExchangeType(), + RelDistribution.Type.HASH_DISTRIBUTED, List.of(0), false, collations, sortOnSender, false, null, null, + "absHashCode"); + } + + private static SortNode sort(PlanNode input) { + return new SortNode(0, SCHEMA, PlanNode.NodeHint.EMPTY, mutable(input), COLLATION, 10, 0); + } + + private static ProjectNode project(PlanNode input) { + return new ProjectNode(0, SCHEMA, PlanNode.NodeHint.EMPTY, mutable(input), List.of()); + } + + private static FilterNode filter(PlanNode input) { + return new FilterNode(0, SCHEMA, PlanNode.NodeHint.EMPTY, mutable(input), null); + } + + private static AggregateNode aggregate(PlanNode input) { + return new AggregateNode(0, SCHEMA, PlanNode.NodeHint.EMPTY, mutable(input), List.of(), List.of(), List.of(0), + AggregateNode.AggType.DIRECT, false, null, 0); + } + + private static WindowNode window(PlanNode input) { + return new WindowNode(0, SCHEMA, PlanNode.NodeHint.EMPTY, mutable(input), List.of(0), COLLATION, List.of(), + WindowNode.WindowFrameType.ROWS, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowExclusion.NO_OTHERS, + List.of()); + } + + private static JoinNode join(PlanNode left, PlanNode right) { + return new JoinNode(0, SCHEMA, PlanNode.NodeHint.EMPTY, new ArrayList<>(List.of(left, right)), + JoinRelType.INNER, List.of(0), List.of(0), List.of(), JoinNode.JoinStrategy.HASH); + } + + private static List mutable(PlanNode input) { + return new ArrayList<>(List.of(input)); + } + + private static TableScanNode scan(String tableName) { + return new TableScanNode(0, SCHEMA, PlanNode.NodeHint.EMPTY, new ArrayList<>(), tableName, List.of("col1", "col2")); + } + + /** + * A table cache holding one offline-only table, one realtime-only table, one hybrid table and one logical table. + */ + private static TableCache mockTableCache() { + TableCache tableCache = Mockito.mock(TableCache.class); + Function actualName = name -> { + String raw = name.replace("_OFFLINE", "").replace("_REALTIME", ""); + return List.of(OFFLINE_ONLY_TABLE, REALTIME_ONLY_TABLE, HYBRID_TABLE, LOGICAL_TABLE).contains(raw) + ? name : null; + }; + Mockito.when(tableCache.getActualTableName(Mockito.anyString())) + .thenAnswer(invocation -> actualName.apply(invocation.getArgument(0))); + Mockito.when(tableCache.isLogicalTable(Mockito.anyString())) + .thenAnswer(invocation -> LOGICAL_TABLE.equals(invocation.getArgument(0))); + Mockito.when(tableCache.getTableConfig(Mockito.anyString())).thenAnswer(invocation -> { + String nameWithType = invocation.getArgument(0); + switch (nameWithType) { + case OFFLINE_ONLY_TABLE + "_OFFLINE": + case REALTIME_ONLY_TABLE + "_REALTIME": + case HYBRID_TABLE + "_OFFLINE": + case HYBRID_TABLE + "_REALTIME": + return tableConfig(nameWithType); + default: + return null; + } + }); + return tableCache; + } + + private static TableConfig tableConfig(String tableNameWithType) { + TableType type = tableNameWithType.endsWith("_REALTIME") ? TableType.REALTIME : TableType.OFFLINE; + return new TableConfigBuilder(type).setTableName(tableNameWithType).build(); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java index 635fd3ace2aa..030c3409796b 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java @@ -122,6 +122,16 @@ public MultiStageQueryStats calculateUpstreamStats() { return _multiConsumer.calculateStats(); } + /** + * Returns one per-sender {@link BlockingMultiStreamConsumer.StreamHandle} so a subclass can read each sender mailbox + * independently (used by the streaming k-way merge). Latches the underlying consumer into per-stream mode: a subclass + * must use either these handles or {@code readMseBlockBlocking()}, never both. Returns an empty list when there are + * no mailboxes. + */ + protected List> streamHandles() { + return _multiConsumer.streamHandles(); + } + @Override public StatMap copyStatMaps() { return new StatMap<>(_statMap); @@ -266,7 +276,28 @@ public int merge(int value1, int value2) { /** * Time spent on GC while this operator or its children in the same stage were running. */ - GC_TIME_MS(StatMap.Type.LONG); + GC_TIME_MS(StatMap.Type.LONG), + /** + * Whether this receive operator served its rows with the streaming k-way merge rather than the + * accumulate-then-sort path. + *

+ * Only {@code SortedMailboxReceiveOperator} ever sets this; the unsorted receive operator always leaves it at + * {@code false}. It is recorded at construction time (not when the first row is merged), so it reports the path + * the operator was configured to take even when the result is empty. + *

+ * Reporting is presence-based: {@link StatMap} drops boolean keys whose value is {@code false}, so this renders + * into the response {@code stageStats} as {@code kWayMergeUsed: true} on the merge path and is simply absent on + * the accumulate-then-sort path (and on any server predating this stat). + *

+ * NOTE: this key is intentionally the last constant in this enum. {@link StatMap} serializes keys by ordinal, so + * new keys must only ever be appended, never inserted or reordered. + */ + K_WAY_MERGE_USED(StatMap.Type.BOOLEAN) { + @Override + public boolean includeDefaultInJson() { + return true; + } + }; private final StatMap.Type _type; 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 739a13c88689..98f34303ba11 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 @@ -19,15 +19,24 @@ package org.apache.pinot.query.runtime.operator; import com.google.common.base.Preconditions; +import java.util.ArrayDeque; import java.util.ArrayList; +import java.util.Comparator; +import java.util.Deque; +import java.util.Iterator; import java.util.List; +import java.util.PriorityQueue; +import javax.annotation.Nullable; import org.apache.calcite.rel.RelFieldCollation; import org.apache.commons.collections4.CollectionUtils; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.query.mailbox.ReceivingMailbox; import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; +import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; +import org.apache.pinot.query.runtime.operator.utils.BlockingMultiStreamConsumer.StreamHandle; import org.apache.pinot.query.runtime.operator.utils.SortUtils; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; @@ -35,29 +44,77 @@ /** - * This {@code SortedMailboxReceiveOperator} receives data from a {@link ReceivingMailbox} and serve it out from the - * {@link #nextBlock()} API in a sorted manner. + * This {@code SortedMailboxReceiveOperator} receives data from a {@link ReceivingMailbox} and serves it out from the + * {@link #nextBlock()} API in a globally sorted manner. * - * TODO: Once sorting on the {@code MailboxSendOperator} is available, modify this to use a k-way merge instead of - * resorting via the PriorityQueue. + *

It supports two strategies, selected at construction time: + *

    + *
  • Accumulate-then-sort (default): every row from every mailbox is buffered, sorted once at EOS, and + * returned in a single data block. This is the historical behavior and is used whenever the k-way merge is not + * enabled.
    • + *
  • Streaming k-way merge: when the sender is known to emit each mailbox already sorted on the receiver's + * collation, the rows are merged incrementally with a min-heap and emitted in bounded blocks (of at most + * {@code blockSize} rows). Global order is preserved across block boundaries because the heap state carries over + * between {@link #getNextBlock()} calls. Senders are deliberately allowed to backpressure — that is where the + * memory advantage comes from — but no sender is left parked indefinitely: every {@link #refill} takes at most + * one element from each sibling stream, and a stalled merge drains ready siblings fully into a per-stream + * backlog. Under sustained key skew that backlog trades some of the memory advantage for liveness.
    • + *
+ * The k-way merge is enabled only when the {@code streamingSortedMailboxReceive} query option is + * {@code true} and {@link MailboxReceiveNode#isSortedOnSender()} is true. All other combinations fall back to + * the accumulate-then-sort path. Which path was taken is reported in the query response {@code stageStats} as the + * {@code kWayMergeUsed} stat (see {@link BaseMailboxReceiveOperator.StatKey#K_WAY_MERGE_USED}). + * + *

Like the rest of the receive operators, this class is driven by a single consumer thread; it is not thread-safe. */ public class SortedMailboxReceiveOperator extends BaseMailboxReceiveOperator { private static final Logger LOGGER = LoggerFactory.getLogger(SortedMailboxReceiveOperator.class); private static final String EXPLAIN_NAME = "SORTED_MAILBOX_RECEIVE"; + /** + * Default upper bound on the number of rows emitted per block in the streaming k-way merge, used when the + * {@code streamingSortedMailboxReceiveBlockSize} query option is not set. Defined locally to avoid introducing a + * dependency on {@code pinot-core} (where {@code SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY} lives). + */ + private static final int DEFAULT_BLOCK_SIZE = 10_000; + private final DataSchema _dataSchema; private final List _collations; private final List _rows = new ArrayList<>(); + // Streaming k-way merge state. The merge-only fields are meaningful only when _useKWayMerge is true. + private final boolean _useKWayMerge; + private final int _blockSize; + private final Comparator _comparator; + // Built lazily on the first merge call so priming (driving every handle to first-data/EOS/error) happens once. + private PriorityQueue _heap; + // Per-sender merge state (backlog + exhaustion), built once alongside the heap during priming. + private List _streams; + // Mutable view of _streams used only for the sibling drain scan in refill(): entries are removed once their stream + // is exhausted so a stalled call doesn't keep re-scanning senders that can never produce more data. + private List _activeStreams; + private boolean _primed; + // Last row handed out by the merge, used to verify the sender-sorted precondition (see checkNonDecreasing). + @Nullable + private Object[] _lastEmitted; + private MseBlock _eosBlock; - // TODO: Support merge sort when sender side sort is supported. public SortedMailboxReceiveOperator(OpChainExecutionContext context, MailboxReceiveNode node) { super(context, node); Preconditions.checkState(!CollectionUtils.isEmpty(node.getCollations()), "Field collations must be set"); _dataSchema = node.getDataSchema(); _collations = node.getCollations(); + // reverse=false => the collation minimum sits at the min-heap head (honors per-field ASC/DESC + null direction). + _comparator = new SortUtils.SortComparator(_collations, false); + _useKWayMerge = QueryOptionsUtils.isStreamingSortedMailboxReceiveEnabled(context.getOpChainMetadata()) + && node.isSortedOnSender(); + // Recorded eagerly (rather than on the first merged block) so the stat reports the configured path even when the + // query returns no rows. Surfaces in the query response stageStats as "kWayMergeUsed". + _statMap.merge(StatKey.K_WAY_MERGE_USED, _useKWayMerge); + Integer blockSize = QueryOptionsUtils.getStreamingSortedMailboxReceiveBlockSize(context.getOpChainMetadata()); + _blockSize = blockSize != null ? blockSize : DEFAULT_BLOCK_SIZE; } @Override @@ -75,6 +132,9 @@ protected MseBlock getNextBlock() { if (_eosBlock != null) { return _eosBlock; } + if (_useKWayMerge) { + return getNextMergedBlock(); + } // Collect all the rows from the mailbox and sort them while (true) { MseBlock block = _multiConsumer.readMseBlockBlocking(); @@ -89,9 +149,7 @@ protected MseBlock getNextBlock() { return eosBlock; } else { 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)); + _rows.sort(_comparator); return new RowHeapDataBlock(_rows, _dataSchema); } else { return block; @@ -100,15 +158,356 @@ protected MseBlock getNextBlock() { } } + /** + * Streaming k-way merge over the per-sender {@link StreamHandle}s. Emits at most {@link #_blockSize} rows per call, + * keeping the heap state between calls so global order is preserved across blocks. + */ + private MseBlock getNextMergedBlock() { + if (_isEarlyTerminated) { + // Stop pulling new data; drive every handle to EOS (so receiving stats are folded in) and finish. + return drainToEos(); + } + if (!_primed) { + List> handles = streamHandles(); + Comparator comparator = _comparator; + _heap = new PriorityQueue<>(Math.max(1, handles.size()), + (a, b) -> comparator.compare(a.head(), b.head())); + _streams = new ArrayList<>(handles.size()); + for (StreamHandle handle : handles) { + _streams.add(new StreamState(handle)); + } + // Built before priming: refill()'s sibling drain scan runs during priming too (a stream can be starved on its + // very first poll), so _activeStreams must already reflect every stream. + _activeStreams = new ArrayList<>(_streams); + // Prime: drive every stream to its first data block / EOS / error before the first pop, so the heap holds a head + // for every still-active mailbox and the min is the global min. + for (StreamState state : _streams) { + Cursor cursor = refill(state); + if (_eosBlock != null) { + // An error was found while priming; refill already cached it and folded stats. + return _eosBlock; + } + if (cursor != null) { + _heap.add(cursor); + } + } + _primed = true; + } + + // Initial capacity is capped at DEFAULT_BLOCK_SIZE so a very large configured _blockSize does not eagerly allocate + // a huge backing array up front; for larger blocks the list grows amortized as rows are appended. + List out = new ArrayList<>(Math.min(_blockSize, DEFAULT_BLOCK_SIZE)); + while (out.size() < _blockSize) { + if (_heap.isEmpty()) { + onEos(); + _eosBlock = SuccessMseBlock.INSTANCE; + return out.isEmpty() ? _eosBlock : new RowHeapDataBlock(out, _dataSchema); + } + Cursor cursor = _heap.poll(); + Object[] row = cursor.head(); + checkNonDecreasing(row, cursor); + out.add(row); + _lastEmitted = row; + if (cursor.advance()) { + // Still has rows in the current block: reseat with the new head. + _heap.add(cursor); + } else { + // Current block exhausted: refill THIS mailbox before the next pop to restore the heap invariant. + Cursor refilled = refill(cursor._state); + if (_eosBlock != null) { + // An error was found while refilling; short-circuit immediately. + return _eosBlock; + } + if (refilled != null) { + _heap.add(refilled); + } + // else this mailbox reached EOS and is dropped from the merge. + } + } + return new RowHeapDataBlock(out, _dataSchema); + } + + /** + * Verifies the precondition the whole merge rests on: each mailbox stream is already sorted on the receiver's + * collation. Only that guarantee makes "min of the heads" the global min, and nothing downstream re-sorts — + * {@code SortOperator} skips its priority queue when its input is a {@code SortedMailboxReceiveOperator}. If a sender + * violates it (a plan shape the fragmenter gate should have rejected, a leaf that concatenates two independently + * sorted runs, a collation mismatch), the merge would otherwise emit misordered rows and, under LIMIT/OFFSET, simply + * return the wrong rows with no error anywhere. One comparison per emitted row converts that into a hard failure. + */ + private void checkNonDecreasing(Object[] row, Cursor cursor) { + if (_lastEmitted != null && _comparator.compare(_lastEmitted, row) > 0) { + throw new IllegalStateException( + "Sorted mailbox receive got out-of-order rows from mailbox: " + cursor._state._handle.getId() + + ". The sender was declared sorted-on-sender but is not sorted on the receiver's collation"); + } + } + + /** + * Produces the next {@link Cursor} for {@code state} without head-of-line blocking. The merge can only emit once it + * has the next row from {@code state}, but it must never park on {@code state} alone while sibling mailboxes fill up: + * that lets senders backpressure and deadlocks the single-threaded pipeline. So while {@code state} has nothing ready + * this drains any other ready stream into its backlog (relieving that sender), and only parks (on the shared + * new-data signal) when no stream anywhere has data. Buffered rows keep per-stream order, so global sort order is + * preserved. Returns {@code null} when {@code state} reaches success EOS (dropped from the merge) or on error (which + * is cached in {@link #_eosBlock} after folding stats via {@link #onEos()}). + */ + @Nullable + private Cursor refill(StreamState state) { + // Bounded relief pass: take at most one element from every other active stream before serving this one. Without + // it, siblings are polled only while the merge is stalled, so with disjoint key ranges a fast sender can sit + // parked on a full mailbox (capacity ReceivingMailbox.DEFAULT_MAX_PENDING_BLOCKS) for the whole query, holding an + // MSE worker thread. One poll per refill (i.e. per consumed block, not per row) keeps every sender advancing + // without eagerly draining a fast sender's entire output into the backlog, which would give back the memory + // advantage the merge exists for. + relieveSiblingsOnce(state); + if (_eosBlock != null) { + return null; + } + while (true) { + if (!state._backlog.isEmpty()) { + return new Cursor(state, state._backlog.poll()); + } + if (_eosBlock != null) { + return null; + } + if (state._handle.isExhausted()) { + // Success EOS already seen and backlog drained: drop this mailbox from the merge. + return null; + } + // Try to advance THIS stream without blocking. + if (pollOnce(state)) { + if (state._handle.isExhausted()) { + _activeStreams.remove(state); + } + // Buffered rows (loop serves the backlog), hit success EOS (loop returns null), or read an empty block (retry). + continue; + } + // This stream has nothing ready. Drain any OTHER ready stream to relieve its sender's backpressure; that may in + // turn unblock the sender feeding this stream. Streams that reach exhaustion are pruned from _activeStreams so + // later calls (for any mailbox) don't keep re-scanning senders that can never produce more data. + boolean progressed = false; + Iterator it = _activeStreams.iterator(); + while (it.hasNext()) { + StreamState other = it.next(); + if (other == state) { + continue; + } + while (pollOnce(other)) { + progressed = true; + if (_eosBlock != null) { + return null; + } + } + if (other._handle.isExhausted()) { + it.remove(); + } + } + if (progressed) { + // Draining may have delivered data (or woken this stream's sender); retry before parking. + continue; + } + // Nothing ready anywhere: park until any stream signals new data (or the deadline is hit). + ReceivingMailbox.MseBlockWithStats timedOut = _multiConsumer.awaitDataOrTerminal(); + if (timedOut != null) { + onEos(); + _eosBlock = timedOut.getBlock(); + return null; + } + // Woken: loop and retry. + } + } + + /** + * Polls every active stream other than {@code state} at most once, non-blocking, pruning any that become exhausted. + * See {@link #refill} for why this runs unconditionally rather than only when the merge stalls. + */ + private void relieveSiblingsOnce(StreamState state) { + if (_eosBlock != null) { + return; + } + Iterator it = _activeStreams.iterator(); + while (it.hasNext()) { + StreamState other = it.next(); + if (other == state) { + continue; + } + pollOnce(other); + if (_eosBlock != null) { + return; + } + if (other._handle.isExhausted()) { + it.remove(); + } + } + } + + /** + * Polls one stream once (non-blocking). Buffers any non-empty data rows into the stream's backlog and caches an error + * into {@link #_eosBlock}. Exhaustion itself is tracked by the underlying {@link StreamHandle#isExhausted()}, not + * duplicated here. Returns {@code true} if any element (data, success EOS, or error) was read, {@code false} if the + * stream had nothing ready. + */ + private boolean pollOnce(StreamState state) { + if (state._handle.isExhausted() || _eosBlock != null) { + return false; + } + ReceivingMailbox.MseBlockWithStats element = state._handle.poll(); + if (element == null) { + return false; + } + MseBlock block = element.getBlock(); + if (block.isError()) { + onEos(); + _eosBlock = block; + return true; + } + if (block.isSuccess()) { + return true; + } + List rows = ((MseBlock.Data) block).asRowHeap().getRows(); + if (!rows.isEmpty()) { + state._backlog.add(rows); + } + // Empty data blocks carry no head; returning true lets refill loop and poll again. + return true; + } + + /** + * Drains every handle to its terminal element after early termination, folding receiving stats. Like {@link + * #refill}, this must not head-of-line block on one handle while sibling mailboxes still have data buffered: doing so + * would let their senders backpressure and deadlock the pipeline, exactly as it would during normal merging. So this + * polls every not-yet-exhausted handle in round-robin passes (discarding data, since early termination means the + * result is no longer needed) and only parks when a full pass makes no progress on any handle. Returns the cached + * error block if any handle yields one, otherwise a success EOS. + */ + private MseBlock drainToEos() { + List> handles = streamHandles(); + int numRemaining = 0; + boolean[] exhausted = new boolean[handles.size()]; + for (int i = 0; i < handles.size(); i++) { + if (handles.get(i).isExhausted()) { + exhausted[i] = true; + } else { + numRemaining++; + } + } + while (numRemaining > 0) { + boolean progressed = false; + for (int i = 0; i < handles.size(); i++) { + if (exhausted[i]) { + continue; + } + StreamHandle handle = handles.get(i); + ReceivingMailbox.MseBlockWithStats element = handle.poll(); + if (element == null) { + continue; + } + progressed = true; + MseBlock block = element.getBlock(); + if (block.isError()) { + onEos(); + _eosBlock = block; + return block; + } + if (handle.isExhausted()) { + exhausted[i] = true; + numRemaining--; + } + // Data blocks are discarded; a still-active handle is retried on a later pass. + } + if (!progressed) { + // No handle had anything ready this pass: park until any stream signals new data (or the deadline is hit). + ReceivingMailbox.MseBlockWithStats timedOut = _multiConsumer.awaitDataOrTerminal(); + if (timedOut != null) { + onEos(); + _eosBlock = timedOut.getBlock(); + return _eosBlock; + } + } + } + onEos(); + _eosBlock = SuccessMseBlock.INSTANCE; + return _eosBlock; + } + @Override public void close() { super.close(); _rows.clear(); + clearMergeState(); } @Override public void cancel(Throwable t) { super.cancel(t); _rows.clear(); + clearMergeState(); + } + + private void clearMergeState() { + if (_heap != null) { + _heap.clear(); + } + if (_streams != null) { + for (StreamState state : _streams) { + state._backlog.clear(); + } + } + if (_activeStreams != null) { + _activeStreams.clear(); + } + } + + /** + * Per-sender merge state: the stream handle plus a backlog of data blocks buffered ahead of the merge's current + * position. Rows are staged here (in arrival order, which the sender guarantees is sorted) when the merge drains this + * mailbox while waiting on another stream. Exhaustion is tracked by {@link StreamHandle#isExhausted()} on the handle + * itself, not duplicated here. + */ + private static final class StreamState { + final StreamHandle _handle; + final Deque> _backlog = new ArrayDeque<>(); + + StreamState(StreamHandle handle) { + _handle = handle; + } + } + + /** + * A cursor over one mailbox's current data block. Holds the owning {@link StreamState} so the merge can refill this + * specific mailbox (from its backlog or the stream) when the block is exhausted. Created only for non-empty blocks, + * so {@link #head()} is valid until {@link #advance()} returns {@code false}. + */ + private static final class Cursor { + final StreamState _state; + final List _rows; + int _idx; + // Cached _rows.get(_idx). The heap comparator reads this on every comparison (~2*log2(k) per emitted row), so it + // is kept as a field rather than re-resolved through the List interface each time. + Object[] _head; + + Cursor(StreamState state, List rows) { + _state = state; + _rows = rows; + _head = rows.get(0); + } + + Object[] head() { + return _head; + } + + /** + * Advances past the current row. Returns {@code true} if a new head is available in this block. + */ + boolean advance() { + _idx++; + if (_idx < _rows.size()) { + _head = _rows.get(_idx); + return true; + } + _head = null; + return false; + } } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java index 3a11d73bddef..e7b35ba4c07b 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java @@ -19,6 +19,8 @@ package org.apache.pinot.query.runtime.operator.utils; import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.concurrent.ArrayBlockingQueue; import java.util.concurrent.TimeUnit; @@ -63,6 +65,24 @@ public abstract class BlockingMultiStreamConsumer implements AutoCloseable { @Nullable private E _errorBlock = null; + /** + * A consumer instance reads either in round-robin mode (via {@link #readBlockBlocking()}) or in per-stream mode + * (via {@link #streamHandles()} / {@link StreamHandle#readBlocking()}), never both. The mode is latched on first use + * and mixing the two throws {@link IllegalStateException}. Both modes share the same EOS/error/timeout bookkeeping + * (the abstract hooks below), so stats stay correct regardless of which mode is used. + */ + private enum Mode { + UNSET, ROUND_ROBIN, PER_STREAM + } + + private Mode _mode = Mode.UNSET; + /** + * Lazily built per-stream handles, used only in {@link Mode#PER_STREAM}. Built once from {@link #_mailboxes} (which + * the mode guard keeps the round-robin path from mutating) and cached so {@link #streamHandles()} is idempotent. + */ + @Nullable + private List> _handles = null; + public BlockingMultiStreamConsumer(Object id, long deadlineMs, List> asyncProducers) { _id = id; _deadlineMs = deadlineMs; @@ -199,6 +219,7 @@ protected void onData() { * This method is called by the consumer thread. */ public E readBlockBlocking() { + latchMode(Mode.ROUND_ROBIN); if (LOGGER.isTraceEnabled()) { String mailboxIds = _mailboxes.stream() .map(AsyncStream::getId) @@ -327,6 +348,228 @@ private E readBlockOrNull() { return null; } + /** + * Parks the consumer thread until any stream signals new data or the deadline is reached. Used by the per-stream + * (k-way merge) read mode: after finding every stream momentarily empty, a caller waits here for progress instead of + * committing to a single stream (which would head-of-line block while sibling mailboxes fill up and their senders + * backpressure, deadlocking the pipeline). Returns {@code null} when woken by new data (the caller should re-poll the + * streams), or the terminal error element (already routed through {@link #onTimeout()}) when the deadline is + * exceeded. + * + *

This method is called by the consumer thread. + */ + @Nullable + public E awaitDataOrTerminal() { + latchMode(Mode.PER_STREAM); + if (_errorBlock != null) { + return _errorBlock; + } + try { + long timeoutMs = _deadlineMs - System.currentTimeMillis(); + if (timeoutMs <= 0 || _newDataReady.poll(timeoutMs, TimeUnit.MILLISECONDS) == null) { + _errorBlock = onTimeout(); + return _errorBlock; + } + return null; + } catch (Exception e) { + _errorBlock = onException(e); + return _errorBlock; + } + } + + /** + * Latches the read mode on first use and enforces that a single consumer instance is read in exactly one mode. + * + * @throws IllegalStateException if a different mode was already latched. + */ + private void latchMode(Mode mode) { + if (_mode == Mode.UNSET) { + _mode = mode; + } else if (_mode != mode) { + throw new IllegalStateException("BlockingMultiStreamConsumer mixes round-robin and per-stream reads"); + } + } + + /** + * A narrow per-stream handle for the k-way-merge read mode. Returned by {@link #streamHandles()}. + * + * Unlike {@link #readBlockBlocking()}, which reads from all mailboxes in a fair round-robin and hides which mailbox a + * block came from, a handle reads from one specific stream so a caller (the k-way merge) can advance each sender + * independently. All terminal bookkeeping (success EOS, error, timeout, exception) still routes through the same + * hooks the round-robin path uses, so {@code calculateStats()} stays correct. + * + * All methods are called by the single consumer thread only. + * + * @param the element type, matching the enclosing consumer. + */ + public interface StreamHandle { + /** + * The id of the underlying stream. Mostly used for logging. + */ + Object getId(); + + /** + * Blocking read of the next element from this stream only. Returns a data element, a success-EOS element (after + * which {@link #isExhausted()} is true), or an error/timeout element (after which the whole consumer is in error + * and every handle returns that same error element on subsequent calls). Never returns null. + */ + T readBlocking(); + + /** + * Non-blocking read of the next element from this stream only. Returns a data element, a success-EOS element (after + * which {@link #isExhausted()} is true), an error element (after which the whole consumer is in error), or + * {@code null} when nothing is ready yet. Unlike {@link #readBlocking()} this never parks the consumer thread, so + * the k-way merge can drain whichever siblings are ready while waiting for the specific stream it needs. + */ + @Nullable + T poll(); + + /** + * Returns true once this stream has emitted a success EOS, meaning no more data will come from it. + */ + boolean isExhausted(); + + /** + * Sets the underlying stream to early-terminate state, asking for the metadata block. + */ + void earlyTerminate(); + } + + /** + * Returns one {@link StreamHandle} per mailbox in declaration order (an empty list when there are no mailboxes). + * + * The first call latches this consumer into per-stream mode; subsequent calls to {@link #readBlockBlocking()} throw. + * The returned list is built once and cached, so repeated calls return the same handles. All returned handles must be + * driven by the single consumer thread (they share this consumer's wakeup and error state with no extra + * synchronization). + */ + public List> streamHandles() { + latchMode(Mode.PER_STREAM); + if (_handles == null) { + List> handles = new ArrayList<>(_mailboxes.size()); + for (AsyncStream mailbox : _mailboxes) { + handles.add(new Handle(mailbox)); + } + _handles = Collections.unmodifiableList(handles); + } + return _handles; + } + + /** + * Per-stream handle implementation. Reads from a single captured {@link AsyncStream} (not an index into the + * round-robin-mutated {@link #_mailboxes}), reusing the shared {@link #_newDataReady} wakeup and {@link #_deadlineMs} + * deadline. A wakeup meant for another stream simply causes a re-poll that returns null and loops, which is correct + * because {@link AsyncStream#poll()} reads from the per-mailbox queue, not from {@link #_newDataReady}. + */ + private class Handle implements StreamHandle { + private final AsyncStream _stream; + private boolean _exhausted; + /** + * The success-EOS element seen on this stream. Cached so that, once exhausted, we return it without polling an + * already-released mailbox again (and without re-merging its stats). + */ + @Nullable + private E _eosElement; + + Handle(AsyncStream stream) { + _stream = stream; + } + + @Override + public Object getId() { + return _stream.getId(); + } + + @Override + public boolean isExhausted() { + return _exhausted; + } + + @Override + public void earlyTerminate() { + _stream.earlyTerminate(); + } + + @Override + public E readBlocking() { + if (_errorBlock != null) { + // A global error (from this or any other handle) short-circuits every handle. + return _errorBlock; + } + if (_exhausted) { + // EOS already seen; do not poll an already-released mailbox again. Stats were merged once when first seen. + assert _eosElement != null : "_eosElement must be set whenever _exhausted is true"; + return _eosElement; + } + // Mirror the round-robin path (readDroppingSuccessEos): a deadline already in the past times out before any read, + // so both modes report a timeout rather than racing a last-moment block. + if (System.currentTimeMillis() > _deadlineMs) { + _errorBlock = onTimeout(); + return _errorBlock; + } + // Optimistic read without waiting. + E block = pollThisStream(); + if (block != null) { + return block; + } + try { + while (true) { + long timeoutMs = _deadlineMs - System.currentTimeMillis(); + if (_newDataReady.poll(timeoutMs, TimeUnit.MILLISECONDS) == null) { + _errorBlock = onTimeout(); + return _errorBlock; + } + block = pollThisStream(); + if (block != null) { + return block; + } + } + } catch (Exception e) { + _errorBlock = onException(e); + return _errorBlock; + } + } + + @Nullable + @Override + public E poll() { + if (_errorBlock != null) { + // A global error (from this or any other handle) short-circuits every handle. + return _errorBlock; + } + if (_exhausted) { + // Success EOS already seen; do not poll an already-released mailbox again. + return null; + } + return pollThisStream(); + } + + /** + * Polls this stream once, routing any terminal element through the shared hooks. + * + * @return the element read (data, success EOS, or error), or null if nothing is ready yet. + */ + @Nullable + private E pollThisStream() { + E block = _stream.poll(); + if (block == null) { + return null; + } + if (isError(block)) { + _errorBlock = block; + onError(block); + return block; + } + if (isSuccess(block)) { + _exhausted = true; + _eosElement = block; + onMailboxSuccess(block); + return block; + } + return block; + } + } + /// A [BlockingMultiStreamConsumer] that reads [ReceivingMailbox.MseBlockWithStats]s. /// /// This class is also the entry point for diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorTestUtil.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorTestUtil.java index 7ba84f28a4cc..bb8ce0b3ef2b 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorTestUtil.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorTestUtil.java @@ -109,7 +109,12 @@ public static ReceivingMailbox.MseBlockWithStats eosWithStats(List s public static OpChainExecutionContext getOpChainContext(MailboxService mailboxService, long deadlineMs, StageMetadata stageMetadata) { - return new OpChainExecutionContext(mailboxService, 0, "cid", deadlineMs, deadlineMs, "brokerId", Map.of(), + return getOpChainContext(mailboxService, deadlineMs, stageMetadata, Map.of()); + } + + public static OpChainExecutionContext getOpChainContext(MailboxService mailboxService, long deadlineMs, + StageMetadata stageMetadata, Map opChainMetadata) { + return new OpChainExecutionContext(mailboxService, 0, "cid", deadlineMs, deadlineMs, "brokerId", opChainMetadata, stageMetadata, stageMetadata.getWorkerMetadataList().get(0), null, true, true); } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java index 0f7cff7c02e2..1505b2200173 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java @@ -18,18 +18,35 @@ */ package org.apache.pinot.query.runtime.operator; +import java.util.ArrayList; import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.apache.calcite.rel.RelDistribution; import org.apache.calcite.rel.RelFieldCollation; import org.apache.calcite.rel.RelFieldCollation.Direction; import org.apache.calcite.rel.RelFieldCollation.NullDirection; +import org.apache.pinot.common.datatable.StatMap; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.mailbox.MailboxService; +import org.apache.pinot.query.mailbox.ReceivingMailbox; +import org.apache.pinot.query.planner.physical.MailboxIdUtils; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.planner.plannode.SortNode; +import org.apache.pinot.query.routing.MailboxInfo; +import org.apache.pinot.query.routing.MailboxInfos; +import org.apache.pinot.query.routing.SharedMailboxInfos; +import org.apache.pinot.query.routing.StageMetadata; import org.apache.pinot.query.routing.VirtualServerAddress; +import org.apache.pinot.query.routing.WorkerMetadata; import org.apache.pinot.query.runtime.blocks.ErrorMseBlock; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.spi.utils.CommonConstants; import org.mockito.Mock; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; @@ -38,10 +55,12 @@ 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.Mockito.eq; import static org.mockito.Mockito.mock; 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; @@ -51,6 +70,12 @@ public class SortOperatorTest { private MultiStageOperator _input; @Mock private VirtualServerAddress _serverAddress; + @Mock + private MailboxService _mailboxService; + @Mock + private ReceivingMailbox _mailbox1; + @Mock + private ReceivingMailbox _mailbox2; @BeforeMethod public void setUp() { @@ -499,6 +524,72 @@ public void shouldPreservePrecision() { assertTrue(operator.nextBlock().isSuccess(), "expected EOS block to propagate"); } + /** + * End-to-end fast-path test: feed a real {@link SortedMailboxReceiveOperator} running in k-way MERGE mode (block + * size = 2, so it emits MULTIPLE bounded data blocks) into a {@link SortOperator} with a LIMIT. Because the input is + * a {@code SortedMailboxReceiveOperator}, the SortOperator skips re-sorting and just slices to the limit. The final + * output must be globally sorted and honor the limit across the bounded blocks. + */ + @Test + public void shouldSliceLimitOverMergedBoundedBlocksWithoutResorting() { + DataSchema schema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT}); + List collations = List.of(new RelFieldCollation(0, Direction.ASCENDING, NullDirection.LAST)); + // Two pre-sorted sender streams; the k-way merge produces a globally sorted stream 1..6. + String mailboxId1 = MailboxIdUtils.toMailboxId(0, 1, 0, 0, 0); + String mailboxId2 = MailboxIdUtils.toMailboxId(0, 1, 1, 0, 0); + when(_mailboxService.getHostname()).thenReturn("localhost"); + when(_mailboxService.getPort()).thenReturn(1234); + when(_mailbox1.getStatMap()).thenReturn(new StatMap<>(ReceivingMailbox.StatKey.class)); + when(_mailbox2.getStatMap()).thenReturn(new StatMap<>(ReceivingMailbox.StatKey.class)); + when(_mailboxService.getReceivingMailbox(eq(mailboxId1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(schema, new Object[]{1, 1}, new Object[]{3, 3}, new Object[]{5, 5}), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(mailboxId2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(schema, new Object[]{2, 2}, new Object[]{4, 4}, new Object[]{6, 6}), + OperatorTestUtil.eosWithEmptyStats()); + + MailboxInfos mailboxInfos = new SharedMailboxInfos(new MailboxInfo("localhost", 1234, List.of(0, 1))); + StageMetadata stageMetadata = new StageMetadata(0, + Stream.of(0, 1).map(workerId -> new WorkerMetadata(workerId, Map.of(1, mailboxInfos), Map.of())) + .collect(Collectors.toList()), Map.of()); + Map opChainMetadata = Map.of( + CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true", + CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE, "2"); + OpChainExecutionContext receiveContext = + OperatorTestUtil.getOpChainContext(_mailboxService, Long.MAX_VALUE, stageMetadata, opChainMetadata); + MailboxReceiveNode receiveNode = mock(MailboxReceiveNode.class); + when(receiveNode.getDistributionType()).thenReturn(RelDistribution.Type.HASH_DISTRIBUTED); + when(receiveNode.getSenderStageId()).thenReturn(1); + when(receiveNode.getDataSchema()).thenReturn(schema); + when(receiveNode.getCollations()).thenReturn(collations); + when(receiveNode.isSortedOnSender()).thenReturn(true); + + try (SortedMailboxReceiveOperator receiveOperator = new SortedMailboxReceiveOperator(receiveContext, receiveNode)) { + // fetch = 4, offset = 1 => keep merged rows at indices 1..4 (values 2, 3, 4, 5). + SortOperator operator = new SortOperator(OperatorTestUtil.getTracingContext(), receiveOperator, + new SortNode(-1, schema, PlanNode.NodeHint.EMPTY, List.of(), collations, 4, 1)); + + List resultRows = new ArrayList<>(); + MseBlock block = operator.nextBlock(); + while (block.isData()) { + resultRows.addAll(((MseBlock.Data) block).asRowHeap().getRows()); + block = operator.nextBlock(); + } + assertTrue(block.isSuccess(), "expected EOS block to propagate"); + assertEquals(resultRows.size(), 4, "limit (fetch=4) must be honored across bounded blocks"); + assertEquals(resultRows.get(0), new Object[]{2, 2}); + assertEquals(resultRows.get(1), new Object[]{3, 3}); + assertEquals(resultRows.get(2), new Object[]{4, 4}); + assertEquals(resultRows.get(3), new Object[]{5, 5}); + // Prove the fast-path was taken: the SortOperator must NOT have built a priority queue (no re-sort) because the + // input is a SortedMailboxReceiveOperator. REQUIRE_SORT reflects (_priorityQueue != null). + assertFalse(operator.copyStatMaps().getBoolean(SortOperator.StatKey.REQUIRE_SORT), + "SortOperator must skip re-sorting when input is a SortedMailboxReceiveOperator"); + } + } + private SortOperator getOperator(DataSchema schema, List collations, int fetch, int offset) { return new SortOperator(OperatorTestUtil.getTracingContext(), _input, new SortNode(-1, schema, PlanNode.NodeHint.EMPTY, List.of(), collations, fetch, offset)); 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 59aeccd6d5d5..373cf13f44a1 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 @@ -18,8 +18,17 @@ */ package org.apache.pinot.query.runtime.operator; +import com.fasterxml.jackson.databind.JsonNode; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Random; import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.calcite.rel.RelDistribution; @@ -39,8 +48,12 @@ import org.apache.pinot.query.routing.WorkerMetadata; import org.apache.pinot.query.runtime.blocks.ErrorMseBlock; import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.operator.utils.SortUtils; +import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.segment.spi.memory.DataBuffer; import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.utils.CommonConstants; import org.mockito.Mock; import org.mockito.MockitoAnnotations; import org.testng.annotations.AfterMethod; @@ -49,11 +62,19 @@ import org.testng.annotations.Test; 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.Mockito.any; +import static org.mockito.Mockito.atLeast; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; @@ -252,6 +273,874 @@ public void shouldReceiveMailboxFromTwoServersWithCollationKeyTwoColumns() { } } + @Test + public void shouldMergeFromTwoServersInOrder() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1}, new Object[]{3, 3}, new Object[]{5, 5}), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2}, new Object[]{4, 4}, new Object[]{6, 6}), + OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED)) { + List rows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + assertEquals(rows.size(), 6); + for (int i = 0; i < 6; i++) { + assertEquals(rows.get(i)[0], i + 1); + } + assertTrue(operator.nextBlock().isSuccess()); + } + } + + @Test + public void shouldMergeWithTiedKeysPreservingMultiset() { + // Rows with equal collation keys (col0) but distinct payloads (col1). The k-way merge (PriorityQueue) is not a + // stable sort, so tie order may differ from the accumulate-then-sort path; assert what IS guaranteed: the output + // is globally non-decreasing by the collation key and the row multiset is preserved. + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 10}, new Object[]{1, 11}, new Object[]{2, 12}), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 20}, new Object[]{2, 21}, new Object[]{2, 22}), + OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED)) { + List rows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + assertEquals(rows.size(), 6); + for (int i = 1; i < rows.size(); i++) { + assertTrue((int) rows.get(i)[0] >= (int) rows.get(i - 1)[0], "merge output must be non-decreasing by key"); + } + // Multiset equality, independent of tie order: sort both sides by (col0, col1) and compare element-wise. + List actualSorted = new ArrayList<>(rows); + actualSorted.sort((x, y) -> (int) x[0] != (int) y[0] ? (int) x[0] - (int) y[0] : (int) x[1] - (int) y[1]); + List expected = List.of(new Object[]{1, 10}, new Object[]{1, 11}, new Object[]{1, 20}, + new Object[]{2, 12}, new Object[]{2, 21}, new Object[]{2, 22}); + assertRowsEqual(actualSorted, expected); + assertTrue(operator.nextBlock().isSuccess()); + } + } + + @Test + public void shouldMergeWithStaggeredEos() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2}), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1}, new Object[]{3, 3}, new Object[]{4, 4}), + OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED)) { + List rows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + assertEquals(rows.size(), 4); + for (int i = 0; i < 4; i++) { + assertEquals(rows.get(i)[0], i + 1); + } + assertTrue(operator.nextBlock().isSuccess()); + } + } + + @Test + public void shouldMergeInterleavedNotReady() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + // Capture the reader BEFORE constructing the operator (the ctor registers it). Single-consumer-thread test, so + // reading the captured reference from a later poll Answer needs no synchronization. + ReceivingMailbox.Reader[] reader1 = new ReceivingMailbox.Reader[1]; + doAnswer(inv -> { + reader1[0] = inv.getArgument(0); + return null; + }).when(_mailbox1).registeredReader(any()); + // data -> (self-wake + null) -> data -> EOS. The self-wake makes the blocking poll return immediately. + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1})) + .thenAnswer(inv -> { + reader1[0].blockReadyToRead(); + return null; + }) + .thenReturn(OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2})) + .thenReturn(OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadata1, RelDistribution.Type.SINGLETON, + DATA_SCHEMA, FIELD_COLLATIONS, System.currentTimeMillis() + 30_000L, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"))) { + List all = new ArrayList<>(); + MseBlock block = operator.nextBlock(); + while (block.isData()) { + all.addAll(((MseBlock.Data) block).asRowHeap().getRows()); + block = operator.nextBlock(); + } + assertTrue(block.isSuccess()); + assertEquals(all.size(), 2); + assertEquals(all.get(0)[0], 1); + assertEquals(all.get(1)[0], 2); + } + } + + /** + * Regression for the streaming k-way merge pipeline deadlock: when the merge needs the next row from one stream that + * is momentarily empty, it must drain the OTHER ready stream (relieving that sender's backpressure) instead of + * head-of-line blocking on the starved stream. Here mailbox1 is starved (returns null) and only becomes ready when + * mailbox2 is polled during draining; a merge that blocks on mailbox1 alone never polls mailbox2 and times out. + */ + @Test + public void shouldDrainReadySiblingWhenOtherStreamStarved() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + // Capture mailbox1's reader (registered in the operator ctor) so mailbox2's poll can wake it. + ReceivingMailbox.Reader[] reader1 = new ReceivingMailbox.Reader[1]; + doAnswer(inv -> { + reader1[0] = inv.getArgument(0); + return null; + }).when(_mailbox1).registeredReader(any()); + // mailbox1: {1}, then starved (null), then {3}, then EOS. The null forces the merge to look elsewhere. + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1})) + .thenReturn(null) + .thenReturn(OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{3, 3})) + .thenReturn(OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + // mailbox2: {2}; polling for the next block ({4}) wakes mailbox1 (models the sibling drain unblocking the starved + // sender); then EOS. + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2})) + .thenAnswer(inv -> { + reader1[0].blockReadyToRead(); + return OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{4, 4}); + }) + .thenReturn(OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, FIELD_COLLATIONS, + System.currentTimeMillis() + 30_000L, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"))) { + List all = new ArrayList<>(); + MseBlock block = operator.nextBlock(); + while (block.isData()) { + all.addAll(((MseBlock.Data) block).asRowHeap().getRows()); + block = operator.nextBlock(); + } + assertTrue(block.isSuccess()); + assertEquals(all.size(), 4); + for (int i = 0; i < 4; i++) { + assertEquals(all.get(i)[0], i + 1); + } + // Prove the sibling drain actually ran (not just that the output happens to be correct): mailbox2 must have been + // polled past its first block (drained for row {4,4} and its EOS) while mailbox1 was starved, and mailbox1 must + // have been re-polled after the starved (null) response instead of parking on it forever. + verify(_mailbox2, atLeast(3)).poll(); + verify(_mailbox1, atLeast(3)).poll(); + } + } + + /** + * Regression for error propagation during the sibling drain: while refilling a starved stream, the merge polls + * sibling streams non-blocking; if a sibling yields an error mid-drain, that error must short-circuit the merge + * immediately rather than being swallowed or causing a hang while the starved stream is still awaited. + */ + @Test + public void shouldPropagateErrorFromSiblingDuringDrain() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + // mailbox1: {1}, then starved forever (null) -- it never itself produces the error or EOS. + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1})) + .thenReturn(null); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + // mailbox2: {2}, then an error -- surfaced while mailbox2 is drained as a sibling of the starved mailbox1. + String errorMessage = "SIBLING ERROR"; + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2})) + .thenReturn(OperatorTestUtil.errorWithEmptyStats(new RuntimeException(errorMessage))); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, FIELD_COLLATIONS, + System.currentTimeMillis() + 30_000L, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"))) { + MseBlock block = operator.nextBlock(); + assertTrue(block.isError()); + assertTrue(((ErrorMseBlock) block).getErrorMessages().get(QueryErrorCode.UNKNOWN).contains(errorMessage)); + } + } + + /** + * Regression for the early-termination drain path (drainToEos), which was rewritten to be cooperative for the same + * reason as the merge itself: it must not head-of-line block on one handle while a sibling still has buffered data, + * or the sibling's sender backpressures and the pipeline deadlocks. mailbox1 is starved (null) until mailbox2 is + * polled during the round-robin drain; a drain that blocks on mailbox1 alone never polls mailbox2 and times out. + */ + @Test + public void shouldDrainSiblingsCooperativelyOnEarlyTerminate() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + ReceivingMailbox.Reader[] reader1 = new ReceivingMailbox.Reader[1]; + doAnswer(inv -> { + reader1[0] = inv.getArgument(0); + return null; + }).when(_mailbox1).registeredReader(any()); + // mailbox1: starved (null), then EOS. It only makes progress once woken by mailbox2's drain. + when(_mailbox1.poll()).thenReturn(null).thenReturn(OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + // mailbox2: a data block (discarded on early-termination) whose poll wakes mailbox1, then EOS. + when(_mailbox2.poll()) + .thenAnswer(inv -> { + reader1[0].blockReadyToRead(); + return OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1}); + }) + .thenReturn(OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, FIELD_COLLATIONS, + System.currentTimeMillis() + 10_000L, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"))) { + operator.earlyTerminate(); + assertTrue(operator.nextBlock().isSuccess()); + verify(_mailbox1).earlyTerminate(); + verify(_mailbox2).earlyTerminate(); + // Both mailboxes were drained to EOS cooperatively (round-robin), not head-of-line blocked on the starved one. + verify(_mailbox1, atLeast(2)).poll(); + verify(_mailbox2, atLeast(2)).poll(); + } + } + + /** + * Regression proving the per-stream backlog is served in FIFO order. While mailbox1 (the current min source) is + * starved, the greedy sibling drain buffers ALL of mailbox2's ready blocks ({3},{5}) into mailbox2's backlog at once; + * mailbox2 also reaches EOS during that drain. The merge must then serve that multi-block backlog oldest-first so the + * global output stays sorted (1,2,3,4,5). A LIFO backlog would emit 5 before 3. + */ + @Test + public void shouldServeMultiBlockBacklogInOrderWhenStreamStarved() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + // mailbox1 (min source): {1}, starved (null), {4}, EOS. + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1})) + .thenReturn(null) + .thenReturn(OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{4, 4})) + .thenReturn(OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + // mailbox2 (fast sibling): {2}, then {3},{5} buffered together during the drain, then EOS (also during the drain). + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2}), + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{3, 3}), + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{5, 5}), + OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, FIELD_COLLATIONS, + System.currentTimeMillis() + 10_000L, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"))) { + List all = new ArrayList<>(); + MseBlock block = operator.nextBlock(); + while (block.isData()) { + all.addAll(((MseBlock.Data) block).asRowHeap().getRows()); + block = operator.nextBlock(); + } + assertTrue(block.isSuccess()); + assertEquals(all.size(), 5); + for (int i = 0; i < 5; i++) { + assertEquals(all.get(i)[0], i + 1); + } + } + } + + @Test + public void shouldMergeBoundedMultiBlock() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1}, new Object[]{3, 3}, new Object[]{5, 5}), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2}, new Object[]{4, 4}, new Object[]{6, 6}), + OperatorTestUtil.eosWithEmptyStats()); + SortUtils.SortComparator comparator = new SortUtils.SortComparator(FIELD_COLLATIONS, false); + Map hints = Map.of( + CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true", + CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE, "2"); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, FIELD_COLLATIONS, Long.MAX_VALUE, hints)) { + Object[] prevLast = null; + int dataBlocks = 0; + List all = new ArrayList<>(); + MseBlock block = operator.nextBlock(); + while (block.isData()) { + List rows = ((MseBlock.Data) block).asRowHeap().getRows(); + assertEquals(rows.size(), 2); + if (prevLast != null) { + assertTrue(comparator.compare(prevLast, rows.get(0)) <= 0); + } + prevLast = rows.get(rows.size() - 1); + all.addAll(rows); + dataBlocks++; + block = operator.nextBlock(); + } + assertEquals(dataBlocks, 3); + for (int i = 0; i < 6; i++) { + assertEquals(all.get(i)[0], i + 1); + } + assertTrue(block.isSuccess()); + } + } + + @Test + public void shouldMergeWithDescAndNullDirection() { + DataSchema dataSchema = + new DataSchema(new String[]{"col1", "col2", "col3"}, new DataSchema.ColumnDataType[]{INT, INT, STRING}); + List collations = List.of(new RelFieldCollation(2, Direction.DESCENDING, NullDirection.FIRST), + new RelFieldCollation(0, Direction.ASCENDING, NullDirection.LAST)); + SortUtils.SortComparator comparator = new SortUtils.SortComparator(collations, false); + Object[] row1 = new Object[]{3, 3, "queen"}; + Object[] row2 = new Object[]{1, 1, "pink floyd"}; + Object[] row3 = new Object[]{4, 2, "pink floyd"}; + Object[] row4 = new Object[]{2, 4, "aerosmith"}; + Object[] row5 = new Object[]{-1, 95, null}; + List mb1 = new ArrayList<>(List.of(row1, row2)); + mb1.sort(comparator); + List mb2 = new ArrayList<>(List.of(row3, row4, row5)); + mb2.sort(comparator); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(dataSchema, mb1.toArray(new Object[0][])), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(dataSchema, mb2.toArray(new Object[0][])), + OperatorTestUtil.eosWithEmptyStats()); + List reference = new ArrayList<>(List.of(row1, row2, row3, row4, row5)); + reference.sort(comparator); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, dataSchema, collations, Long.MAX_VALUE, Map.of( + CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"))) { + assertRowsEqual(((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(), reference); + assertTrue(operator.nextBlock().isSuccess()); + } + } + + @Test + public void shouldReturnErrorWhenMergeMailboxErrors() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + String errorMessage = "TEST ERROR"; + when(_mailbox1.poll()).thenReturn(OperatorTestUtil.errorWithEmptyStats(new RuntimeException(errorMessage))); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{3, 3}), + OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED)) { + MseBlock block = operator.nextBlock(); + assertTrue(block.isError()); + assertTrue(((ErrorMseBlock) block).getErrorMessages().get(QueryErrorCode.UNKNOWN).contains(errorMessage)); + } + } + + @Test + public void shouldTimeoutInMergeMode() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadata1, RelDistribution.Type.SINGLETON, + DATA_SCHEMA, FIELD_COLLATIONS, System.currentTimeMillis() + 100L, Map.of( + CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"))) { + MseBlock block = operator.nextBlock(); + assertTrue(block.isError()); + assertTrue(((ErrorMseBlock) block).getErrorMessages().containsKey(QueryErrorCode.EXECUTION_TIMEOUT)); + } + } + + @Test + public void shouldReturnSuccessOnEarlyTerminateInMergeMode() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn(OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadata1, RelDistribution.Type.SINGLETON)) { + operator.earlyTerminate(); + assertTrue(operator.nextBlock().isSuccess()); + verify(_mailbox1).earlyTerminate(); + } + } + + /** + * Early termination in production ({@code SortOperator} hitting its LIMIT) arrives mid-merge: the heap is + * live, mailboxes still hold undelivered blocks, and no handle has reached EOS. That is the state {@code drainToEos} + * was rewritten for — it must round-robin every remaining handle to EOS rather than head-of-line block on one while + * a sibling's sender sits on a full mailbox. Terminating before the first {@code nextBlock()} (as the other early + * termination tests do) leaves that path unexercised, because every handle is already exhausted on the first pass. + */ + @Test + public void shouldDrainBothMailboxesWhenEarlyTerminatedMidMerge() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1}), + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{3, 3}), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2}), + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{4, 4}), + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{6, 6}), + OperatorTestUtil.eosWithEmptyStats()); + // Block size 1 so the merge is primed and mid-flight (heap live, both mailboxes unexhausted) after two calls. + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, FIELD_COLLATIONS, Long.MAX_VALUE, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true", + CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE, "1"))) { + MseBlock first = operator.nextBlock(); + assertTrue(first.isData()); + assertRowsEqual(((MseBlock.Data) first).asRowHeap().getRows(), List.of(new Object[]{1, 1})); + MseBlock second = operator.nextBlock(); + assertTrue(second.isData()); + assertRowsEqual(((MseBlock.Data) second).asRowHeap().getRows(), List.of(new Object[]{2, 2})); + + operator.earlyTerminate(); + assertTrue(operator.nextBlock().isSuccess()); + // Both senders must have been driven to EOS, not just the one the merge happened to be waiting on. + verify(_mailbox1).earlyTerminate(); + verify(_mailbox2).earlyTerminate(); + verify(_mailbox1, atLeast(3)).poll(); + verify(_mailbox2, atLeast(4)).poll(); + } + } + + /** + * The merge's whole correctness rests on each mailbox stream already being sorted, and nothing downstream re-sorts. + * If a sender violates that (a plan shape the fragmenter gate should have rejected, or a leaf concatenating two + * independently sorted runs), the merge must fail loudly rather than silently emit misordered rows. + */ + @Test + public void shouldFailFastOnOutOfOrderSenderRows() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + // A single mailbox whose second block sorts before its first: two independently sorted runs concatenated. + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{5, 5}, new Object[]{9, 9}), + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1}, new Object[]{2, 2}), + OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadata1, RelDistribution.Type.SINGLETON, + DATA_SCHEMA, FIELD_COLLATIONS, Long.MAX_VALUE, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"))) { + // MultiStageOperator.nextBlock() converts the thrown IllegalStateException into an error block, so the query + // fails with a diagnosable message instead of returning silently misordered rows. + MseBlock block = operator.nextBlock(); + assertTrue(block.isError()); + assertTrue(((ErrorMseBlock) block).getErrorMessages().values().stream() + .anyMatch(message -> message.contains("out-of-order")), + "Expected an out-of-order error, got: " + ((ErrorMseBlock) block).getErrorMessages()); + } + } + + @Test + public void shouldPreserveStatsInMergeMode() + throws Exception { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + List stats1 = new MultiStageQueryStats.Builder(1).addLast( + open -> open.addLastOperator(MultiStageOperator.Type.MAILBOX_SEND, + new StatMap<>(MailboxSendOperator.StatKey.class)) + .addLastOperator(MultiStageOperator.Type.LEAF, new StatMap<>(LeafOperator.StatKey.class)) + .close()).build().serialize(); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1}), + OperatorTestUtil.eosWithStats(stats1)); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + List stats2 = new MultiStageQueryStats.Builder(1).addLast( + open -> open.addLastOperator(MultiStageOperator.Type.MAILBOX_SEND, + new StatMap<>(MailboxSendOperator.StatKey.class)) + .addLastOperator(MultiStageOperator.Type.LEAF, new StatMap<>(LeafOperator.StatKey.class)) + .close()).build().serialize(); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2}), + OperatorTestUtil.eosWithStats(stats2)); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED)) { + while (!operator.nextBlock().isEos()) { + // drain + } + MultiStageQueryStats stats = operator.calculateStats(); + assertNotNull(stats); + // Both mailboxes carry identically-shaped upstream stats, so they merge cleanly into stage 1. + MultiStageQueryStats.StageStats.Closed upstreamStats = stats.getUpstreamStageStats(1); + assertNotNull(upstreamStats); + } + } + + @Test + public void shouldMatchFullSortParityWithMerge() { + Random random = new Random(42L); + SortUtils.SortComparator comparator = new SortUtils.SortComparator(FIELD_COLLATIONS, false); + // Use unique col0 (the only collation key) across both lists so the comparator has no ties; with ties the merge + // order and a single List.sort order of equal-key rows are both valid but need not match. + List keys = new ArrayList<>(); + for (int i = 0; i < 16; i++) { + keys.add(i); + } + Collections.shuffle(keys, random); + List mb1 = sortedRowsFromKeys(keys.subList(0, 7), random, comparator); + List mb2 = sortedRowsFromKeys(keys.subList(7, 16), random, comparator); + + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, mb1.toArray(new Object[0][])), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, mb2.toArray(new Object[0][])), + OperatorTestUtil.eosWithEmptyStats()); + List mergeRows = new ArrayList<>(); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED)) { + MseBlock block = operator.nextBlock(); + while (block.isData()) { + mergeRows.addAll(((MseBlock.Data) block).asRowHeap().getRows()); + block = operator.nextBlock(); + } + assertTrue(block.isSuccess()); + } + + // Re-stub the same mailboxes with the same data and run the default accumulate-then-sort path. + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, mb1.toArray(new Object[0][])), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, mb2.toArray(new Object[0][])), + OperatorTestUtil.eosWithEmptyStats()); + List defaultRows = new ArrayList<>(); + try (SortedMailboxReceiveOperator operator = getOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED)) { + MseBlock block = operator.nextBlock(); + while (block.isData()) { + defaultRows.addAll(((MseBlock.Data) block).asRowHeap().getRows()); + block = operator.nextBlock(); + } + assertTrue(block.isSuccess()); + } + assertRowsEqual(mergeRows, defaultRows); + } + + @Test + public void shouldMatchMultiColumnParityWithMerge() { + // Schema: key1 (STRING), key2 (INT), key3 (INT), payload (LONG) + DataSchema schema = new DataSchema( + new String[]{"key1", "key2", "key3", "payload"}, + new DataSchema.ColumnDataType[]{STRING, INT, INT, LONG}); + + // Collation: key1 DESC NULLS_FIRST, key2 ASC NULLS_LAST, key3 DESC NULLS_LAST + List collations = List.of( + new RelFieldCollation(0, Direction.DESCENDING, NullDirection.FIRST), + new RelFieldCollation(1, Direction.ASCENDING, NullDirection.LAST), + new RelFieldCollation(2, Direction.DESCENDING, NullDirection.LAST)); + SortUtils.SortComparator comparator = new SortUtils.SortComparator(collations, false); + + // ~20 rows with: tied key1 values, tied key1+key2 pairs, nulls in key1, + // all (key1, key2, key3) composite keys unique for deterministic ordering. + List allRows = List.of( + new Object[]{null, 5, 10, 100L}, // null key1 — exercises NULLS_FIRST in DESC + new Object[]{null, 5, 20, 101L}, // null key1 — different key3 breaks tie + new Object[]{null, 10, 30, 102L}, // null key1, different key2 + new Object[]{"delta", 1, 50, 103L}, + new Object[]{"delta", 1, 40, 104L}, // tied key1+key2, key3 breaks tie + new Object[]{"delta", 2, 60, 105L}, + new Object[]{"delta", 3, 70, 106L}, + new Object[]{"charlie", 1, 10, 107L}, + new Object[]{"charlie", 1, 20, 108L}, + new Object[]{"charlie", 2, 30, 109L}, + new Object[]{"bravo", 5, 15, 110L}, + new Object[]{"bravo", 5, 25, 111L}, + new Object[]{"bravo", 10, 35, 112L}, + new Object[]{"alpha", 1, 100, 113L}, + new Object[]{"alpha", 1, 200, 114L}, + new Object[]{"alpha", 2, 50, 115L}, + new Object[]{"alpha", 3, 10, 116L}, + new Object[]{"alpha", 3, 20, 117L}, + new Object[]{"alpha", 4, 5, 118L}, + new Object[]{"alpha", 4, 15, 119L} + ); + + // Shuffle deterministically and split across 2 mailboxes + Random random = new Random(99L); + List shuffled = new ArrayList<>(allRows); + Collections.shuffle(shuffled, random); + List mb1 = new ArrayList<>(shuffled.subList(0, 10)); + List mb2 = new ArrayList<>(shuffled.subList(10, 20)); + mb1.sort(comparator); + mb2.sort(comparator); + + // --- Merge path --- + RelDistribution.Type distributionType = RelDistribution.Type.HASH_DISTRIBUTED; + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(schema, mb1.toArray(new Object[0][])), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(schema, mb2.toArray(new Object[0][])), + OperatorTestUtil.eosWithEmptyStats()); + List mergeRows = new ArrayList<>(); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + distributionType, schema, collations, Long.MAX_VALUE, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"))) { + MseBlock block = operator.nextBlock(); + while (block.isData()) { + mergeRows.addAll(((MseBlock.Data) block).asRowHeap().getRows()); + block = operator.nextBlock(); + } + assertTrue(block.isSuccess()); + } + + // --- Accumulate path (re-stub same data) --- + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(schema, mb1.toArray(new Object[0][])), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(schema, mb2.toArray(new Object[0][])), + OperatorTestUtil.eosWithEmptyStats()); + List defaultRows = new ArrayList<>(); + try (SortedMailboxReceiveOperator operator = getOperator(_stageMetadataBoth, + distributionType, schema, collations, Long.MAX_VALUE)) { + MseBlock block = operator.nextBlock(); + while (block.isData()) { + defaultRows.addAll(((MseBlock.Data) block).asRowHeap().getRows()); + block = operator.nextBlock(); + } + assertTrue(block.isSuccess()); + } + + assertRowsEqual(mergeRows, defaultRows); + } + + @Test + public void shouldAccumulateWhenHintTrueButNotSortedOnSender() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{3, 3}, new Object[]{1, 1}), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2}), + OperatorTestUtil.eosWithEmptyStats()); + // hint=true but isSortedOnSender=false => accumulate-then-sort (AND gate blocks the unsafe arm). + // blockSize=2 proves accumulate: merge would emit 2 blocks of 2, but accumulate returns all 3 rows in one block. + try (SortedMailboxReceiveOperator operator = getSenderSortedOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, false, "2", + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true", + CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE, "2"))) { + assertRowsEqual(((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(), + List.of(new Object[]{1, 1}, new Object[]{2, 2}, new Object[]{3, 3})); + assertTrue(operator.nextBlock().isSuccess()); + } + } + + @Test + public void shouldUseAccumulatePathWhenHintFalse() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{3, 3}, new Object[]{1, 1}), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2}), + OperatorTestUtil.eosWithEmptyStats()); + try (SortedMailboxReceiveOperator operator = getMergeOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, FIELD_COLLATIONS, Long.MAX_VALUE, Map.of( + CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "false"))) { + // Accumulate-then-sort path returns a single globally sorted data block. + assertRowsEqual(((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(), + List.of(new Object[]{1, 1}, new Object[]{2, 2}, new Object[]{3, 3})); + assertTrue(operator.nextBlock().isSuccess()); + } + } + + @Test + public void shouldAccumulateWhenHintUnsetEvenIfSortedOnSender() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1}, new Object[]{3, 3}), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2}, new Object[]{4, 4}), + OperatorTestUtil.eosWithEmptyStats()); + // No hint in opChainMetadata + isSortedOnSender()==true => accumulate-then-sort (AND gate requires explicit hint). + // blockSize=2 proves accumulate: merge would emit 2 blocks of 2, but accumulate returns all 4 rows in one block. + try (SortedMailboxReceiveOperator operator = getSenderSortedOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, true, "2")) { + assertRowsEqual(((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(), + List.of(new Object[]{1, 1}, new Object[]{2, 2}, new Object[]{3, 3}, new Object[]{4, 4})); + assertTrue(operator.nextBlock().isSuccess()); + } + } + + @Test + public void shouldAccumulateWhenHintUnsetAndNotSortedOnSender() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{3, 3}, new Object[]{1, 1}), + OperatorTestUtil.eosWithEmptyStats()); + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2); + when(_mailbox2.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{2, 2}), + OperatorTestUtil.eosWithEmptyStats()); + // No hint + isSortedOnSender()==false => accumulate-then-sort path: one globally sorted block. + try (SortedMailboxReceiveOperator operator = getSenderSortedOperator(_stageMetadataBoth, + RelDistribution.Type.HASH_DISTRIBUTED, false, null)) { + assertRowsEqual(((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(), + List.of(new Object[]{1, 1}, new Object[]{2, 2}, new Object[]{3, 3})); + assertTrue(operator.nextBlock().isSuccess()); + } + } + + @Test + public void shouldReportKWayMergeStatWhenMergeUsed() { + assertKWayMergeStat(true, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"), true); + } + + @Test + public void shouldNotReportKWayMergeStatWhenSenderNotSorted() { + assertKWayMergeStat(false, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true"), false); + } + + @Test + public void shouldNotReportKWayMergeStatWhenOptionOff() { + assertKWayMergeStat(true, Map.of(), false); + } + + /** + * The stat travels to the broker through {@link StatMap#serialize}/{@link StatMap#deserialize}, which encode keys by + * ordinal. Asserting the round trip (and the ordinal position) here means a future reordering or an accidental + * change to the presence-based boolean encoding fails loudly rather than silently breaking cluster diagnosis and + * mixed-version decoding. + */ + @Test + public void shouldRoundTripKWayMergeStatThroughSerialization() + throws Exception { + BaseMailboxReceiveOperator.StatKey[] keys = BaseMailboxReceiveOperator.StatKey.values(); + assertEquals(keys[keys.length - 1], BaseMailboxReceiveOperator.StatKey.K_WAY_MERGE_USED, + "K_WAY_MERGE_USED must stay the last key: StatMap serializes by ordinal, so keys may only be appended"); + + StatMap merged = + new StatMap<>(BaseMailboxReceiveOperator.StatKey.class); + merged.merge(BaseMailboxReceiveOperator.StatKey.K_WAY_MERGE_USED, true); + assertTrue(roundTrip(merged).getBoolean(BaseMailboxReceiveOperator.StatKey.K_WAY_MERGE_USED)); + + // On the accumulate-then-sort path the key must not be written at all, which is what keeps the new ordinal off + // the wire for peers that predate it. + StatMap notMerged = + new StatMap<>(BaseMailboxReceiveOperator.StatKey.class); + notMerged.merge(BaseMailboxReceiveOperator.StatKey.K_WAY_MERGE_USED, false); + notMerged.merge(BaseMailboxReceiveOperator.StatKey.FAN_IN, 1); + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (DataOutputStream out = new DataOutputStream(bytes)) { + notMerged.serialize(out); + } + // Serialized form is: key count, then one (ordinal, value) pair per present key. Only FAN_IN is present. + assertEquals(bytes.toByteArray()[0], (byte) 1); + assertFalse(roundTrip(notMerged).getBoolean(BaseMailboxReceiveOperator.StatKey.K_WAY_MERGE_USED)); + } + + private static StatMap roundTrip( + StatMap statMap) + throws Exception { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (DataOutputStream out = new DataOutputStream(bytes)) { + statMap.serialize(out); + } + try (DataInputStream in = new DataInputStream(new ByteArrayInputStream(bytes.toByteArray()))) { + return StatMap.deserialize(in, BaseMailboxReceiveOperator.StatKey.class); + } + } + + /** + * Drives a single-mailbox receive to completion and asserts the {@code K_WAY_MERGE_USED} stat, both on the stat map + * and in the JSON that is rendered into the query response {@code stageStats}. Both paths must return the same rows, + * so the stat is the only thing that distinguishes them. + */ + private void assertKWayMergeStat(boolean sortedOnSender, Map opChainMetadata, boolean expected) { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + when(_mailbox1.poll()).thenReturn( + OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1}, new Object[]{2, 2}), + OperatorTestUtil.eosWithEmptyStats()); + List rows = new ArrayList<>(); + try (SortedMailboxReceiveOperator operator = getSenderSortedOperator(_stageMetadata1, + RelDistribution.Type.SINGLETON, sortedOnSender, null, opChainMetadata)) { + MseBlock block = operator.nextBlock(); + while (block.isData()) { + rows.addAll(((MseBlock.Data) block).asRowHeap().getRows()); + block = operator.nextBlock(); + } + assertTrue(block.isSuccess()); + assertRowsEqual(rows, List.of(new Object[]{1, 1}, new Object[]{2, 2})); + + StatMap statMap = operator.copyStatMaps(); + assertEquals(statMap.getBoolean(BaseMailboxReceiveOperator.StatKey.K_WAY_MERGE_USED), expected); + // The stat must survive into the response stageStats, which is rendered from StatMap.asJson(). Reporting is + // presence-based: rendered as true on the merge path, and absent (not "false") otherwise. + JsonNode kWayMergeUsed = statMap.asJson().get(BaseMailboxReceiveOperator.StatKey.K_WAY_MERGE_USED.getStatName()); + if (expected) { + assertNotNull(kWayMergeUsed, "kWayMergeUsed must be present in the stageStats JSON when the merge is used"); + assertTrue(kWayMergeUsed.booleanValue()); + } else { + assertNull(kWayMergeUsed, + "kWayMergeUsed must be absent from the stageStats JSON on the accumulate-then-sort path"); + } + } + } + + private SortedMailboxReceiveOperator getSenderSortedOperator(StageMetadata stageMetadata, + RelDistribution.Type distributionType, boolean sortedOnSender, String blockSize) { + Map opChainMetadata = blockSize == null ? Map.of() + : Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE, blockSize); + return getSenderSortedOperator(stageMetadata, distributionType, sortedOnSender, blockSize, opChainMetadata); + } + + private SortedMailboxReceiveOperator getSenderSortedOperator(StageMetadata stageMetadata, + RelDistribution.Type distributionType, boolean sortedOnSender, String blockSize, + Map opChainMetadata) { + if (blockSize != null && !opChainMetadata + .containsKey(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE)) { + Map merged = new HashMap<>(opChainMetadata); + merged.put(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE, blockSize); + opChainMetadata = merged; + } + OpChainExecutionContext context = + OperatorTestUtil.getOpChainContext(_mailboxService, Long.MAX_VALUE, stageMetadata, opChainMetadata); + MailboxReceiveNode node = mock(MailboxReceiveNode.class); + when(node.getDistributionType()).thenReturn(distributionType); + when(node.getSenderStageId()).thenReturn(1); + when(node.getDataSchema()).thenReturn(DATA_SCHEMA); + when(node.getCollations()).thenReturn(FIELD_COLLATIONS); + when(node.isSortedOnSender()).thenReturn(sortedOnSender); + return new SortedMailboxReceiveOperator(context, node); + } + + private void assertRowsEqual(List actual, List expected) { + assertEquals(actual.size(), expected.size()); + for (int i = 0; i < actual.size(); i++) { + assertEquals(actual.get(i), expected.get(i)); + } + } + + private List sortedRowsFromKeys(List keys, Random random, SortUtils.SortComparator comparator) { + List rows = new ArrayList<>(keys.size()); + for (int key : keys) { + rows.add(new Object[]{key, random.nextInt(50)}); + } + rows.sort(comparator); + return rows; + } + + private SortedMailboxReceiveOperator getMergeOperator(StageMetadata stageMetadata, + RelDistribution.Type distributionType, DataSchema resultSchema, List collations, + long deadlineMs, Map opChainMetadata) { + OpChainExecutionContext context = + OperatorTestUtil.getOpChainContext(_mailboxService, deadlineMs, stageMetadata, opChainMetadata); + MailboxReceiveNode node = mock(MailboxReceiveNode.class); + when(node.getDistributionType()).thenReturn(distributionType); + when(node.getSenderStageId()).thenReturn(1); + when(node.getDataSchema()).thenReturn(resultSchema); + when(node.getCollations()).thenReturn(collations); + when(node.isSortedOnSender()).thenReturn(true); + return new SortedMailboxReceiveOperator(context, node); + } + + private SortedMailboxReceiveOperator getMergeOperator(StageMetadata stageMetadata, + RelDistribution.Type distributionType) { + return getMergeOperator(stageMetadata, distributionType, DATA_SCHEMA, FIELD_COLLATIONS, Long.MAX_VALUE, + Map.of(CommonConstants.Broker.Request.QueryOptionKey.STREAMING_SORTED_MAILBOX_RECEIVE, "true")); + } + private SortedMailboxReceiveOperator getOperator(StageMetadata stageMetadata, RelDistribution.Type distributionType, DataSchema resultSchema, List 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/utils/BlockingMultiStreamConsumerTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumerTest.java new file mode 100644 index 000000000000..e7dabe3e68f8 --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumerTest.java @@ -0,0 +1,357 @@ +/** + * 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.utils; + +import java.io.IOException; +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.List; +import javax.annotation.Nullable; +import org.apache.pinot.common.datatable.StatMap; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.mailbox.ReceivingMailbox; +import org.apache.pinot.query.runtime.blocks.ErrorMseBlock; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.operator.LeafOperator; +import org.apache.pinot.query.runtime.operator.MailboxSendOperator; +import org.apache.pinot.query.runtime.operator.MultiStageOperator.Type; +import org.apache.pinot.query.runtime.operator.OpChainId; +import org.apache.pinot.query.runtime.operator.OperatorTestUtil; +import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.segment.spi.memory.DataBuffer; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.testng.annotations.Test; + +import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.INT; +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.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +/** + * Unit tests for the per-stream pull mode of {@link BlockingMultiStreamConsumer} (the {@link + * BlockingMultiStreamConsumer.StreamHandle} / {@link BlockingMultiStreamConsumer#streamHandles()} / + * {@code StreamHandle.readBlocking()} API plus the single-mode guard). The round-robin path is covered separately by + * {@code MailboxReceiveOperatorTest}. + */ +public class BlockingMultiStreamConsumerTest { + private static final DataSchema DATA_SCHEMA = + new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT}); + private static final int RECEIVER_STAGE_ID = 0; + private static final int SENDER_STAGE_ID = 1; + private static final long FAR_FUTURE = System.currentTimeMillis() + 60_000L; + + private static BlockingMultiStreamConsumer.OfMseBlock newConsumer(long deadlineMs, List streams) { + OpChainExecutionContext context = mock(OpChainExecutionContext.class); + when(context.getId()).thenReturn(mock(OpChainId.class)); + when(context.getStageId()).thenReturn(RECEIVER_STAGE_ID); + when(context.getPassiveDeadlineMs()).thenReturn(deadlineMs); + return new BlockingMultiStreamConsumer.OfMseBlock(context, streams, SENDER_STAGE_ID); + } + + /** + * Serializes a {@link MultiStageQueryStats} with a fixed leaf-side operator shape, so the consumer can merge it on + * success EOS. Two EOS blocks with the same shape merge cleanly (see {@code + * MailboxReceiveOperatorTest#differentUpstreamStatsProduceEmptyStats} for the differing-shape case). + */ + private static List leafStats() + throws IOException { + return new MultiStageQueryStats.Builder(SENDER_STAGE_ID).addLast( + open -> open.addLastOperator(Type.MAILBOX_SEND, new StatMap<>(MailboxSendOperator.StatKey.class)) + .addLastOperator(Type.LEAF, new StatMap<>(LeafOperator.StatKey.class)) + .close()).build().serialize(); + } + + @Test + public void dataThenEosPerHandleKeepsStats() + throws IOException { + Object[] row0 = new Object[]{0, 0}; + Object[] row1 = new Object[]{1, 1}; + FakeStream s0 = new FakeStream("s0"); + FakeStream s1 = new FakeStream("s1"); + s0.enqueue(OperatorTestUtil.blockWithStats(DATA_SCHEMA, row0)); + s0.enqueue(OperatorTestUtil.eosWithStats(leafStats())); + s1.enqueue(OperatorTestUtil.blockWithStats(DATA_SCHEMA, row1)); + s1.enqueue(OperatorTestUtil.eosWithStats(leafStats())); + + BlockingMultiStreamConsumer.OfMseBlock consumer = newConsumer(FAR_FUTURE, List.of(s0, s1)); + List> handles = + consumer.streamHandles(); + assertEquals(handles.size(), 2); + // streamHandles() is idempotent. + assertSame(consumer.streamHandles(), handles); + + BlockingMultiStreamConsumer.StreamHandle h0 = handles.get(0); + BlockingMultiStreamConsumer.StreamHandle h1 = handles.get(1); + assertEquals(h0.getId(), "s0"); + assertEquals(h1.getId(), "s1"); + + // Each handle reads its own data block first, independent of the other. + assertEquals(((MseBlock.Data) h0.readBlocking().getBlock()).asRowHeap().getRows().get(0), row0); + assertFalse(h0.isExhausted()); + assertEquals(((MseBlock.Data) h1.readBlocking().getBlock()).asRowHeap().getRows().get(0), row1); + assertFalse(h1.isExhausted()); + + // Then the success EOS; the handle is exhausted afterwards. + assertTrue(h0.readBlocking().getBlock().isSuccess()); + assertTrue(h0.isExhausted()); + assertTrue(h1.readBlocking().getBlock().isSuccess()); + assertTrue(h1.isExhausted()); + + // Reading an exhausted handle returns the cached EOS without polling (and without re-merging stats). + assertTrue(h0.readBlocking().getBlock().isSuccess()); + + // Identical upstream shapes merge cleanly, so upstream stats survive. + MultiStageQueryStats stats = consumer.calculateStats(); + assertNotNull(stats.getUpstreamStageStats(SENDER_STAGE_ID), + "Upstream stats should be retained when both EOS blocks share the same shape"); + } + + @Test + public void errorIsCachedAndShortCircuitsEveryHandle() { + FakeStream s0 = new FakeStream("s0"); + FakeStream s1 = new FakeStream("s1"); + s0.enqueue(OperatorTestUtil.errorWithEmptyStats(new RuntimeException("boom"))); + // s1 has data, but the global error must short-circuit it before that data is ever read. + s1.enqueue(OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1})); + + BlockingMultiStreamConsumer.OfMseBlock consumer = newConsumer(FAR_FUTURE, List.of(s0, s1)); + List> handles = + consumer.streamHandles(); + + ReceivingMailbox.MseBlockWithStats err0 = handles.get(0).readBlocking(); + assertTrue(err0.getBlock().isError()); + assertTrue(((ErrorMseBlock) err0.getBlock()).getErrorMessages().get(QueryErrorCode.UNKNOWN).contains("boom")); + + // The other handle returns the very same error element, not its pending data block. + ReceivingMailbox.MseBlockWithStats err1 = handles.get(1).readBlocking(); + assertSame(err1, err0); + } + + @Test + public void timeoutReturnsErrorBlockWithSerializedStats() { + FakeStream s0 = new FakeStream("s0"); + // Deadline already in the past => the deadline loop times out immediately with no data. + BlockingMultiStreamConsumer.OfMseBlock consumer = + newConsumer(System.currentTimeMillis() - 1L, List.of(s0)); + BlockingMultiStreamConsumer.StreamHandle h0 = + consumer.streamHandles().get(0); + + ReceivingMailbox.MseBlockWithStats block = h0.readBlocking(); + assertTrue(block.getBlock().isError()); + assertTrue(((ErrorMseBlock) block.getBlock()).getErrorMessages().containsKey(QueryErrorCode.EXECUTION_TIMEOUT)); + // The timeout element carries the serialized accumulated stats (mirrors onException(code, msg)). + assertNotNull(block.getSerializedStats()); + + // The timeout latched a global error, so the handle keeps returning it. + assertSame(h0.readBlocking(), block); + } + + @Test + public void modeGuardRejectsMixingReads() { + // Round-robin first (empty mailboxes -> immediate success), then per-stream must throw. + BlockingMultiStreamConsumer.OfMseBlock roundRobinFirst = newConsumer(FAR_FUTURE, List.of()); + assertTrue(roundRobinFirst.readBlockBlocking().getBlock().isSuccess()); + assertThrows(IllegalStateException.class, roundRobinFirst::streamHandles); + + // Per-stream first, then round-robin must throw. + BlockingMultiStreamConsumer.OfMseBlock perStreamFirst = newConsumer(FAR_FUTURE, List.of(new FakeStream("s0"))); + perStreamFirst.streamHandles(); + assertThrows(IllegalStateException.class, perStreamFirst::readBlockBlocking); + } + + @Test + public void readBlockingUnblocksOnNewDataNotification() + throws InterruptedException { + FakeStream s0 = new FakeStream("s0"); + BlockingMultiStreamConsumer.OfMseBlock consumer = newConsumer(FAR_FUTURE, List.of(s0)); + BlockingMultiStreamConsumer.StreamHandle h0 = + consumer.streamHandles().get(0); + + ReceivingMailbox.MseBlockWithStats data = OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{7, 7}); + // Producer thread: after the consumer is (very likely) blocked, enqueue a block and fire the new-data callback. + Thread producer = new Thread(() -> { + try { + Thread.sleep(150L); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + s0.enqueue(data); + s0.fireNewData(); + }); + producer.start(); + + // First optimistic poll sees nothing, so this blocks on the shared wakeup until the producer fires it. + ReceivingMailbox.MseBlockWithStats read = h0.readBlocking(); + producer.join(); + assertSame(read, data); + } + + @Test + public void earlyTerminateDelegatesToStream() { + FakeStream s0 = new FakeStream("s0"); + BlockingMultiStreamConsumer.OfMseBlock consumer = newConsumer(FAR_FUTURE, List.of(s0)); + consumer.streamHandles().get(0).earlyTerminate(); + assertTrue(s0._earlyTerminated); + } + + /** + * {@link BlockingMultiStreamConsumer.StreamHandle#poll()} is the non-blocking primitive the k-way merge's cooperative + * drain relies on: it must never park, must surface data/EOS exactly like {@code readBlocking()} does, and must not + * re-poll an already-exhausted stream (mirrors the mailbox-release comment on the {@code Handle} implementation). + */ + @Test + public void pollIsNonBlockingAndTracksExhaustion() { + FakeStream s0 = new FakeStream("s0"); + BlockingMultiStreamConsumer.OfMseBlock consumer = newConsumer(FAR_FUTURE, List.of(s0)); + BlockingMultiStreamConsumer.StreamHandle h0 = + consumer.streamHandles().get(0); + + // Nothing scripted yet: poll() must return null immediately rather than parking. + assertNull(h0.poll()); + assertFalse(h0.isExhausted()); + + Object[] row = new Object[]{9, 9}; + s0.enqueue(OperatorTestUtil.blockWithStats(DATA_SCHEMA, row)); + ReceivingMailbox.MseBlockWithStats data = h0.poll(); + assertNotNull(data); + assertEquals(((MseBlock.Data) data.getBlock()).asRowHeap().getRows().get(0), row); + assertFalse(h0.isExhausted()); + + // Drained again with nothing scripted: back to null, not blocking. + assertNull(h0.poll()); + + s0.enqueue(OperatorTestUtil.eosWithEmptyStats()); + assertTrue(h0.poll().getBlock().isSuccess()); + assertTrue(h0.isExhausted()); + + // An exhausted handle's poll() returns null without touching the (already-released) underlying stream again. + assertNull(h0.poll()); + } + + @Test + public void pollShortCircuitsOnGlobalError() { + FakeStream s0 = new FakeStream("s0"); + FakeStream s1 = new FakeStream("s1"); + s0.enqueue(OperatorTestUtil.errorWithEmptyStats(new RuntimeException("boom"))); + s1.enqueue(OperatorTestUtil.blockWithStats(DATA_SCHEMA, new Object[]{1, 1})); + + BlockingMultiStreamConsumer.OfMseBlock consumer = newConsumer(FAR_FUTURE, List.of(s0, s1)); + List> handles = + consumer.streamHandles(); + + ReceivingMailbox.MseBlockWithStats err = handles.get(0).poll(); + assertNotNull(err); + assertTrue(err.getBlock().isError()); + // The other stream still has data queued, but the global error must short-circuit its poll() too. + assertSame(handles.get(1).poll(), err); + } + + @Test + public void awaitDataOrTerminalReturnsNullOnWakeAndErrorOnTimeout() + throws InterruptedException { + FakeStream s0 = new FakeStream("s0"); + BlockingMultiStreamConsumer.OfMseBlock consumer = newConsumer(FAR_FUTURE, List.of(s0)); + consumer.streamHandles(); + + Thread producer = new Thread(() -> { + try { + Thread.sleep(150L); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + s0.fireNewData(); + }); + producer.start(); + // Parks until the producer's new-data signal wakes it; returns null (caller should re-poll) rather than an + // element, since awaitDataOrTerminal() never reads from any stream itself. + assertNull(consumer.awaitDataOrTerminal()); + producer.join(); + + // A deadline already in the past must return the cached timeout error instead of parking. + BlockingMultiStreamConsumer.OfMseBlock timedOut = + newConsumer(System.currentTimeMillis() - 1L, List.of(new FakeStream("s1"))); + timedOut.streamHandles(); + ReceivingMailbox.MseBlockWithStats element = timedOut.awaitDataOrTerminal(); + assertNotNull(element); + assertTrue(element.getBlock().isError()); + assertTrue(((ErrorMseBlock) element.getBlock()).getErrorMessages().containsKey(QueryErrorCode.EXECUTION_TIMEOUT)); + } + + /** + * A hand-written {@link AsyncStream} whose {@link #poll()} drains a scripted queue (empty queue => not ready yet) and + * which exposes a way to fire the captured new-data callback, so tests can drive the per-stream blocking loop without + * any real mailbox infrastructure. + */ + private static class FakeStream implements AsyncStream { + private final Object _id; + private final Deque _scripted = new ArrayDeque<>(); + @Nullable + private OnNewData _listener; + private volatile boolean _earlyTerminated; + + FakeStream(Object id) { + _id = id; + } + + void enqueue(ReceivingMailbox.MseBlockWithStats block) { + _scripted.addLast(block); + } + + void fireNewData() { + if (_listener != null) { + _listener.newDataAvailable(); + } + } + + @Override + public Object getId() { + return _id; + } + + @Nullable + @Override + public ReceivingMailbox.MseBlockWithStats poll() { + return _scripted.pollFirst(); + } + + @Override + public void addOnNewDataListener(OnNewData onNewData) { + _listener = onNewData; + } + + @Override + public void cancel() { + } + + @Override + public void earlyTerminate() { + _earlyTerminated = true; + } + } +} diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 3ddfaf79c0f7..cc415d7836c7 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -543,6 +543,9 @@ public static class Broker { public static final String CONFIG_OF_MSE_STREAMING_GROUP_BY_FLUSH_THRESHOLD = "pinot.broker.mse.streaming.group.by.flush.threshold"; public static final int DEFAULT_MSE_STREAMING_GROUP_BY_FLUSH_THRESHOLD = -1; + /// Default output block size (rows) for the streaming selection ORDER BY combine + /// ({@link Request.QueryOptionKey#SORTED_SELECTION_MERGE_BLOCK_SIZE}). + public static final int DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE = 10_000; // Whether to infer partition hint by default or not. // This value can always be overridden by INFER_PARTITION_HINT query option public static final String CONFIG_OF_INFER_PARTITION_HINT = "pinot.broker.multistage.infer.partition.hint"; @@ -819,6 +822,40 @@ public static class QueryOptionKey { /// Flush threshold for streaming group-by on MSE leaf stages. public static final String STREAMING_GROUP_BY_FLUSH_THRESHOLD = "streamingGroupByFlushThreshold"; + /// Opt-in: use the streaming k-way-merge selection ORDER BY combine over sorted segments. + /// + /// This option has two effects. On the server it selects the streaming selection ORDER BY combine operator. + /// On the broker it additionally allows the multi-stage planner to mark a validated leaf selection ORDER BY + /// sender fragment as sorted-on-sender, which is one of the two ways the precondition + /// [#STREAMING_SORTED_MAILBOX_RECEIVE] needs can be satisfied. It is required only for that plain + /// leaf-selection path, where the rel plan contains no sort exchange: when the rel plan already carries a + /// sort exchange with sender-side sorting, [#STREAMING_SORTED_MAILBOX_RECEIVE] alone is enough. + /// + /// NOTE: the planner-side marking is a no-op under usePhysicalOptimizer (the v2 path does not go through + /// PlanFragmenter). The server-side combine selection is unaffected. + public static final String SORTED_SELECTION_MERGE_ENABLED = "sortedSelectionMergeEnabled"; + /// Output block size (rows) for the streaming selection ORDER BY combine. + public static final String SORTED_SELECTION_MERGE_BLOCK_SIZE = "sortedSelectionMergeBlockSize"; + + /// Opt-in for the streaming k-way merge in SortedMailboxReceiveOperator. true = use the k-way merge when the + /// planner has also proven senders are sorted (isSortedOnSender); unset or false = always + /// accumulate-then-sort. isSortedOnSender is set either by a rel-level sort exchange that already declares + /// sender-side sorting, or (for a plain leaf selection ORDER BY, where no such exchange exists) by + /// PlanFragmenter under [#SORTED_SELECTION_MERGE_ENABLED]. Only the latter case needs both options. + /// Whether the merge actually ran is reported per receive operator in the response stageStats as + /// `kWayMergeUsed`. + /// Note: like SQL ORDER BY, the order among rows with equal collation keys is unspecified and may differ + /// between the merge path and the accumulate-then-sort path; the output row multiset is identical. + /// NOTE: This is a no-op under usePhysicalOptimizer (the v2 path does not set sorted-on-sender). + /// NOTE: do not enable during a mixed-version rollout. `kWayMergeUsed` is a new StatMap key, and StatMap + /// decodes keys by ordinal without a bounds check, so a peer running a build that predates the key cannot + /// deserialize the stage stats it appears in. The key is only written when this option is on, so leaving + /// the option off (the default) keeps a mixed-version cluster safe. + public static final String STREAMING_SORTED_MAILBOX_RECEIVE = "streamingSortedMailboxReceive"; + /// Output block size (rows) for the streaming sorted mailbox-receive k-way merge; defaults to 10000. + public static final String STREAMING_SORTED_MAILBOX_RECEIVE_BLOCK_SIZE = + "streamingSortedMailboxReceiveBlockSize"; + public static final String NUM_REPLICA_GROUPS_TO_QUERY = "numReplicaGroupsToQuery"; public static final String ORDERED_PREFERRED_POOLS = "orderedPreferredPools"; public static final String USE_FIXED_REPLICA = "useFixedReplica";