Jackie-Jiang commented on code in PR #19120: URL: https://github.com/apache/pinot/pull/19120#discussion_r4251361213
########## pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java: ########## @@ -0,0 +1,575 @@ +/** + * 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.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.function.Function; +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.common.utils.HashUtil; +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.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.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 [SelectionOrderByOperator] (which materializes the segment's whole top-K in a single block) this +/// operator emits one globally-sorted [SelectionResultsBlock] per [#getNextBlock()] call and returns +/// `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 `projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))`). +/// +/// Every instance emits at least one block before `null`: a segment matching no rows emits a single empty block +/// carrying the [DataSchema], so consumers never have to reconstruct a schema the segment already knows. +/// +/// It runs in one of two emission modes: +/// +/// - **No tail to sort** (`numSortedExpressions == numOrderByExpressions`, e.g. `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 `limit + offset` budget). +/// - **Tail to sort** (`numSortedExpressions < numOrderByExpressions`, e.g. +/// `ORDER BY sorted, other`): +/// each call reads forward until the first order-by value changes (a primary-value "run"), retains the run's top +/// `limit + offset` rows by the full comparator, and emits them sorted. This bounds the in-memory run buffer to +/// `limit + offset` rows even when the first order-by column is near-constant (very low cardinality). +/// +/// Like [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 [#getNextBlock()] calls and is **not** 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; + /// Source column contexts for the phase-2 expressions; built only on the zero-match schema path (null otherwise) + private Map<String, ColumnContext> _phase2ColumnContextMap; + + // 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; + /// Whether any block has been emitted yet; drives the zero-match segment's one schema-carrying empty block + private boolean _emittedAnyBlock; + 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; + if (_emittedAnyBlock) { + return null; + } + // Zero-match segment: emit one empty block carrying the schema before signalling end-of-stream, so consumers + // never have to reconstruct a schema the segment already knows. Two-phase leaves _dataSchema null until the + // first fetch, which never happens here, so build it from column metadata alone. + if (_twoPhase && _dataSchema == null) { + _dataSchema = buildTwoPhaseDataSchema(this::resolveResultColumnContext); + } + assert _dataSchema != null; + return new SelectionResultsBlock(_dataSchema, List.of(), _comparator, _queryContext); + } + _emittedAnyBlock = true; + 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 `limit + offset` budget. Returns `null` when exhausted, and otherwise a + /// non-empty list -- empty project blocks are skipped and the row budget is checked up front, so "no rows" always + /// means "no more rows". + @Nullable + private List<Object[]> nextSortedRows() { + int remaining = _numRowsToKeep - _numRowsEmitted; + if (remaining <= 0) { + return null; + } + ValueBlock valueBlock = nextNonEmptyBlock(); + 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 + /// `limit + offset` rows by the full comparator, and return them sorted. Returns `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) { + SelectionOperatorUtils.addToPriorityQueue(row, runHeap, _numRowsToKeep); + } else { + // Run boundary: this row starts the next run, keep it for the next call. + _pendingRow = row; + break; + } + } + List<Object[]> 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 + /// `Object[_numExpressions]`. For two-phase the document id is stashed at index + /// `_numOrderByExpressions` (overwritten in the second pass). Returns `null` when the project operator + /// is exhausted. + @Nullable + private Object[] nextRow() { + if (_currentBlock == null || _currentPos >= _currentNumDocs) { + if (_projectExhausted) { + return null; + } + _currentBlock = nextNonEmptyBlock(); + 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; Review Comment: **[P1] Copy document IDs before retaining them across segment pulls** With streaming enabled, an unsorted ORDER BY tail, and a non-ORDER-BY payload, this retains the producer's document-ID array across `nextRun()` calls. Ordinary plans use `DocIdSetOperator`, whose scratch array is shared by every instance on the same thread. The combine interleaves segment readers, so another segment overwrites the retained IDs while the cached ORDER BY values remain intact. For example, let overlapping segments have matching document IDs A=`[0,2,4,6]` and B=`[1,3,5,7]`, with matching primary-key runs `[0,0,1,1]`. After A yields its first run and B overwrites the shared buffer, A can resume with document 6's cached sort values but document ID 7. Phase two then fetches document 7's payload, potentially from a row excluded by the filter. `ReverseDocIdSetOperator` has the same thread-local sharing model, so DESC also needs coverage. Copy the active document-ID prefix when retaining this block. Please cover overlapping filtered segments with different document positions, multiple runs in one project block, and payload assertions. This defect is established by source tracing; I have not run a reproducing test. ########## pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java: ########## @@ -0,0 +1,655 @@ +/** + * 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.HashSet; +import java.util.List; +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.core.common.Operator; +import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator; +import org.apache.pinot.core.operator.ExplainAttributeBuilder; +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.plan.SelectionPlanNode; +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.MutableSegment; +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 `_operators`, returning +/// globally sorted rows in bounded blocks. Each segment is exposed through a [SegmentCursor] that yields that +/// segment's locally-sorted rows in order: +/// +/// - Segments physically sorted on the first order-by column are backed by +/// [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 +/// [SelectionResultsBlock]-producing operator such as `SelectionOrderByOperator`); the cursor reads that +/// one block and iterates its rows. +/// +/// A [PriorityQueue] of [SegmentCursor] ordered by the [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 `MinMaxValueContext` idea from +/// [MinMaxValueBasedSelectionOrderByCombineOperator]. A cursor is only activated (its segment acquired and first +/// block read) when the merge frontier reaches its min/max, so once `limit + offset` rows are emitted the +/// remaining segments are never acquired or read. See [#activateEligibleCursors(SegmentCursor)] for the +/// correctness argument. +/// A consuming ([MutableSegment]) segment is always given no min/max and activates up front: its bounds are a snapshot +/// that rows ingested before activation can outgrow, since its plan (and doc count) is only built at activation. +/// With null handling on, a segment whose leading column is not flagged non-null by +/// [SelectionPlanNode#isColumnFlaggedNonNull] (it has nulls, the segment predates the flag, or it is a consuming +/// segment) is given no min/max and always activates. Flagged segments keep their min/max. With +/// null handling off, min/max are used as stored. The flag is read from segment metadata, not the null value vector: +/// that vector is a mapped buffer, and probing it here would acquire segments pruning is about to skip. +/// Its effect shows in the stats: a never-activated segment scans no docs, so `numSegmentsMatched` excludes it while +/// `numSegmentsProcessed` counts every segment. Their difference also counts activated segments that matched no rows. +/// TODO: report the activated count separately; that needs a new `DataTable.MetadataKey`. +/// See https://github.com/apache/pinot/pull/19120#discussion_r3871714063 +/// +/// **Segment acquire/release lifecycle.** A cursor acquires its +/// [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 +/// [StreamingSelectionOrderByOperator] retains a buffer-backed `ValueBlock` across `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 `Object[]` (via +/// `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 [#releaseAllCursors()]. +/// +/// **Streaming paths only.** This operator is installed only where a +/// [org.apache.pinot.core.query.executor.ResultsBlockStreamer] is present (the +/// MSE leaf, driven by [org.apache.pinot.core.operator.streaming.StreamingInstanceResponseOperator], and the +/// gRPC streaming endpoint). The merge emits many bounded [SelectionResultsBlock]s from successive +/// [#getNextBlock()] calls followed by a final [MetadataResultsBlock]. On the blocking single-stage +/// path the caller expects one complete block from a single [#getNextBlock()] call, so nothing here could +/// stream: the merge would accumulate the whole `limit + offset` result while holding every activated cursor's +/// materialized block alive. That path keeps [MinMaxValueBasedSelectionOrderByCombineOperator]. +/// +/// **Threading.** This operator overrides [#start()]/[#stop()] to no-ops (other than releasing +/// segments) and runs the merge single-threaded and lazily in [#getNextBlock()] on the consumer thread; it does +/// not use the base worker-queue model. The base entry points that it replaces - [#processSegments()] and +/// [#isQuerySatisfied(SelectionResultsBlock, Object)] - are overridden to fail loud. The base +/// `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 _asc; + private final boolean _pruningEnabled; + /// Whether a cursor whose bound *ties* the merge frontier may be deferred, not only one sorting strictly + /// beyond it. Only correct under a single order-by expression: see [#sortsBeyond]. + private final boolean _deferTiedCursors; + private final int _numRowsToKeep; + private final int _blockSize; + /// Segments whose metadata marks the leading order-by column sorted; -1 when that expression is not a column + private final int _numSortedSegments; + 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. Non-null by the time any rows + /// exist to emit: every streaming child emits at least one block carrying its schema, including a zero-match one. + 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) { + // Pass a null merger: we override the consumption path entirely and never touch the base merger / worker queue. + // Every base entry point that would dereference it - processSegments() and isQuerySatisfied() - is overridden to + // throw, so a future change to the base streaming loop fails at that seam instead of with an NPE at query time. + // The executor is unused here (start() is a no-op) but is still handed to the base so that the worker-queue model + // stays constructible; passing null would trade one latent NPE for another. + super(null, operators, queryContext, executorService); + _numRowsToKeep = queryContext.getLimit() + queryContext.getOffset(); + _blockSize = queryContext.getSortedSelectionMergeBlockSize(); + + List<OrderByExpressionContext> orderByExpressions = queryContext.getOrderByExpressions(); + assert orderByExpressions != null && !orderByExpressions.isEmpty(); + OrderByExpressionContext firstOrderByExpression = orderByExpressions.get(0); + _asc = firstOrderByExpression.isAsc(); + _comparator = OrderByComparatorFactory.getComparator(orderByExpressions, queryContext.isNullHandlingEnabled()); + + // Frontier pruning needs a physical column to read min/max from. Under SortedSelectionMergeMode.ON the leading + // order-by may be an expression (forced, for benchmarking), in which case there is no such column: every cursor + // is activated up front in plan order and the merge is a plain k-way heap over materialized children. + ExpressionContext firstOrderByExpressionContext = firstOrderByExpression.getExpression(); + String firstOrderByColumn = + firstOrderByExpressionContext.getType() == ExpressionContext.Type.IDENTIFIER + ? firstOrderByExpressionContext.getIdentifier() : null; + _pruningEnabled = firstOrderByColumn != null; + // A column-0 tie only proves the segment cannot supply an *earlier* row when column 0 is the whole sort key. + // With two or more expressions a tie on column 0 can hide a row sorting earlier on column 1, so ties must + // still activate (same gate as MinMaxValueBasedSelectionOrderByCombineOperator's numOrderByExpressions == 1). + _deferTiedCursors = orderByExpressions.size() == 1; + + // 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]; + int numSortedSegments = firstOrderByColumn != null ? 0 : -1; + for (int i = 0; i < _numOperators; i++) { + Operator<BaseResultsBlock> operator = _operators.get(i); + if (firstOrderByColumn == null) { + _sortedCursors[i] = new SegmentCursor(operator, null, null); + continue; + } + DataSourceMetadata metadata = + operator.getIndexSegment().getDataSource(firstOrderByColumn, queryContext.getSchema()) + .getDataSourceMetadata(); + Comparable minValue = metadata.getMinValue(); + Comparable maxValue = metadata.getMaxValue(); + // A consuming segment's min/max is a snapshot copied when getDataSource() builds the MutableDataSource, but under + // prefetch (AcquireReleaseColumnsSegmentOperator) the child plan, and with it the doc count, is only built at + // activation. Rows indexed in between can fall outside these bounds and be emitted out of order, so treat the + // bounds as unknown and let activateEligibleCursors force-activate this cursor. + if (operator.getIndexSegment() instanceof MutableSegment) { + minValue = null; + maxValue = null; + } + // isNonNull is loaded with the segment, the same class of access as min/max. False means nulls or unknown, so + // drop the bounds and let activateEligibleCursors force-activate this cursor. Do not read the null vector here. + if (queryContext.isNullHandlingEnabled() + && !SelectionPlanNode.isColumnFlaggedNonNull(operator.getIndexSegment(), firstOrderByColumn)) { + minValue = null; + maxValue = null; + } + _sortedCursors[i] = new SegmentCursor(operator, minValue, maxValue); + if (metadata.isSorted()) { + numSortedSegments++; + } + } + _numSortedSegments = numSortedSegments; + if (firstOrderByColumn != null) { + 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 [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; + } + + /// Settings fixed at construction. Explain never drives the merge, so how many segments activate is not shown here. + /// Query-wide values are idempotent so the broker merges the node across workers only when they agree; the segment + /// counts are summed. `numSortedSegments` is metadata sortedness, not whether a child streams: a DESC scan without + /// `allowReverseOrder`, or a leading column holding nulls, still falls back to a materialized child. + @Override + protected void explainAttributes(ExplainAttributeBuilder attributeBuilder) { + super.explainAttributes(attributeBuilder); + attributeBuilder.putLongIdempotent("blockSize", _blockSize); + attributeBuilder.putBool("frontierPruning", _pruningEnabled); + attributeBuilder.putBool("deferTiedCursors", _deferTiedCursors); + attributeBuilder.putLong("numSegments", _numOperators); + if (_numSortedSegments >= 0) { + attributeBuilder.putLong("numSortedSegments", _numSortedSegments); + } + } + + /// Override to a no-op: the merge is single-threaded and lazy in [#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 ([#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"); + } + + /// Only reachable from the base streaming loop, which this operator replaces; the base implementation would + /// dereference the null merger. Fail loud rather than NPE if a future change routes back through it. + @Override + protected boolean isQuerySatisfied(SelectionResultsBlock resultsBlock, Object tracker) { + throw new IllegalStateException( + "StreamingSelectionOrderByCombineOperator tracks its own limit; isQuerySatisfied() 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(); + // The leader is retained outside the heap: segments are usually near-disjoint on a time-like leading order-by + // column, so one cursor supplies long runs and re-heaping it per row costs two O(log k) sifts to reach the same + // cursor again. While it keeps winning we only peek and compare. + // TODO: move whole runs with one addAll(subList(...)) after a binary search for the cut. Deferred because + // advancing _numRowsEmitted by a run steps over the 0x1FFF mask in the termination check below, losing + // cancellation/pause/deadline observation; needs an unconditional check per run or a run-length clamp. + // See https://github.com/apache/pinot/pull/19120#discussion_r3871714056 + SegmentCursor leader = null; + 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(leader); + // Choose the leader strictly AFTER activation, so a cursor just activated can win instead of being skipped + // for one row. + if (leader == null) { + leader = _priorityQueue.poll(); + if (leader == null) { + // All active cursors exhausted (and pruning guarantees the rest cannot contribute). + break; + } + } else { + SegmentCursor top = _priorityQueue.peek(); + if (top != null && _comparator.compare(leader.currentHead(), top.currentHead()) > 0) { + // Lost the lead; hand it back and take the new minimum. A tie keeps the incumbent, reordering only rows + // equal on every order-by expression -- the heap's own tie-breaking is arbitrary too. + _priorityQueue.offer(leader); + leader = _priorityQueue.poll(); + } + } + _outputRows.add(leader.currentHead()); + _numRowsEmitted++; + leader.advance(); + if (leader.currentHead() == null) { + leader = null; + } + if (_outputRows.size() >= _blockSize) { + // Restore the invariant before handing control back: at every entry to and exit from getNextBlock() the + // heap holds all active cursors, so nothing else has to know about the retained leader. + if (leader != null) { + _priorityQueue.offer(leader); + } + return flushDataBlock(); + } + } + // Merge complete; the leader goes back so finish() sees a consistent heap. + if (leader != null) { + _priorityQueue.offer(leader); + } + finish(); + if (!_outputRows.isEmpty()) { + // The final partial data block; the next call returns the terminal metadata block. + return flushDataBlock(); + } + return attachDataSchemaMismatchErrors(attachExecutionStats(new MetadataResultsBlock())); + } 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 bound can still reach the merge frontier. + /// + /// Cursors are visited in min/max order and `_nextToActivate` advances only on a real activation, so a + /// `break` defers the current cursor to a later call with a risen frontier rather than skipping it. Deferring + /// is safe because the first order-by column is the primary sort key: a cursor whose range starts past the frontier + /// holds no row sorting before it. With no frontier known yet, activation is forced. Under a single order-by + /// expression a cursor whose bound merely *ties* the frontier is deferred too, which is what keeps a long run of + /// equal segment minima from activating every segment at once; see [#sortsBeyond]. + /// + /// The frontier is the row about to be emitted. Understating it defers a cursor that could have supplied that row + /// and the merge emits out of order, where overstating it only activates a segment early. The leader is retained + /// outside the heap, so that row is the smaller of its head and the heap head: a bound past either is past the + /// frontier. Checking the leader alone would overstate it whenever the leader has moved past the heap head, which + /// happens both on entry (the leader just advanced) and within this loop (an activation offers a smaller head). + private void activateEligibleCursors(@Nullable SegmentCursor leader) { + while (_nextToActivate < _sortedCursors.length) { + SegmentCursor cursor = _sortedCursors[_nextToActivate]; + if (_pruningEnabled) { + Comparable bound = _asc ? cursor._minValue : cursor._maxValue; + // Re-read per iteration: each activation below can offer a smaller head into the heap. + SegmentCursor top = _priorityQueue.peek(); + // A null bound always activates, and with no candidate at all there is no frontier to prune against. + if (bound != null && (leader != null || top != null) && sortsBeyond(bound, leader, top)) { + break; + } + } + cursor.activate(); + _nextToActivate++; + if (cursor.currentHead() != null) { + _priorityQueue.offer(cursor); + } + } + } + + /// Whether `bound` sorts past the frontier -- the smaller of the two candidates' heads -- on the first order-by + /// column, so that segment cannot supply the next row. An absent candidate imposes no constraint; a null head cannot + /// be compared, so this reports `false` and the caller activates. Tests column 0 only, as the pruning bound always + /// has, rather than the full-row comparator -- which would put every order-by column on the per-row path. + /// + /// Under a single order-by expression (`_deferTiedCursors`) a *tie* also defers, which only postpones a + /// cursor: `_nextToActivate` does not advance, and the caller force-activates once heap and leader drain. + /// Since `bound` bounds every row in the segment, a deferred cursor holds no row sorting strictly before an + /// emitted one -- only ties, interchangeable when column 0 is the whole sort key. Note this is *not* + /// [MinMaxValueBasedSelectionOrderByCombineOperator]'s justification: its bound is a complete top-K's k-th + /// row, this one is the live frontier. It changes which tied rows are returned, already documented as arbitrary. + private boolean sortsBeyond(Comparable bound, @Nullable SegmentCursor leader, @Nullable SegmentCursor top) { + Object leaderHead = leader != null ? leader.currentHead()[0] : null; + Object topHead = top != null ? top.currentHead()[0] : null; + if ((leader != null && leaderHead == null) || (top != null && topHead == null)) { + return false; + } + // Past either head is past the smaller one. + return (leader != null && sortsBeyond(bound, leaderHead)) || (top != null && sortsBeyond(bound, topHead)); + } + + private boolean sortsBeyond(Comparable bound, Object headValue) { + // 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 (_deferTiedCursors) { + return _asc ? cmp >= 0 : cmp <= 0; + } + return _asc ? cmp > 0 : cmp < 0; + } + + /// 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 [SelectionResultsBlock] and resets the output buffer. The + /// block carries the comparator so the broker-side n-way reduce stays correct. Execution stats are not attached + /// here: they go on the terminal metadata block. + private BaseResultsBlock flushDataBlock() { + List<Object[]> rows = _outputRows; + _outputRows = newOutputList(); + // Rows only reach the output buffer through pullBlock(), which captures the schema before accepting any of them. + assert _dataSchema != null; + SelectionResultsBlock block = new SelectionResultsBlock(_dataSchema, rows, _comparator, _queryContext); + attachDataSchemaMismatchErrors(block); Review Comment: **[P1] Preserve schema-mismatch errors in the final response** When a mismatched segment is dropped and some result rows are emitted, this attaches the error to a data block and clears the pending errors. Neither supported consumer preserves it: MSE's `LeafOperator` queues/converts streamed data without checking its errors and checks only the final `InstanceResponseBlock`; `GrpcResultsBlockStreamer` sends `block.getDataTable()`, which does not copy the errors attached to `SelectionResultsBlock`. The terminal metadata is therefore clean and incomplete results can appear successful. Retain and attach these errors to terminal metadata, or explicitly propagate errors through both streaming consumers. The current test collects errors directly from raw combine blocks, so it does not exercise the response boundary where they disappear. ########## pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java: ########## @@ -0,0 +1,575 @@ +/** + * 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.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.function.Function; +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.common.utils.HashUtil; +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.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.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 [SelectionOrderByOperator] (which materializes the segment's whole top-K in a single block) this +/// operator emits one globally-sorted [SelectionResultsBlock] per [#getNextBlock()] call and returns +/// `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 `projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))`). +/// +/// Every instance emits at least one block before `null`: a segment matching no rows emits a single empty block +/// carrying the [DataSchema], so consumers never have to reconstruct a schema the segment already knows. +/// +/// It runs in one of two emission modes: +/// +/// - **No tail to sort** (`numSortedExpressions == numOrderByExpressions`, e.g. `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 `limit + offset` budget). +/// - **Tail to sort** (`numSortedExpressions < numOrderByExpressions`, e.g. +/// `ORDER BY sorted, other`): +/// each call reads forward until the first order-by value changes (a primary-value "run"), retains the run's top +/// `limit + offset` rows by the full comparator, and emits them sorted. This bounds the in-memory run buffer to +/// `limit + offset` rows even when the first order-by column is near-constant (very low cardinality). +/// +/// Like [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 [#getNextBlock()] calls and is **not** 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; + /// Source column contexts for the phase-2 expressions; built only on the zero-match schema path (null otherwise) + private Map<String, ColumnContext> _phase2ColumnContextMap; + + // 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; + /// Whether any block has been emitted yet; drives the zero-match segment's one schema-carrying empty block + private boolean _emittedAnyBlock; + 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); Review Comment: **[P2] Use the complete sorted prefix to delimit runs** `numSortedExpressions` determines whether a tail needs sorting, but this comparator then ignores every sorted column after the first. With constant physically sorted column `a`, increasing physically sorted column `b`, and unsorted `c`, `ORDER BY a,b,c LIMIT 1` supplies a sorted prefix of length two. This code nevertheless treats the entire segment as one run because every row shares `a`, scanning all N documents before emitting anything. The existing `LinearSelectionOrderByOperator` compares the complete sorted prefix and can stop within the first project block when `b` changes. Build the run-boundary comparator over `[0, numSortedExpressions)`; the run heap can compare only the unsorted suffix. ########## pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java: ########## @@ -0,0 +1,575 @@ +/** + * 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.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.function.Function; +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.common.utils.HashUtil; +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.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.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 [SelectionOrderByOperator] (which materializes the segment's whole top-K in a single block) this +/// operator emits one globally-sorted [SelectionResultsBlock] per [#getNextBlock()] call and returns +/// `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 `projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))`). +/// +/// Every instance emits at least one block before `null`: a segment matching no rows emits a single empty block +/// carrying the [DataSchema], so consumers never have to reconstruct a schema the segment already knows. +/// +/// It runs in one of two emission modes: +/// +/// - **No tail to sort** (`numSortedExpressions == numOrderByExpressions`, e.g. `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 `limit + offset` budget). +/// - **Tail to sort** (`numSortedExpressions < numOrderByExpressions`, e.g. +/// `ORDER BY sorted, other`): +/// each call reads forward until the first order-by value changes (a primary-value "run"), retains the run's top +/// `limit + offset` rows by the full comparator, and emits them sorted. This bounds the in-memory run buffer to +/// `limit + offset` rows even when the first order-by column is near-constant (very low cardinality). +/// +/// Like [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 [#getNextBlock()] calls and is **not** 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; + /// Source column contexts for the phase-2 expressions; built only on the zero-match schema path (null otherwise) + private Map<String, ColumnContext> _phase2ColumnContextMap; + + // 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; + /// Whether any block has been emitted yet; drives the zero-match segment's one schema-carrying empty block + private boolean _emittedAnyBlock; + 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; + if (_emittedAnyBlock) { + return null; + } + // Zero-match segment: emit one empty block carrying the schema before signalling end-of-stream, so consumers + // never have to reconstruct a schema the segment already knows. Two-phase leaves _dataSchema null until the + // first fetch, which never happens here, so build it from column metadata alone. + if (_twoPhase && _dataSchema == null) { + _dataSchema = buildTwoPhaseDataSchema(this::resolveResultColumnContext); + } + assert _dataSchema != null; + return new SelectionResultsBlock(_dataSchema, List.of(), _comparator, _queryContext); + } + _emittedAnyBlock = true; + 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 `limit + offset` budget. Returns `null` when exhausted, and otherwise a + /// non-empty list -- empty project blocks are skipped and the row budget is checked up front, so "no rows" always + /// means "no more rows". + @Nullable + private List<Object[]> nextSortedRows() { + int remaining = _numRowsToKeep - _numRowsEmitted; + if (remaining <= 0) { + return null; + } + ValueBlock valueBlock = nextNonEmptyBlock(); + 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 + /// `limit + offset` rows by the full comparator, and return them sorted. Returns `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) { + SelectionOperatorUtils.addToPriorityQueue(row, runHeap, _numRowsToKeep); + } else { + // Run boundary: this row starts the next run, keep it for the next call. + _pendingRow = row; + break; + } + } + List<Object[]> 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 + /// `Object[_numExpressions]`. For two-phase the document id is stashed at index + /// `_numOrderByExpressions` (overwritten in the second pass). Returns `null` when the project operator + /// is exhausted. + @Nullable + private Object[] nextRow() { + if (_currentBlock == null || _currentPos >= _currentNumDocs) { + if (_projectExhausted) { + return null; + } + _currentBlock = nextNonEmptyBlock(); + 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()); + } + int rowId = _currentPos++; + return materializeRow(_currentFetcher, _currentDocIds, _currentNullBitmaps, rowId); + } + + /// Pulls the next project block carrying documents, skipping any that carry none. Returns `null` only when + /// the project operator is exhausted; an empty block means "nothing in this batch", not "end of segment", and + /// treating one as exhaustion would truncate the scan. + /// + /// No reachable doc-id-set operator emits an empty block today, so this normally returns on its first pull. The + /// skip stays because that is a convention of those implementations, not something + /// [org.apache.pinot.core.operator.BaseProjectOperator] enforces. + @Nullable + private ValueBlock nextNonEmptyBlock() { + while (true) { + ValueBlock valueBlock = _projectOperator.nextBlock(); + if (valueBlock == null) { + return null; + } + if (valueBlock.getNumDocs() > 0) { + return valueBlock; + } + } + } + + /// 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. The second pass fills + /// non-order-by values in place, so the list must at least be settable. + /// + /// It is a growable `ArrayList` rather than a fixed-size view for defensive reasons: `SelectionResultsBlock` is a + /// shared type, and `SelectionOperatorUtils.mergeWithoutOrdering()` structurally adds to the row list of the block + /// it merges into. No current consumer of this operator takes that path - the streaming combine only reads the + /// rows - so this is hardening against a future one, not a fix for a reachable failure. + private List<Object[]> drainAscending(PriorityQueue<Object[]> heap) { + int numRows = heap.size(); + // One exactly-sized backing array and no intermediate copy: drain in the heap's own (descending) order, then + // flip in place. + List<Object[]> rows = new ArrayList<>(numRows); + for (int i = 0; i < numRows; i++) { + rows.add(heap.poll()); + } + Collections.reverse(rows); + return rows; + } + + /// 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 [SelectionOrderByOperator#computePartiallyOrdered()]. + private void fetchNonOrderByColumns(List<Object[]> 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<Object[]> 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); Review Comment: **[P2] Batch small runs before rebuilding the payload-reading pipeline** A unique sorted leading key followed by an unsorted ORDER BY expression produces one-row runs. If additional payload columns are selected, every row consequently constructs and closes a projection, data fetcher, reader contexts, transform tree, and bitmap. The combine's output block size does not amortize these leaf calls. For compressed raw payload columns, this also discards the cached decompressed chunk on every row: `DataFetcher` caches the reader context only within that fetcher, `BaseChunkForwardIndexReader` caches the chunk in that context, and `ChunkReaderContext` allocates/frees a direct buffer. Consecutive rows from the same chunk repeatedly decompress it. Coalesce consecutive small sorted runs into bounded leaf blocks before fetching payloads, preserving reader-context reuse within each batch. This avoids turning setup and chunk decompression into per-row work without coupling the leaf to a one-shot operator. ########## pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java: ########## @@ -0,0 +1,655 @@ +/** + * 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.HashSet; +import java.util.List; +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.core.common.Operator; +import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator; +import org.apache.pinot.core.operator.ExplainAttributeBuilder; +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.plan.SelectionPlanNode; +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.MutableSegment; +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 `_operators`, returning +/// globally sorted rows in bounded blocks. Each segment is exposed through a [SegmentCursor] that yields that +/// segment's locally-sorted rows in order: +/// +/// - Segments physically sorted on the first order-by column are backed by +/// [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 +/// [SelectionResultsBlock]-producing operator such as `SelectionOrderByOperator`); the cursor reads that +/// one block and iterates its rows. +/// +/// A [PriorityQueue] of [SegmentCursor] ordered by the [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 `MinMaxValueContext` idea from +/// [MinMaxValueBasedSelectionOrderByCombineOperator]. A cursor is only activated (its segment acquired and first +/// block read) when the merge frontier reaches its min/max, so once `limit + offset` rows are emitted the +/// remaining segments are never acquired or read. See [#activateEligibleCursors(SegmentCursor)] for the +/// correctness argument. +/// A consuming ([MutableSegment]) segment is always given no min/max and activates up front: its bounds are a snapshot +/// that rows ingested before activation can outgrow, since its plan (and doc count) is only built at activation. +/// With null handling on, a segment whose leading column is not flagged non-null by +/// [SelectionPlanNode#isColumnFlaggedNonNull] (it has nulls, the segment predates the flag, or it is a consuming +/// segment) is given no min/max and always activates. Flagged segments keep their min/max. With +/// null handling off, min/max are used as stored. The flag is read from segment metadata, not the null value vector: +/// that vector is a mapped buffer, and probing it here would acquire segments pruning is about to skip. +/// Its effect shows in the stats: a never-activated segment scans no docs, so `numSegmentsMatched` excludes it while +/// `numSegmentsProcessed` counts every segment. Their difference also counts activated segments that matched no rows. +/// TODO: report the activated count separately; that needs a new `DataTable.MetadataKey`. +/// See https://github.com/apache/pinot/pull/19120#discussion_r3871714063 +/// +/// **Segment acquire/release lifecycle.** A cursor acquires its +/// [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 +/// [StreamingSelectionOrderByOperator] retains a buffer-backed `ValueBlock` across `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 `Object[]` (via +/// `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 [#releaseAllCursors()]. +/// +/// **Streaming paths only.** This operator is installed only where a +/// [org.apache.pinot.core.query.executor.ResultsBlockStreamer] is present (the +/// MSE leaf, driven by [org.apache.pinot.core.operator.streaming.StreamingInstanceResponseOperator], and the +/// gRPC streaming endpoint). The merge emits many bounded [SelectionResultsBlock]s from successive +/// [#getNextBlock()] calls followed by a final [MetadataResultsBlock]. On the blocking single-stage +/// path the caller expects one complete block from a single [#getNextBlock()] call, so nothing here could +/// stream: the merge would accumulate the whole `limit + offset` result while holding every activated cursor's +/// materialized block alive. That path keeps [MinMaxValueBasedSelectionOrderByCombineOperator]. +/// +/// **Threading.** This operator overrides [#start()]/[#stop()] to no-ops (other than releasing +/// segments) and runs the merge single-threaded and lazily in [#getNextBlock()] on the consumer thread; it does +/// not use the base worker-queue model. The base entry points that it replaces - [#processSegments()] and +/// [#isQuerySatisfied(SelectionResultsBlock, Object)] - are overridden to fail loud. The base +/// `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 _asc; + private final boolean _pruningEnabled; + /// Whether a cursor whose bound *ties* the merge frontier may be deferred, not only one sorting strictly + /// beyond it. Only correct under a single order-by expression: see [#sortsBeyond]. + private final boolean _deferTiedCursors; + private final int _numRowsToKeep; + private final int _blockSize; + /// Segments whose metadata marks the leading order-by column sorted; -1 when that expression is not a column + private final int _numSortedSegments; + 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. Non-null by the time any rows + /// exist to emit: every streaming child emits at least one block carrying its schema, including a zero-match one. + 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) { + // Pass a null merger: we override the consumption path entirely and never touch the base merger / worker queue. + // Every base entry point that would dereference it - processSegments() and isQuerySatisfied() - is overridden to + // throw, so a future change to the base streaming loop fails at that seam instead of with an NPE at query time. + // The executor is unused here (start() is a no-op) but is still handed to the base so that the worker-queue model + // stays constructible; passing null would trade one latent NPE for another. + super(null, operators, queryContext, executorService); + _numRowsToKeep = queryContext.getLimit() + queryContext.getOffset(); + _blockSize = queryContext.getSortedSelectionMergeBlockSize(); + + List<OrderByExpressionContext> orderByExpressions = queryContext.getOrderByExpressions(); + assert orderByExpressions != null && !orderByExpressions.isEmpty(); + OrderByExpressionContext firstOrderByExpression = orderByExpressions.get(0); + _asc = firstOrderByExpression.isAsc(); + _comparator = OrderByComparatorFactory.getComparator(orderByExpressions, queryContext.isNullHandlingEnabled()); + + // Frontier pruning needs a physical column to read min/max from. Under SortedSelectionMergeMode.ON the leading + // order-by may be an expression (forced, for benchmarking), in which case there is no such column: every cursor + // is activated up front in plan order and the merge is a plain k-way heap over materialized children. + ExpressionContext firstOrderByExpressionContext = firstOrderByExpression.getExpression(); + String firstOrderByColumn = + firstOrderByExpressionContext.getType() == ExpressionContext.Type.IDENTIFIER + ? firstOrderByExpressionContext.getIdentifier() : null; + _pruningEnabled = firstOrderByColumn != null; + // A column-0 tie only proves the segment cannot supply an *earlier* row when column 0 is the whole sort key. + // With two or more expressions a tie on column 0 can hide a row sorting earlier on column 1, so ties must + // still activate (same gate as MinMaxValueBasedSelectionOrderByCombineOperator's numOrderByExpressions == 1). + _deferTiedCursors = orderByExpressions.size() == 1; + + // 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]; + int numSortedSegments = firstOrderByColumn != null ? 0 : -1; + for (int i = 0; i < _numOperators; i++) { + Operator<BaseResultsBlock> operator = _operators.get(i); + if (firstOrderByColumn == null) { + _sortedCursors[i] = new SegmentCursor(operator, null, null); + continue; + } + DataSourceMetadata metadata = + operator.getIndexSegment().getDataSource(firstOrderByColumn, queryContext.getSchema()) + .getDataSourceMetadata(); + Comparable minValue = metadata.getMinValue(); + Comparable maxValue = metadata.getMaxValue(); + // A consuming segment's min/max is a snapshot copied when getDataSource() builds the MutableDataSource, but under + // prefetch (AcquireReleaseColumnsSegmentOperator) the child plan, and with it the doc count, is only built at + // activation. Rows indexed in between can fall outside these bounds and be emitted out of order, so treat the + // bounds as unknown and let activateEligibleCursors force-activate this cursor. + if (operator.getIndexSegment() instanceof MutableSegment) { + minValue = null; + maxValue = null; + } + // isNonNull is loaded with the segment, the same class of access as min/max. False means nulls or unknown, so + // drop the bounds and let activateEligibleCursors force-activate this cursor. Do not read the null vector here. + if (queryContext.isNullHandlingEnabled() + && !SelectionPlanNode.isColumnFlaggedNonNull(operator.getIndexSegment(), firstOrderByColumn)) { + minValue = null; + maxValue = null; + } + _sortedCursors[i] = new SegmentCursor(operator, minValue, maxValue); + if (metadata.isSorted()) { + numSortedSegments++; + } + } + _numSortedSegments = numSortedSegments; + if (firstOrderByColumn != null) { + 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 [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; + } + + /// Settings fixed at construction. Explain never drives the merge, so how many segments activate is not shown here. + /// Query-wide values are idempotent so the broker merges the node across workers only when they agree; the segment + /// counts are summed. `numSortedSegments` is metadata sortedness, not whether a child streams: a DESC scan without + /// `allowReverseOrder`, or a leading column holding nulls, still falls back to a materialized child. + @Override + protected void explainAttributes(ExplainAttributeBuilder attributeBuilder) { + super.explainAttributes(attributeBuilder); + attributeBuilder.putLongIdempotent("blockSize", _blockSize); + attributeBuilder.putBool("frontierPruning", _pruningEnabled); + attributeBuilder.putBool("deferTiedCursors", _deferTiedCursors); + attributeBuilder.putLong("numSegments", _numOperators); + if (_numSortedSegments >= 0) { + attributeBuilder.putLong("numSortedSegments", _numSortedSegments); + } + } + + /// Override to a no-op: the merge is single-threaded and lazy in [#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 ([#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"); + } + + /// Only reachable from the base streaming loop, which this operator replaces; the base implementation would + /// dereference the null merger. Fail loud rather than NPE if a future change routes back through it. + @Override + protected boolean isQuerySatisfied(SelectionResultsBlock resultsBlock, Object tracker) { + throw new IllegalStateException( + "StreamingSelectionOrderByCombineOperator tracks its own limit; isQuerySatisfied() 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()); Review Comment: **[P2] Accumulate consumer CPU and allocations across streamed blocks** This operator performs scans and merging on the consumer thread, while its no-op `start()` leaves inherited worker resource counters at zero. `StreamingInstanceResponseOperator` calls `InstanceResponseOperator.getBaseBlock()` for every data block and again for terminal metadata. That method measures each call independently and overwrites its CPU/allocation totals rather than accumulating them. The terminal call therefore replaces previous scan/merge measurements with the cost of producing metadata. Successful nonempty queries substantially underreport consumer-thread CPU and allocation measurements in server metrics and MSE statistics, even with just one data block followed by metadata. Accumulate consumer-thread measurements across the streaming response lifecycle, separately from the existing cumulative worker counters. A full response-operator test should verify the final totals across several data blocks. ########## pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java: ########## @@ -0,0 +1,575 @@ +/** + * 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.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.function.Function; +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.common.utils.HashUtil; +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.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.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 [SelectionOrderByOperator] (which materializes the segment's whole top-K in a single block) this +/// operator emits one globally-sorted [SelectionResultsBlock] per [#getNextBlock()] call and returns +/// `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 `projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))`). +/// +/// Every instance emits at least one block before `null`: a segment matching no rows emits a single empty block +/// carrying the [DataSchema], so consumers never have to reconstruct a schema the segment already knows. +/// +/// It runs in one of two emission modes: +/// +/// - **No tail to sort** (`numSortedExpressions == numOrderByExpressions`, e.g. `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 `limit + offset` budget). +/// - **Tail to sort** (`numSortedExpressions < numOrderByExpressions`, e.g. +/// `ORDER BY sorted, other`): +/// each call reads forward until the first order-by value changes (a primary-value "run"), retains the run's top +/// `limit + offset` rows by the full comparator, and emits them sorted. This bounds the in-memory run buffer to +/// `limit + offset` rows even when the first order-by column is near-constant (very low cardinality). +/// +/// Like [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 [#getNextBlock()] calls and is **not** 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; + /// Source column contexts for the phase-2 expressions; built only on the zero-match schema path (null otherwise) + private Map<String, ColumnContext> _phase2ColumnContextMap; + + // 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; + /// Whether any block has been emitted yet; drives the zero-match segment's one schema-carrying empty block + private boolean _emittedAnyBlock; + 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; + if (_emittedAnyBlock) { + return null; + } + // Zero-match segment: emit one empty block carrying the schema before signalling end-of-stream, so consumers + // never have to reconstruct a schema the segment already knows. Two-phase leaves _dataSchema null until the + // first fetch, which never happens here, so build it from column metadata alone. + if (_twoPhase && _dataSchema == null) { + _dataSchema = buildTwoPhaseDataSchema(this::resolveResultColumnContext); + } + assert _dataSchema != null; + return new SelectionResultsBlock(_dataSchema, List.of(), _comparator, _queryContext); + } + _emittedAnyBlock = true; + 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 `limit + offset` budget. Returns `null` when exhausted, and otherwise a + /// non-empty list -- empty project blocks are skipped and the row budget is checked up front, so "no rows" always + /// means "no more rows". + @Nullable + private List<Object[]> nextSortedRows() { + int remaining = _numRowsToKeep - _numRowsEmitted; + if (remaining <= 0) { + return null; + } + ValueBlock valueBlock = nextNonEmptyBlock(); + 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 + /// `limit + offset` rows by the full comparator, and return them sorted. Returns `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) { + SelectionOperatorUtils.addToPriorityQueue(row, runHeap, _numRowsToKeep); Review Comment: **[P2] Bound each run's heap by the remaining row budget** After earlier runs emit `K-1` rows of a `LIMIT K` query, the next run can contribute only one row. Both heap insertions still retain up to `K` rows, and `drainAscending()` sorts them before the result is trimmed to one. For a large run of M rows with decreasing secondary values, this requires `O(M log K)` work instead of linear minimum selection, plus retaining and sorting K candidates. The existing `LinearSelectionOrderByOperator` bounds its final partition by the remaining capacity. Pass `remaining` to both `addToPriorityQueue()` calls (here and the initial insertion at line 293). ########## pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java: ########## @@ -0,0 +1,575 @@ +/** + * 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.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.function.Function; +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.common.utils.HashUtil; +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.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.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 [SelectionOrderByOperator] (which materializes the segment's whole top-K in a single block) this +/// operator emits one globally-sorted [SelectionResultsBlock] per [#getNextBlock()] call and returns +/// `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 `projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))`). +/// +/// Every instance emits at least one block before `null`: a segment matching no rows emits a single empty block +/// carrying the [DataSchema], so consumers never have to reconstruct a schema the segment already knows. +/// +/// It runs in one of two emission modes: +/// +/// - **No tail to sort** (`numSortedExpressions == numOrderByExpressions`, e.g. `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 `limit + offset` budget). +/// - **Tail to sort** (`numSortedExpressions < numOrderByExpressions`, e.g. +/// `ORDER BY sorted, other`): +/// each call reads forward until the first order-by value changes (a primary-value "run"), retains the run's top +/// `limit + offset` rows by the full comparator, and emits them sorted. This bounds the in-memory run buffer to +/// `limit + offset` rows even when the first order-by column is near-constant (very low cardinality). +/// +/// Like [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 [#getNextBlock()] calls and is **not** 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; + /// Source column contexts for the phase-2 expressions; built only on the zero-match schema path (null otherwise) + private Map<String, ColumnContext> _phase2ColumnContextMap; + + // 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; + /// Whether any block has been emitted yet; drives the zero-match segment's one schema-carrying empty block + private boolean _emittedAnyBlock; + 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); Review Comment: **[P2] Allocate the run heap lazily** Every leaf allocates a `PriorityQueue` backing array of up to 10,000 references, including the fully sorted path where the heap is never used. With prefetch disabled, `CombinePlanNode` constructs all leaves before frontier pruning. For 5,000 segments and `LIMIT >= 10,000`, this reserves 50 million unused reference slots—approximately 200 MB with compressed references—even if only one segment is activated. The existing linear operator defers its result allocation until execution. Allocate the heap on the first `nextRun()` call and avoid constructing run-only comparators for the fully sorted path. The memory figure is an allocation estimate, not a measured benchmark. -- 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]
