rohityadav1993 commented on code in PR #19120: URL: https://github.com/apache/pinot/pull/19120#discussion_r3719665643
########## pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java: ########## @@ -0,0 +1,529 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.operator.query; + +import com.google.common.base.CaseFormat; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.PriorityQueue; +import java.util.Set; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.OrderByExpressionContext; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.common.RowBasedBlockValueFetcher; +import org.apache.pinot.core.operator.BaseOperator; +import org.apache.pinot.core.operator.BaseProjectOperator; +import org.apache.pinot.core.operator.BitmapDocIdSetOperator; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.operator.ExecutionStatistics; +import org.apache.pinot.core.operator.ExplainAttributeBuilder; +import org.apache.pinot.core.operator.ProjectionOperator; +import org.apache.pinot.core.operator.ProjectionOperatorUtils; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock; +import org.apache.pinot.core.operator.transform.TransformOperator; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.selection.SelectionOperatorUtils; +import org.apache.pinot.core.query.utils.OrderByComparatorFactory; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.spi.query.QueryScanCostContext; +import org.roaringbitmap.RoaringBitmap; + + +/** + * Lazy, incremental selection ORDER BY operator for segments that are physically sorted on the first order-by column. + * + * <p>Unlike {@link SelectionOrderByOperator} (which materializes the segment's whole top-K in a single block) this + * operator emits one globally-sorted {@link SelectionResultsBlock} per {@link #getNextBlock()} call and returns + * {@code null} when the segment is exhausted, so that a downstream k-way-merge combine operator can pull from many + * segments lazily and stop early. It relies on the underlying project operator iterating the first order-by column in + * the query order (the caller must guarantee {@code projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))}). + * + * <p>It runs in one of two emission modes: + * <ul> + * <li><b>No tail to sort</b> ({@code numSortedExpressions == numOrderByExpressions}, e.g. {@code ORDER BY sorted}): + * rows already arrive from the project operator in final order, so each call emits the next project block (trimmed to + * the remaining {@code limit + offset} budget).</li> + * <li><b>Tail to sort</b> ({@code numSortedExpressions < numOrderByExpressions}, e.g. + * {@code ORDER BY sorted, other}): + * each call reads forward until the first order-by value changes (a primary-value "run"), retains the run's top + * {@code limit + offset} rows by the full comparator, and emits them sorted. This bounds the in-memory run buffer to + * {@code limit + offset} rows even when the first order-by column is near-constant (very low cardinality).</li> + * </ul> + * + * <p>Like {@link SelectionOrderByOperator} it preserves the two-phase projection optimization: when there are output + * expressions that are not order-by expressions, the forward scan only fetches the order-by expressions plus the + * document id, and the non-order-by expressions are fetched in a second pass over the retained document ids of each + * emitted block. + * + * <p>This operator is stateful across {@link #getNextBlock()} calls and is <b>not</b> thread-safe; a single consumer + * must drive it. + */ +public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionResultsBlock> { + private static final String EXPLAIN_NAME = "SELECT_ORDERBY_STREAMING"; + + private final IndexSegment _indexSegment; + private final QueryContext _queryContext; + private final boolean _nullHandlingEnabled; + // Deduped order-by expressions followed by output expressions from SelectionOperatorUtils.extractExpressions() + private final List<ExpressionContext> _expressions; + private final BaseProjectOperator<?> _projectOperator; + private final List<OrderByExpressionContext> _orderByExpressions; + private final ColumnContext[] _orderByColumnContexts; + private final int _numExpressions; + private final int _numOrderByExpressions; + private final int _numRowsToKeep; + // Whether there are output expressions that are not order-by expressions (requires the two-phase fetch) + private final boolean _twoPhase; + // Whether the order-by has an unsorted tail that must be sorted in memory per run + private final boolean _tailToSort; + // Expressions fetched during the forward scan: order-by expressions only when two-phase, otherwise all expressions + private final List<ExpressionContext> _phase1Expressions; + private final int _numPhase1Columns; + private final Comparator<Object[]> _comparator; + // Compares only the first order-by column; used to detect primary-value run boundaries + private final Comparator<Object[]> _primaryComparator; + // Pre-allocated run heap (cleared and reused each nextRun() call to avoid per-run allocation) + private final Comparator<Object[]> _reversedComparator; + private final PriorityQueue<Object[]> _runHeap; + + // Pre-computed invariants for the two-phase fetch (null when single-phase) + private final List<ExpressionContext> _nonOrderByExpressions; + private final Map<String, DataSource> _phase2DataSourceMap; + private final int _phase2NumColumns; + + // Lazily built and cached; for two-phase it requires the transform operator's result column contexts + private DataSchema _dataSchema; + + // Forward-scan cursor state (used by the tail-to-sort mode) + private ValueBlock _currentBlock; + private RowBasedBlockValueFetcher _currentFetcher; + private int[] _currentDocIds; + private RoaringBitmap[] _currentNullBitmaps; + private int _currentNumDocs; + private int _currentPos; + // One-row lookahead: the first row of the next run, stashed when a run boundary is crossed + private Object[] _pendingRow; + private boolean _projectExhausted; + + private boolean _exhausted; + private int _numRowsEmitted; + private int _numDocsScanned = 0; + private long _numEntriesScannedPostFilter = 0; + + public StreamingSelectionOrderByOperator(IndexSegment indexSegment, QueryContext queryContext, + List<ExpressionContext> expressions, BaseProjectOperator<?> projectOperator, int numSortedExpressions) { + _indexSegment = indexSegment; + _queryContext = queryContext; + _nullHandlingEnabled = queryContext.isNullHandlingEnabled(); + _expressions = expressions; + _projectOperator = projectOperator; + + _orderByExpressions = queryContext.getOrderByExpressions(); + assert _orderByExpressions != null; + _numExpressions = expressions.size(); + _numOrderByExpressions = _orderByExpressions.size(); + _orderByColumnContexts = new ColumnContext[_numOrderByExpressions]; + for (int i = 0; i < _numOrderByExpressions; i++) { + ExpressionContext expression = _orderByExpressions.get(i).getExpression(); + _orderByColumnContexts[i] = _projectOperator.getResultColumnContext(expression); + } + + _numRowsToKeep = queryContext.getOffset() + queryContext.getLimit(); + _twoPhase = _numExpressions > _numOrderByExpressions; + _tailToSort = numSortedExpressions < _numOrderByExpressions; + _comparator = + OrderByComparatorFactory.getComparator(_orderByExpressions, _orderByColumnContexts, _nullHandlingEnabled); + // The first order-by column is the physically sorted column, so it never contains nulls on this path; comparing + // only index 0 is enough to detect when one primary-value run ends and the next begins. + _primaryComparator = + OrderByComparatorFactory.getComparator(_orderByExpressions, _orderByColumnContexts, _nullHandlingEnabled, 0, 1); + _reversedComparator = _comparator.reversed(); + _runHeap = new PriorityQueue<>( + Math.min(_numRowsToKeep, SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY), _reversedComparator); + + if (_twoPhase) { + _phase1Expressions = new ArrayList<>(_numOrderByExpressions); + for (OrderByExpressionContext orderByExpression : _orderByExpressions) { + _phase1Expressions.add(orderByExpression.getExpression()); + } + _nonOrderByExpressions = _expressions.subList(_numOrderByExpressions, _numExpressions); + Set<String> columns = new HashSet<>(); + for (ExpressionContext expressionContext : _nonOrderByExpressions) { + expressionContext.getColumns(columns); + } + _phase2NumColumns = columns.size(); + _phase2DataSourceMap = new HashMap<>(); + for (String column : columns) { + _phase2DataSourceMap.put(column, _indexSegment.getDataSource(column, _queryContext.getSchema())); + } + } else { + _phase1Expressions = _expressions; + _nonOrderByExpressions = null; + _phase2NumColumns = 0; + _phase2DataSourceMap = null; + // Single-phase: all output expressions are order-by expressions, so their types are known up front. + _dataSchema = buildSinglePhaseDataSchema(); + } + _numPhase1Columns = _phase1Expressions.size(); + } + + @Override + protected SelectionResultsBlock getNextBlock() { + if (_exhausted) { + return null; + } + List<Object[]> rows = _tailToSort ? nextRun() : nextSortedRows(); + if (rows == null || rows.isEmpty()) { + _exhausted = true; + return null; + } + if (_twoPhase) { + fetchNonOrderByColumns(rows); + } + // Single-phase builds the schema in the constructor; two-phase builds it during fetchNonOrderByColumns above. + assert _dataSchema != null; + return new SelectionResultsBlock(_dataSchema, rows, _comparator, _queryContext); + } + + /** + * No-tail-to-sort mode: the project operator already returns rows in final order, so emit the next project block, + * trimmed to the remaining {@code limit + offset} budget. Returns {@code null} when exhausted. + */ + @Nullable + private List<Object[]> nextSortedRows() { + int remaining = _numRowsToKeep - _numRowsEmitted; + if (remaining <= 0) { + return null; + } + ValueBlock valueBlock = _projectOperator.nextBlock(); + if (valueBlock == null) { + return null; + } + int numDocsFetched = valueBlock.getNumDocs(); + BlockValSet[] blockValSets = new BlockValSet[_numPhase1Columns]; + for (int i = 0; i < _numPhase1Columns; i++) { + blockValSets[i] = valueBlock.getBlockValueSet(_phase1Expressions.get(i)); + } + RowBasedBlockValueFetcher blockValueFetcher = new RowBasedBlockValueFetcher(blockValSets); + int[] docIds = _twoPhase ? valueBlock.getDocIds() : null; + RoaringBitmap[] nullBitmaps = null; + if (_nullHandlingEnabled) { + nullBitmaps = new RoaringBitmap[_numPhase1Columns]; + for (int i = 0; i < _numPhase1Columns; i++) { + nullBitmaps[i] = blockValSets[i].getNullBitmap(); + } + } + _numDocsScanned += numDocsFetched; + _numEntriesScannedPostFilter += (long) numDocsFetched * _projectOperator.getNumColumnsProjected(); + reportScanCost(numDocsFetched, (long) numDocsFetched * _projectOperator.getNumColumnsProjected()); + + // Rows arrive sorted; we only need the first 'remaining' of them globally. + int numRows = Math.min(numDocsFetched, remaining); + List<Object[]> rows = new ArrayList<>(numRows); + for (int i = 0; i < numRows; i++) { + rows.add(materializeRow(blockValueFetcher, docIds, nullBitmaps, i)); + } + _numRowsEmitted += rows.size(); + return rows; + } + + /** + * Tail-to-sort mode: read forward until the first order-by value changes, retain the run's top {@code limit + offset} + * rows by the full comparator, and return them sorted. Returns {@code null} when exhausted. + */ + @Nullable + private List<Object[]> nextRun() { + int remaining = _numRowsToKeep - _numRowsEmitted; + if (remaining <= 0) { + return null; + } + if (_pendingRow == null) { + _pendingRow = nextRow(); + if (_pendingRow == null) { + return null; + } + } + PriorityQueue<Object[]> runHeap = _runHeap; + runHeap.clear(); + Object[] runFirstRow = _pendingRow; + SelectionOperatorUtils.addToPriorityQueue(_pendingRow, runHeap, _numRowsToKeep); + _pendingRow = null; + Object[] row; + while ((row = nextRow()) != null) { + if (_primaryComparator.compare(row, runFirstRow) == 0) { Review Comment: Self review: This is going to be inefficient when segment sort expression is not same as orderBy experssion and there are too few rows per key. Consciously keeping it out of scope for now to keep the logic simple for happy path. -- 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]
