rohityadav1993 commented on code in PR #19120: URL: https://github.com/apache/pinot/pull/19120#discussion_r4125620614
########## 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<SelectionResultsBlock> { + 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<Object[]> _comparator; + private final SegmentCursor[] _sortedCursors; + private final PriorityQueue<SegmentCursor> _priorityQueue; + + // Merge progress (single-threaded; mutated only by the consumer thread driving getNextBlock()) + private int _nextToActivate; + private int _numRowsEmitted; + private List<Object[]> _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<String> _dataSchemaMismatchErrors; + /// Subset of the above not yet attached to an emitted block; drained on each attach + @Nullable + private List<String> _unreportedDataSchemaMismatchErrors; + + public StreamingSelectionOrderByCombineOperator(List<Operator> 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(); Review Comment: Done in ff800d6a20, per segment, but via ColumnMetadata.isNonNull() (#19072) instead of the null value vector. Pruning is decided before any segment is acquired, and a tiered SegmentDirectory may only materialize readers on acquire, so probing the vector would fetch the segments pruning is meant to skip. The flag is loaded with segment metadata, like the min/max already used here. With null handling on, a flagged segment prunes normally; an unflagged one (nulls, a pre-#19072 segment, or a consuming segment) always activates, as every segment did before. AUTO uses the same predicate, which also closes the null gap from [1]. I kept null-bearing segments on always-activate rather than pruning them in the safe direction: that depends on each column's default null value, and can follow if it matters. Tests cover flagged, unflagged, null-bearing, DESC reverse-scan, consuming-segment and AUTO cases, each watched failing first. ########## 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<SelectionResultsBlock> { + 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<Object[]> _comparator; + private final SegmentCursor[] _sortedCursors; + private final PriorityQueue<SegmentCursor> _priorityQueue; + + // Merge progress (single-threaded; mutated only by the consumer thread driving getNextBlock()) + private int _nextToActivate; + private int _numRowsEmitted; + private List<Object[]> _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<String> _dataSchemaMismatchErrors; + /// Subset of the above not yet attached to an emitted block; drained on each attach + @Nullable + private List<String> _unreportedDataSchemaMismatchErrors; + + public StreamingSelectionOrderByCombineOperator(List<Operator> 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<OrderByExpressionContext> 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<BaseResultsBlock> 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) { Review Comment: Agreed, including that capping active cursors would be wrong. I'd like to defer the levers to a follow-up: - The win you point out is now in the description. - The cost is now measured (https://github.com/rohityadav1993/pinot/pull/1). At LIMIT 1M, MinMax needs about 950 MB of heap at any overlap; streaming needs under 64 MB disjoint and 237 MB at full overlap. The only adverse memory case is full overlap at LIMIT 10K with a two-column ORDER BY, about 1.27x MinMax; listed under Known gaps. - Two changes this round lower depth: pruning stays on under null handling ([9], ff800d6a20), and tied minima under a single-column ORDER BY are deferred (82731f46d1). Follow-up order I'd propose: overlap-aware AUTO fallback, smaller maxDocsPerCall for streaming children, then late materialisation. Parallel merge (b) as its own change with a benchmark, since it brings back the Phaser fence. ########## 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<SelectionResultsBlock> { + 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<Object[]> _comparator; + private final SegmentCursor[] _sortedCursors; + private final PriorityQueue<SegmentCursor> _priorityQueue; + + // Merge progress (single-threaded; mutated only by the consumer thread driving getNextBlock()) + private int _nextToActivate; + private int _numRowsEmitted; + private List<Object[]> _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<String> _dataSchemaMismatchErrors; + /// Subset of the above not yet attached to an emitted block; drained on each attach + @Nullable + private List<String> _unreportedDataSchemaMismatchErrors; + + public StreamingSelectionOrderByCombineOperator(List<Operator> 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<OrderByExpressionContext> 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<BaseResultsBlock> 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() { Review Comment: Done in e3a908aca5. Explain: blockSize, frontierPruning, deferTiedCursors, numSegments and numSortedSegments, visible on CombineSelectOrderbyStreaming with explainAskingServers=true. blockSize is idempotent, so servers with different defaults stay separate nodes. Stats: a never-activated segment scans no docs, so numSegmentsProcessed - numSegmentsMatched counts skipped segments (plus activated ones with no matches); a test pins this. A dedicated activated-count key needs a new DataTable MetadataKey, so it is a TODO pointing here. No overlap estimate: it is per worker and would make the explain node unmergeable across workers. ########## pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java: ########## @@ -0,0 +1,529 @@ +/** + * 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. + * + * <p>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))}). + * + * <p>It runs in one of two emission modes: + * <ul> + * <li><b>No tail to sort</b> ({@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).</li> + * <li><b>Tail to sort</b> ({@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).</li> + * </ul> + * + * <p>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. + * + * <p>This operator is stateful across {@link #getNextBlock()} calls and is <b>not</b> thread-safe; a single consumer + * must drive it. + */ +public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionResultsBlock> { + 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<ExpressionContext> _expressions; + private final BaseProjectOperator<?> _projectOperator; + private final List<OrderByExpressionContext> _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<ExpressionContext> _phase1Expressions; + private final int _numPhase1Columns; + private final Comparator<Object[]> _comparator; + // Compares only the first order-by column; used to detect primary-value run boundaries + private final Comparator<Object[]> _primaryComparator; + // Pre-allocated run heap (cleared and reused each nextRun() call to avoid per-run allocation) + private final Comparator<Object[]> _reversedComparator; + private final PriorityQueue<Object[]> _runHeap; + + // Pre-computed invariants for the two-phase fetch (null when single-phase) + private final List<ExpressionContext> _nonOrderByExpressions; + private final Map<String, DataSource> _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<ExpressionContext> 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<String> 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<Object[]> 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<Object[]> 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<Object[]> 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<Object[]> nextRun() { + int remaining = _numRowsToKeep - _numRowsEmitted; + if (remaining <= 0) { + return null; + } + if (_pendingRow == null) { + _pendingRow = nextRow(); + if (_pendingRow == null) { + return null; + } + } + PriorityQueue<Object[]> 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) { Review Comment: Still out of scope. When the segment is sorted on only a prefix of the ORDER BY and each prefix value has few rows, the per-run sort pass dominates. Results stay correct, and the path is opt-in. Tracked as a follow-up with [10]'s levers. ########## pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java: ########## @@ -170,13 +164,85 @@ public void testCancelGroupByOrderByCombineOperator() { testCancelCombineOperator(combineOperator, ready); } + @Test + public void testCancelStreamingSelectionOrderByCombineOperator() { + CountDownLatch ready = new CountDownLatch(1); + List<Operator> 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<Operator> 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 <b>before</b> 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<Operator> operators = getOperators(null, minMaxSegmentSupplier()); Review Comment: Follow-up: the offered timeOut is in (e3a908aca5), so a regression fails the deadline test instead of hanging the build. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
