rohityadav1993 commented on code in PR #19120:
URL: https://github.com/apache/pinot/pull/19120#discussion_r4044526362
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java:
##########
@@ -88,11 +89,36 @@ public Operator<SelectionResultsBlock> run() {
maxDocsPerCall = Math.min(limit + _queryContext.getOffset(),
DocIdSetPlanNode.MAX_DOC_PER_CALL);
}
- BaseProjectOperator<?> projectOperator = getSortedByProject(expressions,
maxDocsPerCall, orderByExpressions);
boolean asc = orderByExpressions.get(0).isAsc();
// Remember that we cannot use asc == projectOperator.isAscending()
because empty operators are considered
// both ascending and descending
DocIdOrderedOperator.DocIdOrder queryOrder =
DocIdOrderedOperator.DocIdOrder.fromAsc(asc);
+
+ // Opt-in streaming path: emit one globally-sorted block at a time so a
downstream k-way-merge combine can pull
+ // lazily. Only build it when the first order-by column is an identifier
(kept consistent with the combine-side
+ // gate) and the forward-scan project is order-compatible; the
DESC-incompatible sorted case still falls back to
+ // the materialized SelectionPartiallyOrderedByDescOperation below so
global order stays correct.
+ if (_queryContext.isSortedSelectionMergeEnabled()
+ && orderByExpressions.get(0).getExpression().getType() ==
ExpressionContext.Type.IDENTIFIER) {
+ // When there are non-order-by output expressions, only fetch the
order-by expressions during the forward scan
+ // (the streaming operator fetches the rest in a second pass);
otherwise fetch all expressions.
+ List<ExpressionContext> projectExpressions = expressions;
+ if (expressions.size() > numOrderByExpressions) {
+ projectExpressions = new ArrayList<>(numOrderByExpressions);
+ for (OrderByExpressionContext orderByExpression :
orderByExpressions) {
+ projectExpressions.add(orderByExpression.getExpression());
+ }
+ }
+ BaseProjectOperator<?> streamingProjectOperator =
+ getSortedByProject(projectExpressions, maxDocsPerCall,
orderByExpressions);
+ if (streamingProjectOperator.isCompatibleWith(queryOrder)) {
+ return new StreamingSelectionOrderByOperator(_indexSegment,
_queryContext, expressions,
+ streamingProjectOperator, sortedColumnsPrefixSize);
+ }
+ // DESC-incompatible: fall through to the materialized fallback
(rebuilds the project over all expressions).
Review Comment:
Only the duplicate build is fixed in f3f2db3 - the leak is not. The DESC
fall-through now reuses streamingProjectOperator when it already covers the
full expression list; the wider order-by-prefix case still builds a second one.
On the leak: ProjectPlanNode#run already has a pre-existing "TODO: figure
out a way to close this operator" covering every call site in the module, so a
narrow fix here would not actually close it - left as-is.
Tests: testDescIncompatibleFallbackBuildsTheProjectOnce (counts data source
lookups, 4 vs 3 baseline) and
testDescIncompatibleFallbackMatchesTheMaterializedPlan (pins schema/rows across
reuse, wider-rebuild, zero-match cases).
[addressed by agent]
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java:
##########
@@ -0,0 +1,514 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.operator.query;
+
+import com.google.common.base.CaseFormat;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.OrderByExpressionContext;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.core.common.BlockValSet;
+import org.apache.pinot.core.common.Operator;
+import org.apache.pinot.core.common.RowBasedBlockValueFetcher;
+import org.apache.pinot.core.operator.BaseOperator;
+import org.apache.pinot.core.operator.BaseProjectOperator;
+import org.apache.pinot.core.operator.BitmapDocIdSetOperator;
+import org.apache.pinot.core.operator.ColumnContext;
+import org.apache.pinot.core.operator.ExecutionStatistics;
+import org.apache.pinot.core.operator.ExplainAttributeBuilder;
+import org.apache.pinot.core.operator.ProjectionOperator;
+import org.apache.pinot.core.operator.ProjectionOperatorUtils;
+import org.apache.pinot.core.operator.blocks.ValueBlock;
+import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
+import org.apache.pinot.core.operator.transform.TransformOperator;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.query.selection.SelectionOperatorUtils;
+import org.apache.pinot.core.query.utils.OrderByComparatorFactory;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.spi.query.QueryScanCostContext;
+import org.roaringbitmap.RoaringBitmap;
+
+
+/// Lazy, incremental selection ORDER BY operator for segments that are
physically sorted on the first order-by column.
+///
+/// Unlike {@link SelectionOrderByOperator} (which materializes the segment's
whole top-K in a single block) this
+/// operator emits one globally-sorted {@link SelectionResultsBlock} per
{@link #getNextBlock()} call and returns
+/// {@code null} when the segment is exhausted, so that a downstream
k-way-merge combine operator can pull from many
+/// segments lazily and stop early. It relies on the underlying project
operator iterating the first order-by column in
+/// the query order (the caller must guarantee {@code
projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))}).
+///
+/// It runs in one of two emission modes:
+///
+/// - **No tail to sort** ({@code numSortedExpressions ==
numOrderByExpressions}, e.g. {@code ORDER BY sorted}):
+/// rows already arrive from the project operator in final order, so each
call emits the next project block
+/// (trimmed to the remaining {@code limit + offset} budget).
+/// - **Tail to sort** ({@code numSortedExpressions < numOrderByExpressions},
e.g.
+/// {@code ORDER BY sorted, other}):
+/// each call reads forward until the first order-by value changes (a
primary-value "run"), retains the run's top
+/// {@code limit + offset} rows by the full comparator, and emits them
sorted. This bounds the in-memory run buffer to
+/// {@code limit + offset} rows even when the first order-by column is
near-constant (very low cardinality).
+///
+/// Like {@link SelectionOrderByOperator} it preserves the two-phase
projection optimization: when there are output
+/// expressions that are not order-by expressions, the forward scan only
fetches the order-by expressions plus the
+/// document id, and the non-order-by expressions are fetched in a second pass
over the retained document ids of each
+/// emitted block.
+///
+/// This operator is stateful across {@link #getNextBlock()} calls and is
**not** thread-safe; a single consumer
+/// must drive it.
+public class StreamingSelectionOrderByOperator extends
BaseOperator<SelectionResultsBlock> {
+ private static final String EXPLAIN_NAME = "SELECT_ORDERBY_STREAMING";
+
+ private final IndexSegment _indexSegment;
+ private final QueryContext _queryContext;
+ private final boolean _nullHandlingEnabled;
+ /// Deduped order-by expressions followed by output expressions from
SelectionOperatorUtils.extractExpressions()
+ private final List<ExpressionContext> _expressions;
+ private final BaseProjectOperator<?> _projectOperator;
+ private final List<OrderByExpressionContext> _orderByExpressions;
+ private final ColumnContext[] _orderByColumnContexts;
+ private final int _numExpressions;
+ private final int _numOrderByExpressions;
+ private final int _numRowsToKeep;
+ /// Whether there are output expressions that are not order-by expressions
(requires the two-phase fetch)
+ private final boolean _twoPhase;
+ /// Whether the order-by has an unsorted tail that must be sorted in memory
per run
+ private final boolean _tailToSort;
+ /// Expressions fetched during the forward scan: order-by expressions only
when two-phase, otherwise all expressions
+ private final List<ExpressionContext> _phase1Expressions;
+ private final int _numPhase1Columns;
+ private final Comparator<Object[]> _comparator;
+ /// Compares only the first order-by column; used to detect primary-value
run boundaries
+ private final Comparator<Object[]> _primaryComparator;
+ /// Pre-allocated run heap (cleared and reused each nextRun() call to avoid
per-run allocation)
+ private final Comparator<Object[]> _reversedComparator;
+ private final PriorityQueue<Object[]> _runHeap;
+
+ // Pre-computed invariants for the two-phase fetch (null when single-phase)
+ private final List<ExpressionContext> _nonOrderByExpressions;
+ private final Map<String, DataSource> _phase2DataSourceMap;
+ private final int _phase2NumColumns;
+
+ /// Lazily built and cached; for two-phase it requires the transform
operator's result column contexts
+ private DataSchema _dataSchema;
+
+ // Forward-scan cursor state (used by the tail-to-sort mode)
+ private ValueBlock _currentBlock;
+ private RowBasedBlockValueFetcher _currentFetcher;
+ private int[] _currentDocIds;
+ private RoaringBitmap[] _currentNullBitmaps;
+ private int _currentNumDocs;
+ private int _currentPos;
+ /// One-row lookahead: the first row of the next run, stashed when a run
boundary is crossed
+ private Object[] _pendingRow;
+ private boolean _projectExhausted;
+
+ private boolean _exhausted;
+ private int _numRowsEmitted;
+ private int _numDocsScanned = 0;
+ private long _numEntriesScannedPostFilter = 0;
+
+ public StreamingSelectionOrderByOperator(IndexSegment indexSegment,
QueryContext queryContext,
+ List<ExpressionContext> expressions, BaseProjectOperator<?>
projectOperator, int numSortedExpressions) {
+ _indexSegment = indexSegment;
+ _queryContext = queryContext;
+ _nullHandlingEnabled = queryContext.isNullHandlingEnabled();
+ _expressions = expressions;
+ _projectOperator = projectOperator;
+
+ _orderByExpressions = queryContext.getOrderByExpressions();
+ assert _orderByExpressions != null;
+ _numExpressions = expressions.size();
+ _numOrderByExpressions = _orderByExpressions.size();
+ _orderByColumnContexts = new ColumnContext[_numOrderByExpressions];
+ for (int i = 0; i < _numOrderByExpressions; i++) {
+ ExpressionContext expression =
_orderByExpressions.get(i).getExpression();
+ _orderByColumnContexts[i] =
_projectOperator.getResultColumnContext(expression);
+ }
+
+ _numRowsToKeep = queryContext.getOffset() + queryContext.getLimit();
+ _twoPhase = _numExpressions > _numOrderByExpressions;
+ _tailToSort = numSortedExpressions < _numOrderByExpressions;
+ _comparator =
+ OrderByComparatorFactory.getComparator(_orderByExpressions,
_orderByColumnContexts, _nullHandlingEnabled);
+ // The first order-by column is the physically sorted column, so it never
contains nulls on this path; comparing
+ // only index 0 is enough to detect when one primary-value run ends and
the next begins.
+ _primaryComparator =
+ OrderByComparatorFactory.getComparator(_orderByExpressions,
_orderByColumnContexts, _nullHandlingEnabled, 0, 1);
+ _reversedComparator = _comparator.reversed();
+ _runHeap = new PriorityQueue<>(
+ Math.min(_numRowsToKeep,
SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY), _reversedComparator);
+
+ if (_twoPhase) {
+ _phase1Expressions = new ArrayList<>(_numOrderByExpressions);
+ for (OrderByExpressionContext orderByExpression : _orderByExpressions) {
+ _phase1Expressions.add(orderByExpression.getExpression());
+ }
+ _nonOrderByExpressions = _expressions.subList(_numOrderByExpressions,
_numExpressions);
+ Set<String> columns = new HashSet<>();
+ for (ExpressionContext expressionContext : _nonOrderByExpressions) {
+ expressionContext.getColumns(columns);
+ }
+ _phase2NumColumns = columns.size();
+ _phase2DataSourceMap = new HashMap<>();
+ for (String column : columns) {
+ _phase2DataSourceMap.put(column, _indexSegment.getDataSource(column,
_queryContext.getSchema()));
+ }
+ } else {
+ _phase1Expressions = _expressions;
+ _nonOrderByExpressions = null;
+ _phase2NumColumns = 0;
+ _phase2DataSourceMap = null;
+ // Single-phase: all output expressions are order-by expressions, so
their types are known up front.
+ _dataSchema = buildSinglePhaseDataSchema();
+ }
+ _numPhase1Columns = _phase1Expressions.size();
+ }
+
+ @Override
+ protected SelectionResultsBlock getNextBlock() {
+ if (_exhausted) {
+ return null;
+ }
+ List<Object[]> rows = _tailToSort ? nextRun() : nextSortedRows();
+ if (rows == null || rows.isEmpty()) {
+ _exhausted = true;
+ return null;
+ }
+ if (_twoPhase) {
+ fetchNonOrderByColumns(rows);
+ }
+ // Single-phase builds the schema in the constructor; two-phase builds it
during fetchNonOrderByColumns above.
+ assert _dataSchema != null;
+ return new SelectionResultsBlock(_dataSchema, rows, _comparator,
_queryContext);
+ }
+
+ /// No-tail-to-sort mode: the project operator already returns rows in final
order, so emit the next project block,
+ /// trimmed to the remaining {@code limit + offset} budget. Returns {@code
null} when exhausted.
+ @Nullable
+ private List<Object[]> nextSortedRows() {
+ int remaining = _numRowsToKeep - _numRowsEmitted;
+ if (remaining <= 0) {
+ return null;
+ }
+ ValueBlock valueBlock = _projectOperator.nextBlock();
+ if (valueBlock == null) {
+ return null;
+ }
+ int numDocsFetched = valueBlock.getNumDocs();
+ BlockValSet[] blockValSets = new BlockValSet[_numPhase1Columns];
+ for (int i = 0; i < _numPhase1Columns; i++) {
+ blockValSets[i] = valueBlock.getBlockValueSet(_phase1Expressions.get(i));
+ }
+ RowBasedBlockValueFetcher blockValueFetcher = new
RowBasedBlockValueFetcher(blockValSets);
+ int[] docIds = _twoPhase ? valueBlock.getDocIds() : null;
+ RoaringBitmap[] nullBitmaps = null;
+ if (_nullHandlingEnabled) {
+ nullBitmaps = new RoaringBitmap[_numPhase1Columns];
+ for (int i = 0; i < _numPhase1Columns; i++) {
+ nullBitmaps[i] = blockValSets[i].getNullBitmap();
+ }
+ }
+ _numDocsScanned += numDocsFetched;
+ _numEntriesScannedPostFilter += (long) numDocsFetched *
_projectOperator.getNumColumnsProjected();
+ reportScanCost(numDocsFetched, (long) numDocsFetched *
_projectOperator.getNumColumnsProjected());
+
+ // Rows arrive sorted; we only need the first 'remaining' of them globally.
+ int numRows = Math.min(numDocsFetched, remaining);
+ List<Object[]> rows = new ArrayList<>(numRows);
+ for (int i = 0; i < numRows; i++) {
+ rows.add(materializeRow(blockValueFetcher, docIds, nullBitmaps, i));
+ }
+ _numRowsEmitted += rows.size();
+ return rows;
+ }
+
+ /// Tail-to-sort mode: read forward until the first order-by value changes,
retain the run's top
+ /// {@code limit + offset} rows by the full comparator, and return them
sorted. Returns {@code null} when
+ /// exhausted.
+ @Nullable
+ private List<Object[]> nextRun() {
+ int remaining = _numRowsToKeep - _numRowsEmitted;
+ if (remaining <= 0) {
+ return null;
+ }
+ if (_pendingRow == null) {
+ _pendingRow = nextRow();
+ if (_pendingRow == null) {
+ return null;
+ }
+ }
+ PriorityQueue<Object[]> runHeap = _runHeap;
+ runHeap.clear();
+ Object[] runFirstRow = _pendingRow;
+ SelectionOperatorUtils.addToPriorityQueue(_pendingRow, runHeap,
_numRowsToKeep);
+ _pendingRow = null;
+ Object[] row;
+ while ((row = nextRow()) != null) {
+ if (_primaryComparator.compare(row, runFirstRow) == 0) {
+ 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
+ /// {@code Object[_numExpressions]}. For two-phase the document id is
stashed at index
+ /// {@code _numOrderByExpressions} (overwritten in the second pass). Returns
{@code null} when the project operator
+ /// is exhausted.
+ @Nullable
+ private Object[] nextRow() {
+ while (true) {
+ if (_currentBlock == null || _currentPos >= _currentNumDocs) {
+ if (_projectExhausted) {
+ return null;
+ }
+ _currentBlock = _projectOperator.nextBlock();
+ if (_currentBlock == null) {
+ _projectExhausted = true;
+ return null;
+ }
+ BlockValSet[] blockValSets = new BlockValSet[_numPhase1Columns];
+ for (int i = 0; i < _numPhase1Columns; i++) {
+ blockValSets[i] =
_currentBlock.getBlockValueSet(_phase1Expressions.get(i));
+ }
+ _currentFetcher = new RowBasedBlockValueFetcher(blockValSets);
+ _currentNumDocs = _currentBlock.getNumDocs();
+ _currentDocIds = _twoPhase ? _currentBlock.getDocIds() : null;
+ if (_nullHandlingEnabled) {
+ _currentNullBitmaps = new RoaringBitmap[_numPhase1Columns];
+ for (int i = 0; i < _numPhase1Columns; i++) {
+ _currentNullBitmaps[i] = blockValSets[i].getNullBitmap();
+ }
+ }
+ _currentPos = 0;
+ _numDocsScanned += _currentNumDocs;
+ _numEntriesScannedPostFilter += (long) _currentNumDocs *
_projectOperator.getNumColumnsProjected();
+ reportScanCost(_currentNumDocs, (long) _currentNumDocs *
_projectOperator.getNumColumnsProjected());
+ if (_currentNumDocs == 0) {
+ _currentBlock = null;
+ continue;
+ }
+ }
+ int rowId = _currentPos++;
+ return materializeRow(_currentFetcher, _currentDocIds,
_currentNullBitmaps, rowId);
+ }
+ }
+
+ /// Materializes a single phase-1 row (deep-copied out of the value block
buffers) from the given fetcher.
+ private Object[] materializeRow(RowBasedBlockValueFetcher fetcher, @Nullable
int[] docIds,
+ @Nullable RoaringBitmap[] nullBitmaps, int rowId) {
+ Object[] row = new Object[_numExpressions];
+ fetcher.getRow(rowId, row, 0);
+ if (_twoPhase) {
+ row[_numOrderByExpressions] = docIds[rowId];
+ }
+ if (_nullHandlingEnabled) {
+ for (int colId = 0; colId < _numPhase1Columns; colId++) {
+ if (nullBitmaps[colId] != null && nullBitmaps[colId].contains(rowId)) {
+ row[colId] = null;
+ }
+ }
+ }
+ return row;
+ }
+
+ /// Drains a max-heap (created with the reversed comparator) into an
ascending list, mutable so the second pass can
+ /// fill non-order-by values in place.
+ private List<Object[]> drainAscending(PriorityQueue<Object[]> heap) {
+ int numRows = heap.size();
+ Object[][] sortedRows = new Object[numRows][];
+ for (int i = numRows - 1; i >= 0; i--) {
+ sortedRows[i] = heap.poll();
+ }
+ return Arrays.asList(sortedRows);
Review Comment:
Done in f3f2db3 - drainAscending now returns a real, growable ArrayList.
Correction: the addAll/mergeWithOrdering throw site you cited isn't
reachable from this operator today (mergeWithOrdering swaps references via
setRows(); only mergeWithoutOrdering mutates, for non-ORDER-BY queries).
Treating this as hardening rather than a live-bug fix, and said so in the
javadoc.
New test testRunPathEmitsGrowableRowLists appends to every emitted row list;
throws pre-fix, passes post-fix.
[addressed by agent]
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java:
##########
@@ -0,0 +1,543 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.operator.combine;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.OrderByExpressionContext;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.common.Operator;
+import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator;
+import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
+import org.apache.pinot.core.operator.blocks.results.MetadataResultsBlock;
+import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
+import org.apache.pinot.core.operator.query.StreamingSelectionOrderByOperator;
+import org.apache.pinot.core.operator.streaming.BaseStreamingCombineOperator;
+import org.apache.pinot.core.operator.transform.function.TransformFunction;
+import
org.apache.pinot.core.operator.transform.function.TransformFunctionFactory;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.query.selection.SelectionOperatorUtils;
+import org.apache.pinot.core.query.utils.OrderByComparatorFactory;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.segment.spi.datasource.DataSourceMetadata;
+import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.exception.QueryErrorMessage;
+import org.apache.pinot.spi.query.QueryThreadContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/// Streaming, lazy combine operator for selection ORDER BY queries whose
first order-by expression is an identifier.
+///
+/// It performs an incremental k-way heap merge across the per-segment
operators in {@code _operators}, returning
+/// globally sorted rows in bounded blocks. Each segment is exposed through a
{@link SegmentCursor} that yields that
+/// segment's locally-sorted rows in order:
+///
+/// - Segments physically sorted on the first order-by column are backed by
+/// {@link StreamingSelectionOrderByOperator}, which is pulled lazily one
run/block at a time.
+/// - Other (e.g. consuming/unsorted) segments are backed by a single
materialized top-K block (any
+/// {@link SelectionResultsBlock}-producing operator such as {@code
SelectionOrderByOperator}); the cursor reads that
+/// one block and iterates its rows.
+///
+/// A {@link PriorityQueue} of {@link SegmentCursor} ordered by the {@link
OrderByComparatorFactory} comparator on
+/// each
+/// cursor's current head row drives the merge with an
at-most-one-head-per-active-segment invariant (the heap holds the
+/// cursors themselves, never all rows, which would degenerate into a full
heap-sort that materializes everything). Each
+/// cycle pops the global-min cursor, appends its head to the current output
block, advances that one cursor by a single
+/// row, and re-offers it if it still has a head.
+///
+/// **Min/max lazy segment activation (pruning).** Cursors are sorted by the
first order-by column's min value
+/// (ASC) / max value (DESC) reusing the {@code MinMaxValueContext} idea from
+/// {@link MinMaxValueBasedSelectionOrderByCombineOperator}. A cursor is only
activated (its segment acquired and first
+/// block read) when the merge frontier reaches its min/max, so once {@code
limit + offset} rows are emitted the
+/// remaining segments are never acquired or read. See {@link
#activateEligibleCursors()} for the correctness argument.
+/// Pruning is disabled when null handling is enabled (an unsorted segment's
first order-by column may then contain
+/// nulls whose ordering position the raw min/max cannot capture), in which
case every segment is activated.
+///
+/// **Segment acquire/release lifecycle.** A cursor acquires its
+/// {@link AcquireReleaseColumnsSegmentOperator} on activation and releases it
only when its child operator is fully
+/// drained (acquire-on-activate / release-on-exhaust), rather than per run.
This is intentional: the backing
+/// {@link StreamingSelectionOrderByOperator} retains a buffer-backed {@code
ValueBlock} across {@code nextBlock()}
+/// calls in its tail-to-sort mode, so releasing between interleaved runs
could read segment buffers after a release
+/// under prefetch. Holding the acquire for the cursor's lifetime guarantees
no release happens between a cursor's
+/// own reads; min/max pruning bounds the number of simultaneously-active
(acquired) segments to the merge frontier.
+/// The rows handed out by the child operators are already deep-copied to heap
{@code Object[]} (via
+/// {@code RowBasedBlockValueFetcher}), so they remain valid after the segment
is released. Any cursors still
+/// acquired when the merge ends early (LIMIT reached) or errors out are
released via {@link #releaseAllCursors()}.
+///
+/// **Streaming vs single-block.** When {@code _streaming} is {@code true}
(MSE leaf path, driven by
+/// {@link
org.apache.pinot.core.operator.streaming.StreamingInstanceResponseOperator})
the merge emits many bounded
+/// {@link SelectionResultsBlock}s from successive {@link #getNextBlock()}
calls followed by a final
+/// {@link MetadataResultsBlock}. When {@code false} (classic single-stage
path) the merge runs to completion and the
+/// first {@link #getNextBlock()} call returns a single block with execution
stats attached.
+///
+/// **Threading.** This operator overrides {@link #start()}/{@link #stop()} to
no-ops (other than releasing
+/// segments) and runs the merge single-threaded and lazily in {@link
#getNextBlock()} on the consumer thread; it does
+/// not use the base worker-queue model, and {@link #processSegments()} is
overridden to fail loud. The base
+/// {@code Phaser} (which exists only to fence worker threads against segment
release) is intentionally bypassed because
+/// all child/segment access is synchronous on the single consumer thread that
holds the segment references; no async
+/// work may be introduced here without restoring that fence. The instance is
single-use (driven once to completion) and
+/// is not thread-safe.
+@SuppressWarnings({"rawtypes", "unchecked"})
+public class StreamingSelectionOrderByCombineOperator extends
BaseStreamingCombineOperator<SelectionResultsBlock> {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StreamingSelectionOrderByCombineOperator.class);
+ private static final String EXPLAIN_NAME =
"COMBINE_SELECT_ORDERBY_STREAMING";
+
+ private final boolean _streaming;
+ private final boolean _asc;
+ private final boolean _pruningEnabled;
+ private final int _numRowsToKeep;
+ private final int _blockSize;
+ private final Comparator<Object[]> _comparator;
+ private final SegmentCursor[] _sortedCursors;
+ private final PriorityQueue<SegmentCursor> _priorityQueue;
+
+ // Merge progress (single-threaded; mutated only by the consumer thread
driving getNextBlock())
+ private int _nextToActivate;
+ private int _numRowsEmitted;
+ private List<Object[]> _outputRows;
+ private boolean _done;
+ /// Captured from the first child block seen; all child blocks share the
same schema
+ private DataSchema _dataSchema;
+ /// Deduplicated MERGE_RESPONSE errors for segment blocks dropped on schema
mismatch; null until the first mismatch
+ @Nullable
+ private Set<String> _dataSchemaMismatchErrors;
+ /// Subset of the above not yet attached to an emitted block; drained on
each attach
+ @Nullable
+ private List<String> _unreportedDataSchemaMismatchErrors;
+
+ public StreamingSelectionOrderByCombineOperator(List<Operator> operators,
QueryContext queryContext,
+ ExecutorService executorService, boolean streaming) {
+ // Pass a null merger: we override the consumption path entirely and never
touch the base merger / worker queue.
+ super(null, operators, queryContext, executorService);
+ _streaming = streaming;
+ _numRowsToKeep = queryContext.getLimit() + queryContext.getOffset();
+ // Streaming mode flushes bounded blocks; single-stage mode flushes once
at the end as a single block.
+ _blockSize = streaming ? queryContext.getSortedSelectionMergeBlockSize() :
Integer.MAX_VALUE;
+ _pruningEnabled = !queryContext.isNullHandlingEnabled();
+
+ List<OrderByExpressionContext> orderByExpressions =
queryContext.getOrderByExpressions();
+ assert orderByExpressions != null && !orderByExpressions.isEmpty();
+ OrderByExpressionContext firstOrderByExpression =
orderByExpressions.get(0);
+ assert firstOrderByExpression.getExpression().getType() ==
ExpressionContext.Type.IDENTIFIER;
+ _asc = firstOrderByExpression.isAsc();
+ String firstOrderByColumn =
firstOrderByExpression.getExpression().getIdentifier();
+ _comparator = OrderByComparatorFactory.getComparator(orderByExpressions,
queryContext.isNullHandlingEnabled());
+
+ // Build one cursor per segment operator and read its first order-by
column min/max for lazy activation ordering.
+ // Reading DataSourceMetadata does not touch column buffers, so no segment
acquire is needed here (mirrors
+ // MinMaxValueBasedSelectionOrderByCombineOperator).
+ _sortedCursors = new SegmentCursor[_numOperators];
+ for (int i = 0; i < _numOperators; i++) {
+ Operator<BaseResultsBlock> operator = _operators.get(i);
+ DataSourceMetadata metadata =
+ operator.getIndexSegment().getDataSource(firstOrderByColumn,
queryContext.getSchema())
+ .getDataSourceMetadata();
+ _sortedCursors[i] = new SegmentCursor(operator, metadata.getMinValue(),
metadata.getMaxValue());
+ }
+ sortCursorsByMinMax();
+
+ _priorityQueue = new PriorityQueue<>(Math.max(1, _numOperators),
+ (o1, o2) -> _comparator.compare(o1.currentHead(), o2.currentHead()));
+ _outputRows = newOutputList();
+ }
+
+ /// Sorts the cursors so the merge can activate them lazily in frontier
order: ascending by the column min value for
+ /// ASC, descending by the column max value for DESC. Cursors without a
min/max are placed first because they must
+ /// always be processed (mirrors {@link
MinMaxValueBasedSelectionOrderByCombineOperator}).
+ private void sortCursorsByMinMax() {
+ if (_asc) {
+ Arrays.sort(_sortedCursors, (o1, o2) -> {
+ if (o1._minValue == null) {
+ return o2._minValue == null ? 0 : -1;
+ }
+ if (o2._minValue == null) {
+ return 1;
+ }
+ return o1._minValue.compareTo(o2._minValue);
+ });
+ } else {
+ Arrays.sort(_sortedCursors, (o1, o2) -> {
+ if (o1._maxValue == null) {
+ return o2._maxValue == null ? 0 : -1;
+ }
+ if (o2._maxValue == null) {
+ return 1;
+ }
+ return o2._maxValue.compareTo(o1._maxValue);
+ });
+ }
+ }
+
+ @Override
+ public String toExplainString() {
+ return EXPLAIN_NAME;
+ }
+
+ /// Override to a no-op: the merge is single-threaded and lazy in {@link
#getNextBlock()}, so we do not spin up the
+ /// base worker threads / blocking-queue model.
+ @Override
+ public void start() {
Review Comment:
Done in f3f2db3. isQuerySatisfied() now throws IllegalStateException,
matching processSegments()'s existing fail-loud behavior.
Kept the executor param rather than dropping it - removing it would just
pass null up to BaseCombineOperator._executorService, the same latent-NPE shape
one level up. Comment on the constructor explains why.
New test testBaseWorkerEntryPointsFailLoud covers both seams.
[addressed by agent]
##########
pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java:
##########
@@ -170,13 +164,85 @@ public void testCancelGroupByOrderByCombineOperator() {
testCancelCombineOperator(combineOperator, ready);
}
+ @Test
+ public void testCancelStreamingSelectionOrderByCombineOperator() {
+ CountDownLatch ready = new CountDownLatch(1);
+ List<Operator> operators = getOperators(ready, minMaxSegmentSupplier());
+ QueryContext queryContext =
QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY
column");
+ queryContext.setEndTimeMs(System.currentTimeMillis() + 10000);
+ // Single-stage mode: nextBlock() drives the whole merge synchronously on
the (cancellable) caller thread.
+ StreamingSelectionOrderByCombineOperator combineOperator =
+ new StreamingSelectionOrderByCombineOperator(operators, queryContext,
_executorService, false);
+ testCancelCombineOperator(combineOperator, ready, operators);
+ }
+
+ @Test
+ public void
testCancelStreamingSelectionOrderByCombineOperatorStreamingMode() {
+ CountDownLatch ready = new CountDownLatch(1);
+ List<Operator> operators = getOperators(ready, minMaxSegmentSupplier());
+ QueryContext queryContext =
QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY
column");
+ queryContext.setEndTimeMs(System.currentTimeMillis() + 10000);
+ // Streaming (MSE-leaf) mode: the first getNextBlock() still drives the
merge on the caller thread, so the same
+ // interrupt-on-cancel path applies and must surface an
ExceptionResultsBlock.
+ StreamingSelectionOrderByCombineOperator combineOperator =
+ new StreamingSelectionOrderByCombineOperator(operators, queryContext,
_executorService, true);
+ testCancelCombineOperator(combineOperator, ready, operators);
+ }
+
+ /// The merge loop drains its heap on the caller thread and, for
single-block cursors, may never re-enter a child
+ /// operator, so it must check the query deadline itself. With an
already-expired deadline the operator must surface a
+ /// timeout <b>before</b> activating any child - asserted via {@code
_operationInProgress}, which also keeps the test
+ /// from passing vacuously if the timeout came from somewhere else.
+ @Test
+ public void testStreamingSelectionOrderByCombineOperatorHonorsDeadline() {
+ List<Operator> operators = getOperators(null, minMaxSegmentSupplier());
Review Comment:
Fixed in f3f2db3 - passes a fresh CountDownLatch(1) that's never awaited
instead of null.
Correction: this can't actually NPE - the null guard predates this PR and
still serves other tests. The real regression mode is worse for CI legibility:
SlowOperator sleeps 3,600,000 ms, so a deadline regression hangs this test for
an hour instead of failing. Not adding a @Test timeout for that since it
changes test semantics and wasn't asked for, but can if you'd like.
[addressed by agent]
--
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]
