rohityadav1993 commented on code in PR #19120:
URL: https://github.com/apache/pinot/pull/19120#discussion_r4044525249
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java:
##########
@@ -161,6 +173,10 @@ private BaseCombineOperator getCombineOperator() {
List<OrderByExpressionContext> orderByExpressions =
_queryContext.getOrderByExpressions();
assert orderByExpressions != null;
if (orderByExpressions.get(0).getExpression().getType() ==
ExpressionContext.Type.IDENTIFIER) {
+ if (_queryContext.isSortedSelectionMergeEnabled()) {
Review Comment:
Accepted, done in f3f2db3. Deleted CombinePlanNode's non-streamer branch;
the streaming combine only appears when _streamer != null. The blocking SSE
path keeps MinMaxValueBasedSelectionOrderByCombineOperator unconditionally.
Also went further than asked: InstancePlanMakerImplV2.makeInstancePlan now
clears sortedSelectionMergeMode for the whole blocking plan, since
SelectionPlanNode gates the leaf on the same option independently and has no
way to know which combine it feeds - without this the option would have been
half-honored (streaming leaf, blocking combine, no actual streaming).
Verified this pairing was not a correctness bug (row counts/contents matched
the baseline in every shape I tried), but removing it is still right for
coherence, and because [10]/[15] propose changes that would break that
accidental correctness.
Reworded per your note: gate is "streaming-capable paths" (_streamer !=
null), not "MSE only".
[addressed by agent]
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java:
##########
@@ -130,6 +131,17 @@ private BaseCombineOperator getCombineOperator() {
// Use streaming operator only for non-empty selection-only query
return new StreamingSelectionOnlyCombineOperator(operators,
_queryContext, _executorService);
}
+ // Streaming selection order-by (opt-in via the
sortedSelectionMergeEnabled hint). Selection-only already
+ // returned above, so reaching here with a non-empty limit and an
order-by present implies selection order-by.
+ if (_queryContext.isSortedSelectionMergeEnabled() &&
QueryContextUtils.isSelectionQuery(_queryContext)
Review Comment:
Done in f3f2db3. Replaced the boolean with sortedSelectionMergeMode
(OFF/ON/AUTO, default OFF), following the getJoinOverflowMode idiom.
ON forces the streaming path unconditionally, including non-identifier
order-bys. AUTO takes it when >= sortedSelectionMergeAutoMinSortedRatio
(default 0.8, unmeasured) of segments are physically sorted, decided from
metadata alone, no acquire needed.
Resolved once in InstancePlanMakerImplV2, not independently per gate -
CombinePlanNode.run() builds leaves before its own gate runs, so a per-gate
AUTO decision would arrive too late (same issue as comment [2]).
Known gap, not fixed here: AUTO's metadata-only sortedness check does not
see nulls in the leading column, so with null handling on it can resolve ON
while the leaf's fuller check declines to stream - correct results, just no
speedup. Pinned by testAutoIgnoresNullsInTheLeadingColumn. Would like a
go-ahead before filing an issue to track it.
Tests: 62 combine/operator tests + 30 in QueryOptionsUtilsTest covering
mode/ratio parsing, AUTO thresholds, ON forcing, and the gap above.
[addressed by agent]
##########
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)) {
Review Comment:
Done in f3f2db3, but not via your first suggestion:
ReverseDocIdSetOperator.initializeBitmap() drains the whole matching docId set
into a RoaringBitmap up front when there's no bitmap index to reuse, which is
likely why allowReverseOrder defaults false. Implying it from the streaming
branch would silently trade away the bounded memory this feature is for, so
reversal stays an explicit user lever.
What changed: InstancePlanMakerImplV2.isSortedEnoughForStreamingMerge is now
direction-aware - DESC without allowReverseOrder resolves AUTO to OFF, in one
place both gates read. ON still force-enables DESC as a benchmark lever.
Tests: testAutoDoesNotSelectStreamingForDescWithoutReverseOrder,
testAutoSelectsStreamingForDescWithReverseOrder,
testAutoHonoursTheMinSortedRatioThresholdForDesc,
testAutoOnlyChecksTheLeadingOrderByDirection.
Residual, not fixed here: allowReverseOrder=true doesn't guarantee a
reversed scan (getSortedByProject silently falls back if withOrder(DESC)
throws), so AUTO can still resolve ON over materialized children in that case.
Same shape as the [1] null-handling gap; deferred pending a go-ahead to track
it.
PR description now has a DESC behavior paragraph.
[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() {
Review Comment:
Done in f3f2db3, steps 2 and 3. Step 1 (the Operator.nextBlock() javadoc)
I'd rather do as its own PR since it touches every implementor's contract -
will link once opened.
Step 2: a zero-match segment now emits one empty schema-carrying block
before its terminal null, instead of null immediately.
Step 3: resolveDataSchema() deleted.
One correction: I don't think step 3 was enabled by step 2 as framed. Rows
only ever reach _outputRows through pullBlock(), which already sets _dataSchema
first, and a true zero-match query never calls flushDataBlock() at all.
Negative control confirmed it: disabling step 2 and deleting
resolveDataSchema() together still passed all 30 combine tests. So the
invariant held by avoidance before, and by construction now - the deletion was
always safe, step 2 just makes it stay safe under future changes.
Not fixed here: a zero-match query (not just segment) still returns no
schema, since the combine then exits through MetadataResultsBlock, whose
getDataSchema() is hardcoded null. That's a convention shared by every
BaseStreamingCombineOperator subclass, so I left it as a pinned gap
(testEmptyResultOnStreamingPath asserts assertNull(result._schema)) rather than
a cross-cutting fix here.
[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()) {
Review Comment:
Done in f3f2db3, framed as a coherence fix rather than a live bug: every
BaseDocIdSetOperator reachable here ends with "if (pos > 0) return block; else
return null", so the case you describe can't currently happen. Kept the
tolerance anyway since it's the dominant pattern elsewhere
(SelectionOnlyOperator, AggregationOperator, etc.) and two open threads ([10],
[15]) propose block-sizing changes this invariant would otherwise depend on.
Extracted one shared nextNonEmptyBlock() helper used by both nextRow() and
nextSortedRows(), instead of a second copy of the skip.
Tests inject a mid-stream zero-doc block via a decorator (single-phase,
two-phase, tail-to-sort, consecutive empties, empties before exhaustion,
OFFSET, null handling). Verified against pre-fix code: the sorted-scan cases
fail with 10 rows instead of 25.
[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() {
+ }
+
+ /// Override the base worker-queue stop: no worker threads / phaser tasks
were started. Release any segments still
+ /// acquired (idempotent) so an early stop by the driver cannot leak
acquires.
+ @Override
+ public void stop() {
+ _done = true;
+ releaseAllCursors();
+ }
+
+ /// The base worker-thread entry point must never run here ({@link #start()}
is a no-op). Fail loud if it ever does.
+ @Override
+ protected void processSegments() {
+ throw new IllegalStateException(
+ "StreamingSelectionOrderByCombineOperator runs single-threaded;
processSegments() must not be called");
+ }
+
+ @Override
+ protected BaseResultsBlock getNextBlock() {
+ if (_done) {
+ // Streaming mode: terminal metadata block after the last data block.
Idempotent if called again.
+ return attachExecutionStats(new MetadataResultsBlock());
+ }
+ try {
+ long endTimeMs = _queryContext.getEndTimeMs();
+ while (_numRowsEmitted < _numRowsToKeep) {
+ // The merge drains the heap on this thread and, for single-block
cursors, may never re-enter a child operator.
+ // Without this check nothing would observe cancellation, pause, or
the query deadline for up to
+ // (limit + offset) iterations -- a regression against
MinMaxValueBasedSelectionOrderByCombineOperator, which
+ // this operator replaces on the same query shape.
+
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(_numRowsEmitted,
EXPLAIN_NAME, endTimeMs);
+ activateEligibleCursors();
+ SegmentCursor cursor = _priorityQueue.poll();
Review Comment:
Done in f3f2db3, thanks for the idea. No benchmark numbers though - the
change lands on the structural argument only (one peek()+compare per row
instead of poll()+offer()). Happy to run BenchmarkSelectionOrderByFilterPruning
as a follow-up if useful.
Two things beyond the sketch: the loop returns mid-merge on block flush, so
the leader is re-offered to the heap before each return (one offer/poll per
block, not per row). And pruning's frontier used to read off the heap head,
which is absent while the leader serves a long run - now tested against both
the leader's head and the heap head.
Caught one bug before shipping: using the last emitted row as the frontier
is unsafe (a lower bound, not the next row), which could reorder output. Every
new test here was watched failing against a broken version first.
Deliberately not pinned: that the incumbent wins a tie. Tie order isn't
promised; the test only asserts ties are harmless.
[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]
